dmd lora 4step

app.py

@@ -99,7 +99,7 @@ def load_video_frames(video_path: str, source_frames: int):
     return input_video, original_height, original_width
 
 class VideoCoF_Controller(Wan_Controller):
-    @spaces.GPU(duration=2000)
+    @spaces.GPU(duration=300)
     @timer
     def generate(
         self,
@@ -141,6 +141,8 @@ class VideoCoF_Controller(Wan_Controller):
         repeat_rope_checkbox=True,
         fps=10,
         is_api=False,
+        # New arg for acceleration
+        enable_acceleration=False,
     ):
         self.clear_cache()
         print(f"VideoCoF Generation started.")
@@ -168,10 +170,21 @@ class VideoCoF_Controller(Wan_Controller):
         self.pipeline.scheduler = self.scheduler_dict[sampler_dropdown].from_config(scheduler_config)
 
         # LoRA merging
+        # 1. Merge VideoCoF LoRA
         if self.lora_model_path != "none":
-            print(f"Merge Lora.")
+            print(f"Merge VideoCoF Lora: {self.lora_model_path}")
             self.pipeline = merge_lora(self.pipeline, self.lora_model_path, multiplier=lora_alpha_slider)
 
+        # 2. Merge Acceleration LoRA (FusionX) if enabled
+        acc_lora_path = os.path.join(self.personalized_model_dir, "Wan2.1_Text_to_Video_14B_FusionX_LoRA.safetensors")
+        if enable_acceleration:
+            if os.path.exists(acc_lora_path):
+                print(f"Merge Acceleration LoRA: {acc_lora_path}")
+                # FusionX LoRA generally uses multiplier 1.0
+                self.pipeline = merge_lora(self.pipeline, acc_lora_path, multiplier=1.0)
+            else:
+                print(f"Warning: Acceleration LoRA not found at {acc_lora_path}")
+
         # Seed
         if int(seed_textbox) != -1 and seed_textbox != "": 
             torch.manual_seed(int(seed_textbox))
@@ -232,12 +245,23 @@ class VideoCoF_Controller(Wan_Controller):
 
         except Exception as e:
             print(f"Error: {e}")
+            # Unmerge in case of error (LIFO order)
+            if enable_acceleration and os.path.exists(acc_lora_path):
+                 print("Unmerging Acceleration LoRA (due to error)")
+                 self.pipeline = unmerge_lora(self.pipeline, acc_lora_path, multiplier=1.0)
+            
             if self.lora_model_path != "none":
+                 print("Unmerging VideoCoF LoRA (due to error)")
                  self.pipeline = unmerge_lora(self.pipeline, self.lora_model_path, multiplier=lora_alpha_slider)
             return gr.update(), gr.update(), f"Error: {str(e)}"
 
-        # Unmerge LoRA
+        # Unmerge LoRAs (LIFO order)
+        if enable_acceleration and os.path.exists(acc_lora_path):
+            print("Unmerging Acceleration LoRA")
+            self.pipeline = unmerge_lora(self.pipeline, acc_lora_path, multiplier=1.0)
+
         if self.lora_model_path != "none":
+            print("Unmerging VideoCoF LoRA")
             self.pipeline = unmerge_lora(self.pipeline, self.lora_model_path, multiplier=lora_alpha_slider)
 
         # Save output
@@ -278,9 +302,15 @@ def ui(GPU_memory_mode, scheduler_dict, config_path, compile_dit, weight_dtype):
                 from huggingface_hub import snapshot_download, hf_hub_download
                 print("Downloading Wan2.1-T2V-14B weights...")
                 snapshot_download(repo_id="Wan-AI/Wan2.1-T2V-14B", local_dir="Wan-AI/Wan2.1-T2V-14B")
-                print("Downloading VideoCoF weights...")
+                
                 os.makedirs("models/Personalized_Model", exist_ok=True)
+                
+                print("Downloading VideoCoF weights...")
                 hf_hub_download(repo_id="XiangpengYang/VideoCoF", filename="videocof.safetensors", local_dir="models/Personalized_Model")
+                
+                print("Downloading FusionX Acceleration LoRA...")
+                hf_hub_download(repo_id="MonsterMMORPG/Wan_GGUF", filename="Wan2.1_Text_to_Video_14B_FusionX_LoRA.safetensors", local_dir="models/Personalized_Model")
+                
             except Exception as e:
                 print(f"Warning: Failed to pre-download weights: {e}")
 
@@ -302,12 +332,17 @@ def ui(GPU_memory_mode, scheduler_dict, config_path, compile_dit, weight_dtype):
                 with gr.Column():
                     sampler_dropdown, sample_step_slider = create_samplers(controller)
                     
+                    # Default steps lowered to 4 for acceleration
+                    sample_step_slider.value = 4
+                    
                     # Custom VideoCoF Params
                     with gr.Group():
                         gr.Markdown("### VideoCoF Parameters")
                         source_frames_slider = gr.Slider(label="Source Frames", minimum=1, maximum=100, value=33, step=1)
                         reasoning_frames_slider = gr.Slider(label="Reasoning Frames", minimum=1, maximum=20, value=4, step=1)
                         repeat_rope_checkbox = gr.Checkbox(label="Repeat RoPE", value=True)
+                        # Add Acceleration Checkbox
+                        enable_acceleration = gr.Checkbox(label="Enable 4-step Acceleration (FusionX LoRA)", value=True)
                         
                     # Use custom height/width creation to hide/customize
                     resize_method, width_slider, height_slider, base_resolution = create_height_width_english(
@@ -338,6 +373,7 @@ def ui(GPU_memory_mode, scheduler_dict, config_path, compile_dit, weight_dtype):
                     # Set default seed to 0
                     cfg_scale_slider, seed_textbox, seed_button = create_cfg_and_seedbox(True)
                     seed_textbox.value = "0"
+                    cfg_scale_slider.value = 1.0
                     
                     generate_button = gr.Button(value="Generate", variant='primary')
 
@@ -366,7 +402,7 @@ def ui(GPU_memory_mode, scheduler_dict, config_path, compile_dit, weight_dtype):
                 cfg_scale_slider, 
                 start_image, 
                 end_image, 
-                validation_video,
+                validation_video, 
                 validation_video_mask,
                 control_video,
                 denoise_strength, 
@@ -382,7 +418,8 @@ def ui(GPU_memory_mode, scheduler_dict, config_path, compile_dit, weight_dtype):
                 # New inputs
                 source_frames_slider,
                 reasoning_frames_slider,
-                repeat_rope_checkbox
+                repeat_rope_checkbox,
+                enable_acceleration
             ],
             outputs=[result_image, result_video, infer_progress]
         )