app_lora.py

import spaces
import torch
from diffusers import AutoencoderKLWan, WanImageToVideoPipeline, UniPCMultistepScheduler
from diffusers.utils import export_to_video
from transformers import CLIPVisionModel
import gradio as gr
import tempfile
import os
import subprocess

from huggingface_hub import hf_hub_download
import numpy as np
from PIL import Image
import random

import warnings
warnings.filterwarnings("ignore", message=".*Attempting to use legacy OpenCV backend.*")
warnings.filterwarnings("ignore", message=".*num_frames - 1.*")

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
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)
pipe = WanImageToVideoPipeline.from_pretrained(
    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()
pipe.enable_vae_slicing()
pipe.enable_vae_tiling()

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()

MOD_VALUE = 32
DEFAULT_H_SLIDER_VALUE = 640
DEFAULT_W_SLIDER_VALUE = 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)

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,
                                 default_h, default_w):
    orig_w, orig_h = pil_image.size
    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):
    if uploaded_pil_image is None:
        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, NEW_FORMULA_MAX_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:
        gr.Warning("Error attempting to calculate new dimensions")
        return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)

def get_duration(input_image, prompt, height, width, 
               negative_prompt, duration_seconds,
               guidance_scale, steps,
               seed, randomize_seed, 
               progress):
    # Adjust timeout based on video length and complexity
    if duration_seconds > 7:
        return 180  # 3 minutes for very long videos
    elif duration_seconds > 5:
        return 120  # 2 minutes for long videos
    elif duration_seconds > 3:
        return 90   # 1.5 minutes for medium videos
    else:
        return 60   # 1 minute for short videos

def export_video_with_ffmpeg(frames, output_path, fps=24):
    """Export video using imageio if available, otherwise fall back to OpenCV"""
    try:
        import imageio
        # Use imageio for better quality
        writer = imageio.get_writer(output_path, fps=fps, codec='libx264', 
                                   pixelformat='yuv420p', quality=8)
        for frame in frames:
            writer.append_data(np.array(frame))
        writer.close()
        return True
    except ImportError:
        # Fall back to OpenCV
        export_to_video(frames, output_path, fps=fps)
        return False

@spaces.GPU(duration=get_duration)
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, 
                   progress=gr.Progress(track_tqdm=True)):
    
    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)
    
    # Calculate frames with proper alignment
    raw_frames = int(round(duration_seconds * FIXED_FPS))
    # Ensure num_frames-1 is divisible by 4 as required by the model
    num_frames = ((raw_frames - 1) // 4) * 4 + 1
    num_frames = np.clip(num_frames, MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
    
    # Additional check for very long videos
    if num_frames > 120:
        # For videos longer than 5 seconds, reduce resolution to manage memory
        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}")
    
    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)

    # Clear GPU cache before generation
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    try:
        with torch.inference_mode():
            # Generate video with autocast for memory efficiency
            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]
    except torch.cuda.OutOfMemoryError:
        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)}")

    # Clear cache after generation
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
        video_path = tmpfile.name
        
        # Export using imageio if available, otherwise OpenCV
        used_imageio = export_video_with_ffmpeg(output_frames_list, video_path, fps=FIXED_FPS)
        
        # Only try FFmpeg optimization if we have a valid video file
        if os.path.exists(video_path) and os.path.getsize(video_path) > 0:
            try:
                # Check if ffmpeg is available
                subprocess.run(['ffmpeg', '-version'], capture_output=True, check=True)
                
                optimized_path = video_path + "_opt.mp4"
                cmd = [
                    'ffmpeg',
                    '-y',  # Overwrite without asking
                    '-i', video_path,  # Input file
                    '-c:v', 'libx264',  # Codec
                    '-pix_fmt', 'yuv420p',  # Pixel format
                    '-profile:v', 'main',  # Compatibility profile
                    '-level', '4.0',  # Support for higher resolutions
                    '-movflags', '+faststart',  # Streaming optimized
                    '-crf', '23',  # Quality level
                    '-preset', 'medium',  # Balance between speed and compression
                    '-maxrate', '10M',  # Max bitrate for large videos
                    '-bufsize', '20M',  # Buffer size
                    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)  # Remove original
                    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")

    return video_path, current_seed

# Gradio Interface
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.")
    
    with gr.Row():
        with gr.Column():
            input_image_component = gr.Image(type="pil", label="Input Image (auto-resized to target H/W)")
            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v)
            duration_seconds_input = gr.Slider(
                minimum=round(MIN_FRAMES_MODEL/FIXED_FPS, 1),  # 0.3s (8 frames)
                maximum=round(MAX_FRAMES_MODEL/FIXED_FPS, 1),  # 10.0s (240 frames)
                step=0.1, 
                value=2,  # Default 2 seconds
                label="Duration (seconds)", 
                info=f"Video length: {MIN_FRAMES_MODEL/FIXED_FPS:.1f}-{MAX_FRAMES_MODEL/FIXED_FPS:.1f}s. Longer videos may take more time and use more memory."
            )            
            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)
                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})")
                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")
        with gr.Column():
            video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
            gr.Markdown("### Tips for best results:")
            gr.Markdown("- For videos longer than 5 seconds, consider using lower resolutions (512-768px)")
            gr.Markdown("- Clear, simple prompts often work better than complex descriptions")
            gr.Markdown("- 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]
    )
    
    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])

    gr.Examples(
        examples=[ 
            ["peng.png", "a penguin playfully dancing in the snow, Antarctica", 896, 512],
            ["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"
    )

if __name__ == "__main__":
    demo.queue(max_size=3).launch()