Update app.py

app.py

@@ -14,76 +14,52 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
 
 model_id = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
-print("Loading VAE...")
-vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
 
-print("Loading WanPipeline in bfloat16...")
-# This will use ZeroGPU/accelerate with meta devices
-pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
+# --- PRE-FUSION STRATEGY ---
+# 1. Load everything on the CPU first. Do NOT use device_map or torch_dtype here yet.
+#    This prevents ZeroGPU/accelerate from taking over too early.
+print("Loading all components on CPU for pre-fusion...")
+vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae")
+pipe = WanPipeline.from_pretrained(model_id, vae=vae)
 
-flow_shift = 1.0
-pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
-
-# Move the base pipeline to the GPU. ZeroGPU will manage this.
-print("Moving pipeline to device (ZeroGPU will handle offloading)...")
-pipe.to(device)
-
-# --- LORA SETUP ---
+# 2. Load and FUSE the LoRA while the entire pipeline is still on the CPU.
+print("Loading and fusing base LoRA on CPU...")
 CAUSVID_LORA_REPO = "Kijai/WanVideo_comfy"
 CAUSVID_LORA_FILENAME = "Wan21_CausVid_14B_T2V_lora_rank32.safetensors"
-BASE_LORA_NAME = "causvid_lora"
-CUSTOM_LORA_NAME = "custom_lora"
-
-print("Downloading base LoRA...")
 try:
     causvid_path = hf_hub_download(repo_id=CAUSVID_LORA_REPO, filename=CAUSVID_LORA_FILENAME)
-    print("✅ Base LoRA downloaded.")
+    pipe.load_lora_weights(causvid_path)
+    print("✅ Base LoRA loaded on CPU.")
+    
+    pipe.fuse_lora()
+    print("✅ Base LoRA fused on CPU.")
+    
+    # After fusing, we must unload the original LoRA weights to save memory
+    pipe.unload_lora_weights()
+    print("✅ Original LoRA weights unloaded.")
+    
 except Exception as e:
-    causvid_path = None
-    print(f"⚠️ Could not download base LoRA: {e}")
+    print(f"⚠️ Could not pre-fuse the LoRA: {e}. The model will run without it.")
 
-print("Initialization complete. Gradio is starting...")
+# 3. NOW, move the final, fused model to the target device and set the dtype.
+#    ZeroGPU will now take over a single, coherent model.
+print("Moving the final, fused pipeline to device and setting dtype...")
+pipe.to(device=device, dtype=torch.bfloat16)
 
-@spaces.GPU()
-def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_steps=30, lora_id=None, progress=gr.Progress(track_tqdm=True)):
-    
-    # --- DYNAMIC LORA MANAGEMENT FOR EACH RUN ---
-    active_adapters = []
-    adapter_weights = []
+# 4. Configure the scheduler
+flow_shift = 1.0
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
 
-    # 1. Load the Base LoRA directly onto the correct device
-    if causvid_path:
-        try:
-            print(f"Loading base LoRA '{BASE_LORA_NAME}'...")
-            pipe.load_lora_weights(causvid_path, adapter_name=BASE_LORA_NAME, device_map={"":device})
-            active_adapters.append(BASE_LORA_NAME)
-            adapter_weights.append(1.0)
-            print("✅ Base LoRA loaded to device.")
-        except Exception as e:
-            print(f"⚠️ Failed to load base LoRA: {e}")
+# Clean up memory before starting the server
+del vae
+gc.collect()
+torch.cuda.empty_cache()
 
-    # 2. Load the Custom LoRA if provided, also directly to the device
-    clean_lora_id = lora_id.strip() if lora_id else ""
-    if clean_lora_id:
-        try:
-            print(f"Loading custom LoRA '{CUSTOM_LORA_NAME}' from '{clean_lora_id}'...")
-            pipe.load_lora_weights(clean_lora_id, adapter_name=CUSTOM_LORA_NAME, device_map={"":device})
-            active_adapters.append(CUSTOM_LORA_NAME)
-            adapter_weights.append(1.0)
-            print("✅ Custom LoRA loaded to device.")
-        except Exception as e:
-            print(f"⚠️ Failed to load custom LoRA '{clean_lora_id}': {e}")
-            if CUSTOM_LORA_NAME in getattr(pipe.transformer, 'peft_config', {}):
-                del pipe.transformer.peft_config[CUSTOM_LORA_NAME]
+print("Initialization complete. Gradio is starting...")
 
-    # 3. Activate the successfully loaded adapters
-    if active_adapters:
-        print(f"Activating adapters: {active_adapters} with weights: {adapter_weights}")
-        pipe.set_adapters(active_adapters, adapter_weights)
-        pipe.transformer.to(device)  # Explicitly move transformer to GPU after setting adapters
-    else:
-        # Ensure LoRA is disabled if no adapters were loaded
-        pipe.disable_lora()
+@spaces.GPU()
+def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_steps=30, progress=gr.Progress(track_tqdm=True)):
+    # The LoRA is permanently fused. No dynamic loading is possible or needed.
     
     apply_cache_on_pipe(pipe)
     
@@ -101,22 +77,18 @@ def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_ste
         image = (image * 255).astype(np.uint8)
         return Image.fromarray(image)
     finally:
-        # --- PROPER CLEANUP ---
-        print("Unloading all LoRAs to ensure a clean state...")
-        pipe.unload_lora_weights()
-        gc.collect()  # Force garbage collection
-        torch.cuda.empty_cache()  # Clear CUDA cache
-        print("✅ LoRAs unloaded and memory cleaned.")
+        # No cleanup needed as the state is static
+        pass
 
+# Interface is simplified as the custom LoRA option is removed for stability.
 iface = gr.Interface(
     fn=generate,
     inputs=[
         gr.Textbox(label="Input prompt"),
-        gr.Textbox(label="Negative prompt", value="Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"),
+        gr.Textbox(label="Negative prompt", value = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"),
         gr.Slider(label="Width", minimum=480, maximum=1280, step=16, value=1024),
         gr.Slider(label="Height", minimum=480, maximum=1280, step=16, value=1024),
         gr.Slider(minimum=1, maximum=80, step=1, label="Inference Steps", value=10),
-        gr.Textbox(label="LoRA ID (Optional, loads dynamically)"),
     ],
     outputs=gr.Image(label="output"),
 )