Update app.py

app.py

@@ -15,7 +15,7 @@ import spaces
 model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
 pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
-flow_shift = 5.0 #5.0  1.0 for image, 5.0 for 720P, 3.0 for 480P
+flow_shift = 1.0 #5.0  1.0 for image, 5.0 for 720P, 3.0 for 480P
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
 
 # Configure DDIMScheduler with a beta schedule
@@ -35,11 +35,22 @@ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow
 #     flow_shift=flow_shift           # Retain flow_shift for WanPipeline compatibility
 # )
 
+CAUSVID_LORA_REPO = "WanVideo_comfy"
+CAUSVID_LORA_FILENAME = "Wan21_CausVid_bidirect2_T2V_1_3B_lora_rank32.safetensors"
+
+try:
+    causvid_path = hf_hub_download(repo_id=CAUSVID_LORA_REPO, filename=CAUSVID_LORA_FILENAME)
+    pipe.load_lora_weights(causvid_path, adapter_name="causvid_lora")
+    print("✅ CausVid LoRA loaded (strength: 1.0)")
+except Exception as e:
+    print(f"⚠️ CausVid LoRA not loaded: {e}")
+    causvid_path = None
+
 
 @spaces.GPU()
 def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_steps=30, lora_id=None, progress=gr.Progress(track_tqdm=True)):
     if lora_id and lora_id.strip() != "":
-        pipe.unload_lora_weights()
+        # pipe.unload_lora_weights()
         pipe.load_lora_weights(lora_id.strip())
     pipe.to("cuda")
     # apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold=0.2))
@@ -55,14 +66,15 @@ def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_ste
             width=width,
             num_frames=1,
             num_inference_steps=num_inference_steps,
-            guidance_scale=5.0, #5.0
+            guidance_scale=1.0, #5.0
         )
         image = output.frames[0][0]
         image = (image * 255).astype(np.uint8)
         return Image.fromarray(image)
     finally:
         if lora_id and lora_id.strip() != "":
-            pipe.unload_lora_weights()
+            pass
+            # pipe.unload_lora_weights()
 
 iface = gr.Interface(
     fn=generate,