Update app.py

app.py

@@ -1,17 +1,18 @@
 import torch
-from diffusers import AutoencoderKLWan, WanPipeline, UniPCMultistepScheduler
-from diffusers.utils import export_to_video
-from diffusers.loaders.lora_conversion_utils import _convert_non_diffusers_wan_lora_to_diffusers # Keep this if it's the base for standard LoRA parts
+from diffusers import AutoencoderKLWan, WanImageToVideoPipeline, UniPCMultistepScheduler
+from diffusers.utils import export_to_video, load_image
+from transformers import CLIPVisionModel
 import gradio as gr
 import tempfile
 import os
-import spaces
+import spaces # Assuming this is for Hugging Face Spaces GPU decorator
 from huggingface_hub import hf_hub_download
-import logging # For better logging
-import re # For key manipulation
+import logging
+import numpy as np
+from PIL import Image # Added for type hinting
 
 # --- Global Model Loading & LoRA Handling ---
-MODEL_ID = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
+MODEL_ID = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
 LORA_REPO_ID = "Kijai/WanVideo_comfy"
 LORA_FILENAME = "Wan21_CausVid_14B_T2V_lora_rank32.safetensors"
 
@@ -20,17 +21,25 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # --- Model Loading ---
+logger.info(f"Loading Image Encoder for {MODEL_ID}...")
+image_encoder = CLIPVisionModel.from_pretrained(
+    MODEL_ID,
+    subfolder="image_encoder",
+    torch_dtype=torch.float32
+)
+
 logger.info(f"Loading VAE for {MODEL_ID}...")
 vae = AutoencoderKLWan.from_pretrained(
     MODEL_ID,
     subfolder="vae",
-    torch_dtype=torch.float32 # float32 for VAE stability
+    torch_dtype=torch.float32
 )
 logger.info(f"Loading Pipeline {MODEL_ID}...")
