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"
LORA_ID = "latent-consistency/lcm-lora-sdxl" 

pipe = None

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

    log = "System Start...\n"
    print("Lade Modell...")
    
    try:
        # Standard float32 für maximale CPU Kompatibilität
        pipe = I2VGenXLPipeline.from_pretrained(
            MODEL_ID, 
            torch_dtype=torch.float32, 
            variant="fp16"
        )
        
        log += "🚀 Lade LCM Turbo LoRA...\n"
        try:
            pipe.load_lora_weights(LORA_ID)
            pipe.fuse_lora()
            pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
            log += "✅ LCM Turbo ist AKTIV!\n"
        except Exception as e:
            log += f"⚠️ Turbo Fehler: {e}\n"

        # Versuche aggressives Speichermanagement
        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, resolution, progress=gr.Progress()):
    global pipe
    log_messages = ""
    
    if image_in is None:
        return None, "Kein Bild!"
    
    # Initialisierung des Balkens
    progress(0, desc="Lade Modell...")
    
    if pipe is None:
        model, msg = load_model_safely()
        log_messages += msg
        pipe = model
        
    if pipe is None:
        return None, log_messages + "Abbruch: Modell nicht geladen."

    gc.collect()
    
    try:
        # EINSTELLUNGEN
        steps = 6         
        guidance = 1.2    
        
        # Slider Wert übernehmen
        target_size = int(resolution)
        
        log_messages += f"Skaliere Bild auf {target_size}x{target_size}...\n"
        
        # Bild resizing (WICHTIG für Speed!)
        image_in = image_in.resize((target_size, target_size))
        
        log_messages += f"Starte Generierung ({steps} Steps)...\n"
        
        generator = torch.manual_seed(42)
        
        # --- FIX FÜR DEN FORTSCHRITTSBALKEN (Alte Methode) ---
        def callback_fn(step, timestep, latents):
            # step ist hier der Index (0, 1, 2...)
            current = step + 1
            progress((current, steps), desc=f"Step {current}/{steps} (Größe: {target_size}px)")
        # -----------------------------------------------------
        
        output = pipe(
            prompt=prompt,
            image=image_in,
            negative_prompt=negative_prompt,
            num_frames=16,          
            num_inference_steps=steps, 
            guidance_scale=guidance, 
            height=target_size, 
            width=target_size,  
            generator=generator,
            callback=callback_fn, # Alte Methode nutzen!
            callback_steps=1
        ).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:
        # Fehler abfangen und anzeigen
        err_msg = str(e)
        return None, log_messages + f"\n❌ Fehler: {err_msg}"

# UI
with gr.Blocks() as demo:
    gr.Markdown("# I2VGen-XL ⚡ LCM TURBO (Repariert)")
    gr.Markdown("Wähle die Auflösung passend zur Geduld. **384px oder 448px** empfohlen!")
    
    with gr.Row():
        with gr.Column():
            img = gr.Image(type="pil", label="Bild")
            txt = gr.Textbox(label="Prompt", value="fireworks in the sky")
            neg = gr.Textbox(value="distortion, blurry", label="Negative")
            
            # DER WICHTIGE REGLER FÜR SPEED
            resolution_slider = gr.Slider(
                minimum=320, 
                maximum=640, 
                value=448, 
                step=64, 
                label="Auflösung (Pixel)",
                info="384 = Schnell | 448 = Mittel | 512+ = Langsam"
            )
            
            btn = gr.Button("Turbo Start")
    
    with gr.Row():
        vid = gr.Video(label="Video")
        logs = gr.Textbox(label="Status Log", lines=10)
        
    btn.click(generate_video, [img, txt, neg, resolution_slider], [vid, logs])

demo.launch()