import torch
from diffusers import UniPCMultistepScheduler
from diffusers import WanPipeline, AutoencoderKLWan
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
from huggingface_hub import hf_hub_download
from PIL import Image
import numpy as np
import gradio as gr
import spaces

device = "cuda" if torch.cuda.is_available() else "cpu"

# --- MODEL SETUP ---
model_id = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
flow_shift = 1.0
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
pipe.to(device)

# --- LORA SETUP ---
CAUSVID_NAME = "causvid_base"
PERSONVID_NAME = "personvid_optional"
ROSMX_NAME = "rosmx_optional"
OVIMX_NAME = "ovimx_optional"
OVIMX2_NAME = "ovimx2_optional"
OVIMX3_NAME = "ovimx3_optional"
PRTHE_NAME = "prthe_optional"

SUCCESSFULLY_LOADED_LORAS = {}

lora_definitions = {
    CAUSVID_NAME: ("joerose/Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32", "Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32.safetensors"),
    PERSONVID_NAME: ("ovi054/p3r5onVid1000", None),
    ROSMX_NAME: ("ovi054/rosmxVid1500", None),
    OVIMX_NAME: ("ovi054/ovimxVid1750", None),
    OVIMX2_NAME: ("ovi054/ovimxVid2250", None),
    OVIMX3_NAME: ("ovi054/ovimxVid2500", None),
    PRTHE_NAME: ("ovi054/prwthxVid", None)
}

for name, (repo, filename) in lora_definitions.items():
    print(f"Attempting to load LoRA '{name}'...")
    try:
        if filename:
            path = hf_hub_download(repo_id=repo, filename=filename)
            pipe.load_lora_weights(path, adapter_name=name, device_map="auto")
        else:
            pipe.load_lora_weights(repo, adapter_name=name, device_map="auto")
        print(f"✅ LoRA '{name}' loaded successfully.")
        SUCCESSFULLY_LOADED_LORAS[name] = repo
    except Exception as e:
        print(f"⚠️ LoRA '{name}' could not be loaded: {e}")

OPTIONAL_LORA_MAP = {
    "ovi054/p3r5onVid1000": PERSONVID_NAME,
    "ovi054/rosmxVid1500": ROSMX_NAME,
    "ovi054/ovimxVid1750": OVIMX_NAME,
    "ovi054/ovimxVid2250": OVIMX2_NAME,
    "ovi054/ovimxVid2500": OVIMX3_NAME,
    "ovi054/prwthxVid": PRTHE_NAME,
}
OPTIONAL_LORA_CHOICES = {k: v for k, v in OPTIONAL_LORA_MAP.items() if v in SUCCESSFULLY_LOADED_LORAS}

# At startup, disable all adapters. They will be selectively enabled during each run.
print("Disabling all LoRAs at startup. They will be activated on-demand.")
pipe.disable_lora()

print("Initialization complete. Gradio is starting...")

@spaces.GPU()
def generate(prompt, negative_prompt, width, height, num_inference_steps, optional_lora_id, progress=gr.Progress(track_tqdm=True)):
    
    # Using a try...finally block is robust for state management in apps.
    try:
        # --- Step 1: Build the list of ACTIVE adapters and their weights for THIS run ---
        
        active_adapters = []
        adapter_weights = []
        
        # Always include the base LoRA if it was loaded successfully
        if CAUSVID_NAME in SUCCESSFULLY_LOADED_LORAS:
            active_adapters.append(CAUSVID_NAME)
            adapter_weights.append(1.0)
            
        # If an optional LoRA is selected, add it to the list
        if optional_lora_id and optional_lora_id != "None":
            internal_name_to_add = OPTIONAL_LORA_CHOICES.get(optional_lora_id)
            if internal_name_to_add:
                active_adapters.append(internal_name_to_add)
                adapter_weights.append(1.0)

        # --- Step 2: Apply the adapters and weights for this run using the correct function ---
        if active_adapters:
            print(f"Activating adapters: {active_adapters} with weights: {adapter_weights}")
            # This is the correct, modern way to set adapters and their weights.
            pipe.set_adapters(active_adapters, adapter_weights=adapter_weights)
        else:
            print("No LoRAs are active for this run.")
            # ensure all are disabled if for some reason none were selected
            pipe.disable_lora()

        apply_cache_on_pipe(pipe)
        
        # --- Step 3: Run inference ---
        output = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=height,
            width=width,
            num_frames=1,
            num_inference_steps=num_inference_steps,
            guidance_scale=1.0, 
        )
        image = output.frames[0][0]
        image = (image * 255).astype(np.uint8)
        return Image.fromarray(image)

    finally:
        print("Disabling LoRAs after run to reset state.")
        pipe.disable_lora()

# --- Gradio Interface ---
iface = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Input prompt"),
        gr.Textbox(label="Negative prompt", value = "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"),
        gr.Slider(label="Width", minimum=480, maximum=1280, step=16, value=1024),
        gr.Slider(label="Height", minimum=480, maximum=1280, step=16, value=1024),
        gr.Slider(minimum=1, maximum=80, step=1, label="Inference Steps", value=10),
        gr.Textbox(
            label="Optional LoRA",
        )
    ],
    outputs=gr.Image(label="output"),
)

iface.launch(debug=True)