app.py

import gradio as gr
import torch
from diffusers import I2VGenXLPipeline, LCMScheduler
from diffusers.utils import export_to_video
import gc
import traceback

MODEL_ID = "ali-vilab/i2vgen-xl"
# Das ist der Turbo-Booster für SDXL-basierte Modelle (wie I2VGen-XL)
LORA_ID = "latent-consistency/lcm-lora-sdxl" 

pipe = None

def load_model_safely():
    global pipe
    if pipe is not None:
        return pipe, "Modell bereits geladen."

    log = "Lade Modell...\n"
    print("Starte Ladevorgang...")
    
    try:
        # 1. Basis Modell laden
        pipe = I2VGenXLPipeline.from_pretrained(
            MODEL_ID, 
            torch_dtype=torch.float32, 
            variant="fp16"
        )
        
        # 2. LCM TURBO ZÜNDEN
        log += "🚀 Lade LCM LoRA...\n"
        try:
            # LoRA laden
            pipe.load_lora_weights(LORA_ID)
            
            # WICHTIG: fuse_lora() verschmilzt die Gewichte. 
            # Das verhindert, dass wir doppelten RAM für Modell + LoRA brauchen.
            pipe.fuse_lora()
            
            # Scheduler austauschen gegen LCM
            pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
            log += "✅ LCM Turbo aktiviert! (Steps reduziert auf 4-8)\n"
        except Exception as e:
            log += f"⚠️ LoRA Fehler: {e}\nEs geht ohne Turbo weiter (langsam).\n"

        # 3. Speicher-Optimierung
        try:
            pipe.enable_model_cpu_offload()
            log += "✅ Model Offloading aktiv.\n"
        except:
            try:
                pipe.enable_sequential_cpu_offload()
                log += "✅ Sequential Offloading aktiv.\n"
            except:
                log += "❌ RAM Warnung: Kein Offloading.\n"

        pipe.enable_vae_slicing()
        pipe.enable_vae_tiling()
        
        return pipe, log
        
    except Exception as e:
        return None, f"Absturz beim Laden: {e}\n{traceback.format_exc()}"

def generate_video(image_in, prompt, negative_prompt):
    global pipe
    log_messages = ""
    
    if image_in is None:
        return None, "Kein Bild!"
    
    if pipe is None:
        model, msg = load_model_safely()
        log_messages += msg
        pipe = model
        
    if pipe is None:
        return None, log_messages + "Abbruch."

    gc.collect()
    
    try:
        # LCM EINSTELLUNGEN
        # Steps: 6 (statt 20-50)
        # Guidance: 1.5 (LCM braucht sehr niedrige Werte, sonst Bildmatsch)
        steps = 6
        guidance = 1.5 
        
        # Auflösung: Wir bleiben bei 448px für Stabilität
        target_size = 448
        
        log_messages += f"Generiere mit {steps} Steps (Turbo Modus)...\n"
        
        image_in = image_in.resize((target_size, target_size))
        
        generator = torch.manual_seed(42)
        
        output = pipe(
            prompt=prompt,
            image=image_in,
            negative_prompt=negative_prompt,
            num_frames=16,          
            num_inference_steps=steps, 
            guidance_scale=guidance, # WICHTIG: Niedrig halten bei LCM!
            height=target_size, 
            width=target_size,  
            generator=generator
        ).frames[0]
        
        video_path = "turbo_output.mp4"
        export_to_video(output, video_path, fps=8)
        
        log_messages += "✅ FERTIG!"
        return video_path, log_messages
        
    except Exception as e:
        return None, log_messages + f"\n❌ Fehler: {e}"

with gr.Blocks() as demo:
    gr.Markdown("# I2VGen-XL ⚡ LCM TURBO")
    gr.Markdown("Nutzt LCM LoRA für extreme Geschwindigkeit (6 Steps).")
    
    with gr.Row():
        with gr.Column():
            img = gr.Image(type="pil", label="Bild")
            txt = gr.Textbox(label="Prompt", value="clouds moving, cinematic")
            neg = gr.Textbox(value="distortion, blurry", label="Negative")
            btn = gr.Button("Turbo Start (ca. 2-3 Min)")
    
    with gr.Row():
        vid = gr.Video(label="Video")
        logs = gr.Textbox(label="Status", lines=8)
        
    btn.click(generate_video, [img, txt, neg], [vid, logs])

demo.launch()