Update app.py

app.py

@@ -70,8 +70,8 @@ if lora_path:
 	)
 	print(f"✓ Configured scheduler (flow_shift=2.0)")
 
-# pipe.enable_model_cpu_offload()
 pipe.to("cuda")
+# pipe.enable_model_cpu_offload()
 end = time.time()
 print(f"Model loaded in {end - start:.2f}s.")
 
@@ -117,7 +117,8 @@ def run_inference(
 ):
 	# Rewriter
 	final_prompt = prompt
-
+    
+	prompt_model.to("cuda")
 	# Enhance prompt with CoT reasoning
 	start = time.time()
 	cot_prompt = enhance_prompt(
@@ -128,6 +129,7 @@ def run_inference(
 	)
 	end = time.time()
 	print(f"Prompt enhanced in {end - start:.2f}s")
+	prompt_model.to("cpu")
 
 	# Print enhanced CoT prompt
 	print("\n" + "=" * 80)
@@ -160,6 +162,7 @@ def run_inference(
 		guidance_scale=guidance_scale,
 		enable_temporal_reasoning=enable_temporal_reasoning,
 		num_temporal_reasoning_steps=num_temporal_reasoning_steps,
+        # offload_model=True
 	).frames[0]
 	end = time.time()
 	print(f"Generated video in {end - start:.2f}s")