import torch 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 # Assuming this is for Hugging Face Spaces GPU decorator from huggingface_hub import hf_hub_download 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-I2V-14B-480P-Diffusers" LORA_REPO_ID = "Kijai/WanVideo_comfy" LORA_FILENAME = "Wan21_CausVid_14B_T2V_lora_rank32.safetensors" # Configure logging 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 ) logger.info(f"Loading Pipeline {MODEL_ID}...") pipe = WanImageToVideoPipeline.from_pretrained( MODEL_ID, vae=vae, image_encoder=image_encoder, torch_dtype=torch.bfloat16 ) flow_shift = 8.0 pipe.scheduler = UniPCMultistepScheduler.from_config( pipe.scheduler.config, flow_shift=flow_shift ) logger.info("Moving pipeline to CUDA...") pipe.to("cuda") # --- LoRA Loading --- 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...") 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 # 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" 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}") target_height = int(height) target_width = int(width) num_frames = int(num_frames) 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=target_height, width=target_width, num_frames=num_frames, 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_val) # For export logger.info(f"Video successfully generated and saved to {video_path}") return video_path # --- Gradio UI Definition --- 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""" # 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): 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", 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=inputs_for_click_and_examples, outputs=video_output ) gr.Examples( examples=[ [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=inputs_for_click_and_examples, outputs=video_output, fn=generate_video, cache_examples=False ) if __name__ == "__main__": demo.queue().launch(share=True, debug=True)