offload rewriter

app.py

@@ -67,16 +67,15 @@ pipe.scheduler = UniPCMultistepScheduler.from_config(
 print(f"✓ Configured scheduler (flow_shift=2.0)")
 
 # Move to device
-pipe.to("cuda")
+# pipe.to("cuda")
 print(f"✓ Models loaded and moved to cuda")
 
 end = time.time()
 print(f"Model loaded in {end - start:.2f}s.")
 
 start = time.time()
-prompt_enhancer_model = "Qwen/Qwen2.5-VL-7B-Instruct"
+prompt_enhancer_model = "Qwen/Qwen3-VL-30B-A3B-Instruct"
 prompt_model, processor = load_model(prompt_enhancer_model)
-prompt_model.to("cuda")
 end = time.time()
 print(f"Prompt enhancer loaded in {end - start:.2f}s.")
 
@@ -114,6 +113,7 @@ def run_inference(
 	num_temporal_reasoning_steps: int = 8,
 ):
 	# Rewriter
+	prompt_model.to("cuda")
 	final_prompt = prompt
 
 	with th.no_grad():
@@ -132,8 +132,10 @@ def run_inference(
 	print(cot_prompt)
 	print("=" * 80 + "\n")
 	final_prompt = cot_prompt
+	prompt_model.to("cpu")
 
 	# Inference
+	pipe.to("cuda")
 	print(f"Loading input image: {image_path}")
 	image = load_image(image_path)
 	mod_value = pipe.vae_scale_factor_spatial * pipe.transformer.config.patch_size[1]
@@ -159,6 +161,7 @@ def run_inference(
 		).frames[0]
 
 		end = time.time()
+	pipe.to("cpu")
 
 	video_tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
 	output_path_video = video_tmp.name
@@ -227,16 +230,16 @@ def build_ui() -> gr.Blocks:
 			examples=[ 
 				[
 					"examples/1.png",
-					"The user wants to change the scene so that the girl in the traditional-style painting, wearing her ornate floral robe and headdress, is now playing a guitar. ",
-					False,
+					"The user wants to change the provided illustration of an elegant woman in a flowing red kimono into a high-end Japanese anime PVC scale figure, rendered photorealistically as a pre-painted collectible. Preserve her long black hair styled with golden hair ornaments and delicate floral accessories, her slightly tilted head and confident gaze, and the detailed red kimono with golden and floral embroidery tied with a wide gold obi.  Cherry blossom petals drift around.   Maintain the pose and camera view point unchanged. The scene should look like a premium finished PVC figure on display, with realistic textures, fine paint detailing, and a polished collectible presentation. Place the figure on a simple round base on a computer desk, with blurred keyboard and monitor glow in the background. Emphasize a strong 3D sense of volume and depth, realistic shadows and lighting, and painted PVC figure textures. Professional studio photography style, shallow depth of field, focus on the figure as a physical collectible. The lighting on the figure is uniform and highlighted, emphasizing every sculpted detail and painted accent.",
+					True,
 				],
 				[
-					"examples/1.png",
-					"The user wants to change the scene so that the girl in the traditional-style painting, wearing her ornate floral robe and headdress, is now playing a guitar.",
-					True,
+					"examples/2.png",
+					"The user wants to change the scene so that the girl in the traditional-style painting, wearing her ornate floral robe and headdress, is now playing a guitar. ",
+					False,
 				],
 			],
-			inputs=[image, prompt, enable_temporal_reasoning], outputs=[output_image, output_video], fn=_on_run, cache_examples=True
+			inputs=[image, prompt, enable_temporal_reasoning], outputs=[output_image, output_video], fn=_on_run, cache_examples="lazy"
 		)
 
 	return demo
@@ -245,4 +248,4 @@ def build_ui() -> gr.Blocks:
 if __name__ == "__main__":
 	demo = build_ui()
 	# demo.launch(server_name="0.0.0.0", server_port=7869)
-	demo.queue().launch(share=True)
+	demo.queue().launch()