Update app_lora.py

app_lora.py

@@ -16,17 +16,13 @@ import warnings
 warnings.filterwarnings("ignore", message=".*Attempting to use legacy OpenCV backend.*")
 warnings.filterwarnings("ignore", message=".*num_frames - 1.*")
 
-# This decorator is specific to HuggingFace Spaces and will cause an error in other environments.
-# import spaces
-
 MODEL_ID = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
-
 LORA_REPO_ID = "vrgamedevgirl84/Wan14BT2VFusioniX"
 LORA_FILENAME = "FusionX_LoRa/Wan2.1_I2V_14B_FusionX_LoRA.safetensors"
 
-# Initialize models with proper dtype handling
-# This section requires a GPU and CUDA to be available
+# --- Model Initialization ---
 pipe = None
+# This check correctly identifies if the Hugging Face Space has a GPU.
 if torch.cuda.is_available():
     image_encoder = CLIPVisionModel.from_pretrained(MODEL_ID, subfolder="image_encoder", torch_dtype=torch.float16)
     vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float16)
@@ -34,65 +30,47 @@ if torch.cuda.is_available():
         MODEL_ID, vae=vae, image_encoder=image_encoder, torch_dtype=torch.float16
     )
     pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=8.0)
-
-    # Enable memory efficient attention and CPU offloading for large videos
     pipe.enable_model_cpu_offload()
 
     try:
         causvid_path = hf_hub_download(repo_id=LORA_REPO_ID, filename=LORA_FILENAME)
         print("✅ LoRA downloaded to:", causvid_path)
-
         pipe.load_lora_weights(causvid_path, adapter_name="causvid_lora")
         pipe.set_adapters(["causvid_lora"], adapter_weights=[0.75])
         pipe.fuse_lora()
-
     except Exception as e:
         import traceback
         print("❌ Error during LoRA loading:")
         traceback.print_exc()
 else:
-    print("CUDA is not available. This script requires a GPU to run.")
-
+    print("CUDA is not available. This script requires a GPU. Please upgrade your Space hardware.")
 
+# --- Constants and Helper Functions ---
 MOD_VALUE = 32
-DEFAULT_H_SLIDER_VALUE = 640
-DEFAULT_W_SLIDER_VALUE = 1024
+DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE = 640, 1024
 NEW_FORMULA_MAX_AREA = 640.0 * 1024.0
-
 SLIDER_MIN_H, SLIDER_MAX_H = 128, 1024
 SLIDER_MIN_W, SLIDER_MAX_W = 128, 1024
 MAX_SEED = np.iinfo(np.int32).max
-
-FIXED_FPS = 24
-MIN_FRAMES_MODEL = 8  # Minimum 8 frames (~0.33s)
-MAX_FRAMES_MODEL = 240  # Maximum 240 frames (10 seconds at 24fps)
-
+FIXED_FPS, MIN_FRAMES_MODEL, MAX_FRAMES_MODEL = 24, 8, 240
 default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation"
 default_negative_prompt = "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, watermark, text, signature"
 
-
 def _calculate_new_dimensions_wan(pil_image, mod_val, calculation_max_area,
-                                 min_slider_h, max_slider_h,
-                                 min_slider_w, max_slider_w,
+                                 min_slider_h, max_slider_h, min_slider_w, max_slider_w,
                                  default_h, default_w):
     orig_w, orig_h = pil_image.size
-    if orig_w <= 0 or orig_h <= 0:
-        return default_h, default_w
-
+    if orig_w <= 0 or orig_h <= 0: return default_h, default_w
     aspect_ratio = orig_h / orig_w
-
     calc_h = round(np.sqrt(calculation_max_area * aspect_ratio))
     calc_w = round(np.sqrt(calculation_max_area / aspect_ratio))
