Update app_lora.py

app_lora.py

@@ -36,17 +36,8 @@ DEFAULT_NAG_NEGATIVE_PROMPT = "Static, motionless, still, ugly, bad quality, wor
 MODEL_ID = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
 SUB_MODEL_ID = "vrgamedevgirl84/Wan14BT2VFusioniX"
 SUB_MODEL_FILENAME = "Wan14BT2VFusioniX_fp16_.safetensors"
-
-
-
-#LORA_REPO_ID = "Kijai/WanVideo_comfy"
-#LORA_FILENAME = "Wan21_CausVid_14B_T2V_lora_rank32.safetensors"
-
-
-# new experiment for future work
-
-LORA_REPO_ID = "vrgamedevgirl84/Wan14BT2VFusioniX"
-LORA_FILENAME = "FusionX_LoRa/Wan2.1_T2V_14B_FusionX_LoRA.safetensors"
+LORA_REPO_ID = "Kijai/WanVideo_comfy"
+LORA_FILENAME = "Wan21_CausVid_14B_T2V_lora_rank32.safetensors"
 
 vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
 wan_path = hf_hub_download(repo_id=SUB_MODEL_ID, filename=SUB_MODEL_FILENAME)
@@ -57,9 +48,15 @@ pipe = NAGWanPipeline.from_pretrained(
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=5.0)
 pipe.to("cuda")
 
-
-
-
+causvid_path = hf_hub_download(repo_id=LORA_REPO_ID, filename=LORA_FILENAME)
+pipe.load_lora_weights(causvid_path, adapter_name="causvid_lora")
+pipe.set_adapters(["causvid_lora"], adapter_weights=[0.95])
+for name, param in pipe.transformer.named_parameters():
+   if "lora_B" in name:
+       if "blocks.0" in name:
+           param.data = param.data * 0.25
+pipe.fuse_lora()
+pipe.unload_lora_weights()
 
 pipe.transformer.__class__.attn_processors = NagWanTransformer3DModel.attn_processors
 pipe.transformer.__class__.set_attn_processor = NagWanTransformer3DModel.set_attn_processor
@@ -134,6 +131,10 @@ def generate_video(
     else:
         baseline_video_path = None
 
+    if torch.cuda.is_available():
+        print("Allocated:", torch.cuda.memory_allocated() / 1024**2, "MB")
+        print("Cached:   ", torch.cuda.memory_reserved() / 1024**2, "MB")
+
     return nag_video_path, baseline_video_path, current_seed
 
 
@@ -150,15 +151,19 @@ def generate_video_with_example(
         seed=DEFAULT_SEED, randomize_seed=False,
         compare=True,
     )
+    if torch.cuda.is_available():
+        print("Allocated:", torch.cuda.memory_allocated() / 1024**2, "MB")
+        print("Cached:   ", torch.cuda.memory_reserved() / 1024**2, "MB")
     return nag_video_path, baseline_video_path, \
         DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE, \
         DEFAULT_DURATION_SECONDS, DEFAULT_STEPS, seed, True
 
 
 with gr.Blocks() as demo:
-    gr.Markdown('''# Normalized Attention Guidance (NAG) for fast 4 steps Wan2.1-T2V-14B with FuxionX Base T2V
+    gr.Markdown('''# Normalized Attention Guidance (NAG) for fast 4 steps Wan2.1-T2V-14B with CausVid LoRA
     Implementation of [Normalized Attention Guidance](https://chendaryen.github.io/NAG.github.io/).
- 
+    
+    [CausVid](https://github.com/tianweiy/CausVid) is a distilled version of Wan2.1 to run faster in just 4-8 steps, [extracted as LoRA by Kijai](https://huggingface.co/Kijai/WanVideo_comfy/blob/main/Wan21_CausVid_14B_T2V_lora_rank32.safetensors).
     ''')
 
     with gr.Row():
@@ -229,4 +234,4 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.queue().launch()