-pipe = WanPipeline.from_pretrained(
+pipe = WanImageToVideoPipeline.from_pretrained(
     MODEL_ID,
     vae=vae,
-    torch_dtype=torch.bfloat16 # bfloat16 for pipeline
+    image_encoder=image_encoder,
+    torch_dtype=torch.bfloat16
 )
 flow_shift = 8.0
 pipe.scheduler = UniPCMultistepScheduler.from_config(
@@ -43,99 +52,202 @@ pipe.to("cuda")
 logger.info(f"Downloading LoRA {LORA_FILENAME} from {LORA_REPO_ID}...")
 causvid_path = hf_hub_download(repo_id=LORA_REPO_ID, filename=LORA_FILENAME)
 
-logger.info("Loading LoRA weights with custom converter...")
-pipe.load_lora_weights(causvid_path,adapter_name="causvid_lora")
+logger.info("Loading LoRA weights...")
+pipe.load_lora_weights(causvid_path, adapter_name="causvid_lora")
+logger.info("Setting LoRA adapter...")
+pipe.set_adapters(["causvid_lora"], adapter_weights=[1.0])
+
+# MOD_VALUE for height/width constraints
+# From WanImageToVideoPipeline docs: height/width must be multiple of vae_scale_factor * transformer.config.patch_size[1 or 2]
+MOD_VALUE = pipe.vae_scale_factor * pipe.transformer.config.patch_size[1] # e.g., 8 * 16 = 128
+logger.info(f"Derived MOD_VALUE for dimensions: {MOD_VALUE}")
+
+
+# --- Helper functions and constants for automatic dimension adjustment ---
+# These constants must match the Gradio slider definitions below
+DEFAULT_H_SLIDER_VALUE = 384
+DEFAULT_W_SLIDER_VALUE = 640
+DEFAULT_TARGET_AREA = float(DEFAULT_H_SLIDER_VALUE * DEFAULT_W_SLIDER_VALUE)
+
+SLIDER_MIN_H = 128
+SLIDER_MAX_H = 512
+SLIDER_MIN_W = 128
+SLIDER_MAX_W = 1024
+
+def _calculate_new_dimensions_wan(pil_image: Image.Image, mod_val: int, target_area: float,
+                                 min_h: int, max_h: int, min_w: int, max_w: int,
+                                 default_h: int, default_w: int) -> tuple[int, int]:
+    orig_w, orig_h = pil_image.size
+
+    if orig_w == 0 or orig_h == 0:
+        logger.warning("Uploaded image has zero width or height. Using default slider dimensions.")
+        return default_h, default_w
+
+    aspect_ratio = orig_h / orig_w
+
+    # Calculate ideal dimensions for the target area, maintaining aspect ratio
+    ideal_h = np.sqrt(target_area * aspect_ratio)
+    ideal_w = np.sqrt(target_area / aspect_ratio)
+
+    # Round to nearest multiple of mod_val
+    calc_h = round(ideal_h / mod_val) * mod_val
+    calc_w = round(ideal_w / mod_val) * mod_val
+
+    # Ensure dimensions are at least mod_val (smallest valid multiple)
+    calc_h = mod_val if calc_h == 0 else calc_h
+    calc_w = mod_val if calc_w == 0 else calc_w
+
+    # Clamp to slider limits
+    new_h = int(np.clip(calc_h, min_h, max_h))
+    new_w = int(np.clip(calc_w, min_w, max_w))
+
+    logger.info(f"Auto-dim: Original {orig_w}x{orig_h} (AR: {aspect_ratio:.2f}). Target Area: {target_area}.")
+    logger.info(f"Auto-dim: Ideal HxW: {ideal_h:.0f}x{ideal_w:.0f}. Rounded (step {mod_val}): {calc_h}x{calc_w}.")
+    logger.info(f"Auto-dim: Clamped HxW: {new_h}x{new_w} (H_range:[{min_h}-{max_h}], W_range:[{min_w}-{max_w}]).")
+
+    return new_h, new_w
+
+def handle_image_upload_for_dims_wan(uploaded_pil_image: Image.Image | None, current_h_val: int, current_w_val: int):
+    if uploaded_pil_image is None: # Image cleared by user
+        logger.info("Image cleared. Resetting dimensions to default slider values.")
+        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
+
+    try:
+        new_h, new_w = _calculate_new_dimensions_wan(
+            uploaded_pil_image,
+            MOD_VALUE,
+            DEFAULT_TARGET_AREA,
+            SLIDER_MIN_H, SLIDER_MAX_H,
+            SLIDER_MIN_W, SLIDER_MAX_W,
+            DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE
+        )
+        return gr.update(value=new_h), gr.update(value=new_w)
+    except Exception as e:
+        logger.error(f"Error auto-adjusting H/W from image: {e}", exc_info=True)
+        # On error, revert to defaults or keep current. Defaults are safer.
+        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
 
 
 # --- Gradio Interface Function ---
-@spaces.GPU
-def generate_video(prompt, negative_prompt, height, width, num_frames, guidance_scale, steps, fps, progress=gr.Progress(track_tqdm=True)):
+@spaces.GPU # type: ignore
+def generate_video(input_image: Image.Image, prompt: str, negative_prompt: str,
+                   height: int, width: int, num_frames: int,
+                   guidance_scale: float, steps: int, fps_for_conditioning_and_export: int,
+                   progress=gr.Progress(track_tqdm=True)):
+    if input_image is None:
+        raise gr.Error("Please upload an input image.")
+
     logger.info("Starting video generation...")
+    logger.info(f"  Input Image: Uploaded (Original size: {input_image.size if input_image else 'N/A'})")
     logger.info(f"  Prompt: {prompt}")
     logger.info(f"  Negative Prompt: {negative_prompt if negative_prompt else 'None'}")
-    logger.info(f"  Height: {height}, Width: {width}")
-    logger.info(f"  Num Frames: {num_frames}, FPS: {fps}")
-    logger.info(f"  Guidance Scale: {guidance_scale}")
+    logger.info(f"  Target Output Height: {height}, Target Output Width: {width}")
+    logger.info(f"  Num Frames: {num_frames}, FPS for conditioning & export: {fps_for_conditioning_and_export}")
+    logger.info(f"  Guidance Scale: {guidance_scale}, Steps: {steps}")
+
 
-    height = (int(height) // 8) * 8
-    width = (int(width) // 8) * 8
+    target_height = int(height)
+    target_width = int(width)
     num_frames = int(num_frames)
-    fps = int(fps)
+    fps_val = int(fps_for_conditioning_and_export)
+    guidance_scale_val = float(guidance_scale)
+    steps_val = int(steps)
+
+    # Resize the input PIL image to the target dimensions for the pipeline
+    resized_image = input_image.resize((target_width, target_height))
+    logger.info(f"  Input image resized to: {resized_image.size} for pipeline input.")
+
 
     with torch.inference_mode():
         output_frames_list = pipe(
+            image=resized_image,
             prompt=prompt,
             negative_prompt=negative_prompt,
-            height=height,
-            width=width,
+            height=target_height,
+            width=target_width,
             num_frames=num_frames,
-            guidance_scale=float(guidance_scale),
-            num_inference_steps=steps
+            guidance_scale=guidance_scale_val,
+            num_inference_steps=steps_val,
+            fps=fps_val, # For conditioning
+            generator=torch.Generator(device="cuda").manual_seed(0) # For reproducibility
         ).frames[0]
 
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
         video_path = tmpfile.name
 
-    export_to_video(output_frames_list, video_path, fps=fps)
+    export_to_video(output_frames_list, video_path, fps=fps_val) # For export
     logger.info(f"Video successfully generated and saved to {video_path}")
     return video_path
 
 # --- Gradio UI Definition ---
-default_prompt = "A cat walks on the grass, realistic"
-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"
+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"
+penguin_image_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png"
+
 
 with gr.Blocks() as demo:
     gr.Markdown(f"""
-    # Text-to-Video with Wan 2.1 (14B) + CausVid LoRA
-    Powered by `diffusers` and `Wan-AI/{MODEL_ID}`.
+    # Image-to-Video with Wan 2.1 I2V (14B) + CausVid LoRA
+    Powered by `diffusers` and `{MODEL_ID}`.
     Model is loaded into memory when the app starts. This might take a few minutes.
     Ensure you have a GPU with sufficient VRAM (e.g., ~24GB+ for these default settings).
+    Output Height and Width must be multiples of **{MOD_VALUE}**. Uploading an image will suggest dimensions based on its aspect ratio and a target area.
     """)
     with gr.Row():
         with gr.Column(scale=2):
-            prompt_input = gr.Textbox(label="Prompt", value=default_prompt, lines=3)
-            negative_prompt_input = gr.Textbox(
-                label="Negative Prompt (Optional)",
-                value=default_negative_prompt,
-                lines=3
-            )
-            with gr.Row():
-                height_input = gr.Slider(minimum=256, maximum=768, step=64, value=480, label="Height (multiple of 8)")
-                width_input = gr.Slider(minimum=256, maximum=1024, step=64, value=832, label="Width (multiple of 8)")
-            with gr.Row():
-                num_frames_input = gr.Slider(minimum=16, maximum=100, step=1, value=25, label="Number of Frames")
-                fps_input = gr.Slider(minimum=5, maximum=30, step=1, value=15, label="Output FPS")
-            steps = gr.Slider(minimum=1.0, maximum=30.0, value=4.0, label="Steps")
-            guidance_scale_input = gr.Slider(minimum=1.0, maximum=20.0, step=0.5, value=1.0, label="Guidance Scale")
+            input_image_component = gr.Image(type="pil", label="Input Image (will be resized to target H/W)")
+            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v, lines=3)
+
+            with gr.Accordion("Advanced Settings", open=False):
+                negative_prompt_input = gr.Textbox(
+                    label="Negative Prompt (Optional)",
+                    value=default_negative_prompt,
+                    lines=3
+                )
+                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})")
+                with gr.Row():
+                    num_frames_input = gr.Slider(minimum=8, maximum=81, step=1, value=25, label="Number of Frames")
+                    fps_input = gr.Slider(minimum=5, maximum=30, step=1, value=16, label="FPS (for conditioning & export)")
+                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")
 
             generate_button = gr.Button("Generate Video", variant="primary")
 
         with gr.Column(scale=3):
-            video_output = gr.Video(label="Generated Video")
+            video_output = gr.Video(label="Generated Video", interactive=False)
+
+    # Event handler for image upload/clear to adjust H/W sliders
+    input_image_component.change(
+        fn=handle_image_upload_for_dims_wan,
+        inputs=[input_image_component, height_input, width_input], # Pass current slider values
+        outputs=[height_input, width_input]
+    )
+
+    inputs_for_click_and_examples = [
+        input_image_component,
+        prompt_input,
+        negative_prompt_input,
+        height_input,
+        width_input,
+        num_frames_input,
+        guidance_scale_input,
+        steps_slider,
+        fps_input
+    ]
 
     generate_button.click(
         fn=generate_video,
-        inputs=[
-            prompt_input,
-            negative_prompt_input,
-            height_input,
-            width_input,
-            num_frames_input,
-            guidance_scale_input,
-            steps,
-            fps_input
-        ],
+        inputs=inputs_for_click_and_examples,
         outputs=video_output
     )
 
     gr.Examples(
         examples=[
-            ["A panda eating bamboo in a lush forest, cinematic lighting", default_negative_prompt, 480, 832, 25, 5.0, 4, 15],
-            ["A majestic eagle soaring over snowy mountains", default_negative_prompt, 512, 768, 30, 7.0, 4, 12],
-            ["Timelapse of a flower blooming, vibrant colors", "static, ugly", 384, 640, 40, 6.0, 4, 20],
-            ["Astronaut walking on the moon, Earth in the background, highly detailed", default_negative_prompt, 480, 832, 20, 5.5, 4, 10],
+            [penguin_image_url, "a penguin playfully dancing in the snow, Antarctica", default_negative_prompt, DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE, 25, 1.0, 4, 16]
         ],
-        inputs=[prompt_input, negative_prompt_input, height_input, width_input, num_frames_input, guidance_scale_input, steps, fps_input],
+        inputs=inputs_for_click_and_examples,
         outputs=video_output,
         fn=generate_video,
         cache_examples=False