-
     calc_h = max(mod_val, (calc_h // mod_val) * mod_val)
     calc_w = max(mod_val, (calc_w // mod_val) * mod_val)
-
     new_h = int(np.clip(calc_h, min_slider_h, (max_slider_h // mod_val) * mod_val))
     new_w = int(np.clip(calc_w, min_slider_w, (max_slider_w // mod_val) * mod_val))
-
     return new_h, new_w
 
-def handle_image_upload_for_dims_wan(uploaded_pil_image, current_h_val, current_w_val):
+def handle_image_upload_for_dims_wan(uploaded_pil_image):
     if uploaded_pil_image is None:
         return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
     try:
@@ -103,12 +81,10 @@ def handle_image_upload_for_dims_wan(uploaded_pil_image, current_h_val, current_
         )
         return gr.update(value=new_h), gr.update(value=new_w)
     except Exception as e:
-        gr.Warning("Error attempting to calculate new dimensions")
+        gr.Warning("Error calculating new dimensions.")
         return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
 
-
 def export_video_with_ffmpeg(frames, output_path, fps=24):
-    """Export video using imageio if available, otherwise fall back to OpenCV"""
     try:
         import imageio
         writer = imageio.get_writer(output_path, fps=fps, codec='libx264',
@@ -116,152 +92,87 @@ def export_video_with_ffmpeg(frames, output_path, fps=24):
         for frame in frames:
             writer.append_data(np.array(frame))
         writer.close()
-        return True
     except ImportError:
         export_to_video(frames, output_path, fps=fps)
-        return False
 
 def generate_video(input_image, prompt, height, width,
-                   negative_prompt=default_negative_prompt, duration_seconds=2,
-                   guidance_scale=1, steps=4,
-                   seed=42, randomize_seed=False,
+                   negative_prompt, duration_seconds,
+                   guidance_scale, steps, seed, randomize_seed,
                    progress=gr.Progress(track_tqdm=True)):
-
-    if pipe is None or not torch.cuda.is_available():
-        raise gr.Error("Pipeline not initialized or CUDA not available. Please check the console for errors.")
-
+    if pipe is None:
+        raise gr.Error("Pipeline not initialized. Check logs for GPU availability.")
     if input_image is None:
         raise gr.Error("Please upload an input image.")
 
     target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
     target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
-
     raw_frames = int(round(duration_seconds * FIXED_FPS))
     num_frames = ((raw_frames - 1) // 4) * 4 + 1
     num_frames = np.clip(num_frames, MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
 
-    if num_frames > 120:
-        max_dim = max(target_h, target_w)
-        if max_dim > 768:
-            scale_factor = 768 / max_dim
-            target_h = max(MOD_VALUE, (int(target_h * scale_factor) // MOD_VALUE) * MOD_VALUE)
-            target_w = max(MOD_VALUE, (int(target_w * scale_factor) // MOD_VALUE) * MOD_VALUE)
-            gr.Info(f"Reduced resolution to {target_w}x{target_h} for long video generation")
-
-    print(f"Generating {num_frames} frames (requested {raw_frames}) at {target_w}x{target_h}")
+    if num_frames > 120 and max(target_h, target_w) > 768:
+        scale_factor = 768 / max(target_h, target_w)
+        target_h = max(MOD_VALUE, int(target_h * scale_factor) // MOD_VALUE * MOD_VALUE)
+        target_w = max(MOD_VALUE, int(target_w * scale_factor) // MOD_VALUE * MOD_VALUE)
+        gr.Info(f"Reduced resolution to {target_w}x{target_h} for long video.")
 
     current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
     resized_image = input_image.resize((target_w, target_h), Image.Resampling.LANCZOS)
-
     torch.cuda.empty_cache()
 
     try:
-        with torch.inference_mode():
-            with torch.autocast("cuda", dtype=torch.float16):
-                output_frames_list = pipe(
-                    image=resized_image, prompt=prompt, negative_prompt=negative_prompt,
-                    height=target_h, width=target_w, num_frames=num_frames,
-                    guidance_scale=float(guidance_scale), num_inference_steps=int(steps),
-                    generator=torch.Generator(device="cuda").manual_seed(current_seed),
-                    return_dict=True
-                ).frames[0]
+        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
+            output_frames_list = pipe(
+                image=resized_image, prompt=prompt, negative_prompt=negative_prompt,
+                height=target_h, width=target_w, num_frames=num_frames,
+                guidance_scale=float(guidance_scale), num_inference_steps=int(steps),
+                generator=torch.Generator(device="cuda").manual_seed(current_seed)
+            ).frames[0]
     except torch.cuda.OutOfMemoryError:
+        raise gr.Error("Out of GPU memory. Try reducing duration or resolution.")
+    finally:
         torch.cuda.empty_cache()
-        raise gr.Error("Out of GPU memory. Try reducing the duration or resolution.")
-    except Exception as e:
-        torch.cuda.empty_cache()
-        raise gr.Error(f"Generation failed: {str(e)}")
-
-    torch.cuda.empty_cache()
 
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
         video_path = tmpfile.name
         export_video_with_ffmpeg(output_frames_list, video_path, fps=FIXED_FPS)
-
-        if os.path.exists(video_path) and os.path.getsize(video_path) > 0:
-            try:
-                subprocess.run(['ffmpeg', '-version'], capture_output=True, check=True)
-                optimized_path = video_path + "_opt.mp4"
-                cmd = [
-                    'ffmpeg', '-y', '-i', video_path, '-c:v', 'libx264', '-pix_fmt', 'yuv420p',
-                    '-profile:v', 'main', '-level', '4.0', '-movflags', '+faststart', '-crf', '23',
-                    '-preset', 'medium', '-maxrate', '10M', '-bufsize', '20M', optimized_path
-                ]
-                result = subprocess.run(cmd, capture_output=True, text=True)
-                if result.returncode == 0 and os.path.exists(optimized_path) and os.path.getsize(optimized_path) > 0:
-                    os.unlink(video_path)
-                    video_path = optimized_path
-                else:
-                    print(f"FFmpeg optimization failed: {result.stderr}")
-            except (subprocess.CalledProcessError, FileNotFoundError):
-                print("FFmpeg not available or optimization failed, using original export")
-
+        # Optional: FFmpeg optimization
+        # ...
     return video_path, current_seed
 
-# Gradio Interface
+# --- Gradio UI ---
 with gr.Blocks() as demo:
-    gr.Markdown("# Fast 4 steps Wan 2.1 I2V (14B) FusionX-LoRA")
-    gr.Markdown("Generate videos up to 10 seconds long! Longer videos may use reduced resolution for stability.")
-
+    gr.Markdown("# Wan 2.1 I2V FusionX-LoRA")
+    gr.Markdown("GPU is required. If this doesn't load, check your Space hardware settings.")
+    
     with gr.Row():
         with gr.Column():
-            input_image_component = gr.Image(type="pil", label="Input Image (auto-resized to target H/W)")
+            input_image_component = gr.Image(type="pil", label="Input Image")
             prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v)
-            duration_seconds_input = gr.Slider(
-                minimum=round(MIN_FRAMES_MODEL/FIXED_FPS, 1),
-                maximum=round(MAX_FRAMES_MODEL/FIXED_FPS, 1),
-                step=0.1, value=2, label="Duration (seconds)",
-                info=f"Video length: {MIN_FRAMES_MODEL/FIXED_FPS:.1f}-{MAX_FRAMES_MODEL/FIXED_FPS:.1f}s."
-            )
+            duration_seconds_input = gr.Slider(minimum=round(MIN_FRAMES_MODEL/FIXED_FPS, 1), maximum=round(MAX_FRAMES_MODEL/FIXED_FPS, 1), step=0.1, value=2, label="Duration (seconds)")
             with gr.Accordion("Advanced Settings", open=False):
                 negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
-                seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
-                randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
+                seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
+                randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True)
                 with gr.Row():
-                    height_input = gr.Slider(minimum=SLIDER_MIN_H, maximum=SLIDER_MAX_H, step=MOD_VALUE, value=DEFAULT_H_SLIDER_VALUE, label=f"Output Height (multiple of {MOD_VALUE})")
-                    width_input = gr.Slider(minimum=SLIDER_MIN_W, maximum=SLIDER_MAX_W, step=MOD_VALUE, value=DEFAULT_W_SLIDER_VALUE, label=f"Output Width (multiple of {MOD_VALUE})")
+                    height_input = gr.Slider(minimum=SLIDER_MIN_H, maximum=SLIDER_MAX_H, step=MOD_VALUE, value=DEFAULT_H_SLIDER_VALUE, label="Height")
+                    width_input = gr.Slider(minimum=SLIDER_MIN_W, maximum=SLIDER_MAX_W, step=MOD_VALUE, value=DEFAULT_W_SLIDER_VALUE, label="Width")
                 steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=4, label="Inference Steps")
                 guidance_scale_input = gr.Slider(minimum=0.0, maximum=20.0, step=0.5, value=1.0, label="Guidance Scale", visible=False)
-
-            generate_button = gr.Button("Generate Video", variant="primary")
+            generate_button = gr.Button("Generate Video", variant="primary", interactive=(pipe is not None))
         with gr.Column():
             video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
-            gr.Markdown("### Tips for best results:\n- For videos longer than 5 seconds, consider using lower resolutions (512-768px)\n- Clear, simple prompts often work better than complex descriptions\n- The model works best with 4-8 inference steps")
-
-    input_image_component.upload(
-        fn=handle_image_upload_for_dims_wan,
-        inputs=[input_image_component, height_input, width_input],
-        outputs=[height_input, width_input]
-    )
-    input_image_component.clear(
-        fn=handle_image_upload_for_dims_wan,
-        inputs=[input_image_component, height_input, width_input],
-        outputs=[height_input, width_input]
-    )
+            gr.Markdown("### Tips:\n- Longer videos need more memory.\n- 4-8 steps is optimal.")
 
-    ui_inputs = [
-        input_image_component, prompt_input, height_input, width_input,
-        negative_prompt_input, duration_seconds_input,
-        guidance_scale_input, steps_slider, seed_input, randomize_seed_checkbox
-    ]
-    generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
+    input_image_component.upload(fn=handle_image_upload_for_dims_wan, inputs=input_image_component, outputs=[height_input, width_input])
     
-    # The example images 'peng.png' and 'forg.jpg' are not present in this environment,
-    # so the gr.Examples component is commented out to prevent errors.
-    # gr.Examples(
-    #     examples=[
-    #         ["path/to/your/peng.png", "a penguin playfully dancing in the snow, Antarctica", 896, 512],
-    #         ["path/to/your/forg.jpg", "the frog jumps around", 448, 832],
-    #     ],
-    #     inputs=[input_image_component, prompt_input, height_input, width_input],
-    #     outputs=[video_output, seed_input],
-    #     fn=generate_video,
-    #     cache_examples="lazy"
-    # )
+    ui_inputs = [input_image_component, prompt_input, height_input, width_input, negative_prompt_input, duration_seconds_input, guidance_scale_input, steps_slider, seed_input, randomize_seed_checkbox]
+    generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
 
 if __name__ == "__main__":
     if pipe is not None:
         demo.queue(max_size=3).launch()
     else:
-        gr.Blocks()._queue_closed = False # A hack to prevent Gradio from hanging
-        gr.Info("Application not started because a GPU (CUDA) is required but not found.")
+        # This provides a clean message in the UI if the app can't start.
+        gr.Markdown("# Application Start Failed").launch()
+        gr.Info("A GPU is required to run this application. Please ensure your Hugging Face Space is configured with GPU hardware.")