app.py

"""Wan 2.2 Space with ZeroGPU, Turbo 4-8 Step Distill, Multi-LoRA (Civitai + HF), Presets, Trigger Words, and 300-Hours Optimizations."""

from __future__ import annotations

import hashlib
import ipaddress
import json
import mimetypes
import os
import random
import re
import shutil
import socket
import tempfile
import time
import traceback
from functools import cache
from urllib.parse import urljoin, urlsplit

import spaces
import gradio as gr
import torch
from PIL import Image, ImageOps
from diffusers.utils import export_to_video
from safetensors import safe_open

# =========================================================================
# Monkey-patch Diffusers Wan LoRA converter for broad community compatibility
# =========================================================================
try:
    import diffusers.loaders.lora_conversion_utils as conv_utils
    import diffusers.loaders.lora_pipeline as lora_pipe

    _orig_wan_converter = getattr(conv_utils, "_convert_non_diffusers_wan_lora_to_diffusers", None)

    def patched_convert_non_diffusers_wan_lora_to_diffusers(original_state_dict):
        # 1. Normalize prefixes: community LoRAs often start directly with "blocks."
        # while diffusers internal converter expects "diffusion_model.blocks."
        prefixed_dict = {}
        for k, v in original_state_dict.items():
            if k.startswith("blocks."):
                prefixed_dict[f"diffusion_model.{k}"] = v
            elif k.startswith("transformer.blocks."):
                prefixed_dict[f"diffusion_model.{k[12:]}"] = v
            else:
                prefixed_dict[k] = v

        if _orig_wan_converter is not None:
            try:
                # Attempt standard conversion with normalized keys
                return _orig_wan_converter(prefixed_dict.copy())
            except Exception as err:
                print(f"[wan] standard converter note: {err}; applying robust fallback mapping", flush=True)

        # 2. Robust fallback mapping for any non-standard Wan LoRA format
        converted_dict = {}
        for k, v in prefixed_dict.items():
            new_k = k
            if new_k.startswith("diffusion_model."):
                new_k = new_k[len("diffusion_model."):]
            if not new_k.startswith("transformer."):
                new_k = f"transformer.{new_k}"
            converted_dict[new_k] = v
        return converted_dict

    if _orig_wan_converter is not None:
        conv_utils._convert_non_diffusers_wan_lora_to_diffusers = patched_convert_non_diffusers_wan_lora_to_diffusers
    if hasattr(lora_pipe, "_convert_non_diffusers_wan_lora_to_diffusers"):
        lora_pipe._convert_non_diffusers_wan_lora_to_diffusers = patched_convert_non_diffusers_wan_lora_to_diffusers
    print("[wan] LoRA converter patch successfully applied.", flush=True)
except Exception as patch_err:
    print(f"[wan] could not apply LoRA converter patch: {patch_err}", flush=True)

DEFAULT_CIVITAI_KEY = "50a9e1bd474c03b856070a7272d8015c"
DEFAULT_MODEL_REPO = os.environ.get("WAN_MODEL_REPO", "Wan-AI/Wan2.2-I2V-A14B-Diffusers")
GPU_SIZE = os.environ.get("WAN_GPU_SIZE", "xlarge")
MAX_GPU_DURATION = int(os.environ.get("WAN_MAX_GPU_DURATION", "300"))
OUTPUT_DIR = os.path.join(tempfile.gettempdir(), "wan-outputs")
LOCAL_LORAS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "loras")
LORA_MAX_BYTES = 2 * 1024**3
DEFAULT_FPS = 16

# =========================================================================
# Turbo Distillation LoRAs (4-Step & 8-Step Acceleration)
# =========================================================================
TURBO_PRESETS = {
    "⚡ Turbo 4-Passos (LightX2V Distill - 4 Steps Ultra Rápido)": {
        "wan2.2_i2v": ("lightx2v/Wan2.2-Distill-Loras", "wan2.2_i2v_A14b_high_noise_lora_rank64_lightx2v_4step_1022.safetensors"),
        "wan2.2_t2v": ("lightx2v/Wan2.2-Distill-Loras", "wan2.2_t2v_A14b_high_noise_lora_rank64_lightx2v_4step_1217.safetensors"),
        "wan2.1_i2v": ("lightx2v/Wan2.1-Distill-Loras", "wan2.1_i2v_lora_rank64_lightx2v_4step.safetensors"),
        "wan2.1_t2v": ("lightx2v/Wan2.1-Distill-Loras", "wan2.1_t2v_14b_lora_rank64_lightx2v_4step.safetensors"),
        "default_steps": 4,
        "default_cfg": 1.5,
        "description": "Destilação 4-passos LightX2V. Renderiza o vídeo em ~15-25 segundos mantendo alta qualidade!",
    },
    "⚡ Turbo 8-Passos (8 Steps Balanceado)": {
        "wan2.2_i2v": ("lightx2v/Wan2.2-Distill-Loras", "wan2.2_i2v_A14b_low_noise_lora_rank64_lightx2v_4step_1022.safetensors"),
        "wan2.2_t2v": ("lightx2v/Wan2.2-Distill-Loras", "wan2.2_t2v_A14b_low_noise_lora_rank64_lightx2v_4step_1217.safetensors"),
        "wan2.1_i2v": ("lightx2v/Wan2.1-Distill-Loras", "wan2.1_i2v_lora_rank64_lightx2v_4step.safetensors"),
        "wan2.1_t2v": ("lightx2v/Wan2.1-Distill-Loras", "wan2.1_t2v_14b_lora_rank64_lightx2v_4step.safetensors"),
        "default_steps": 8,
        "default_cfg": 2.0,
        "description": "Destilação 8-passos para equilíbrio perfeito entre velocidade e detalhes.",
    },
    "🚫 Desativado (Modo Normal - 15 a 25 Passos)": {
        "default_steps": 20,
        "default_cfg": 5.0,
        "description": "Modo padrão sem aceleração Turbo (requer 15-25 passos).",
    },
}

# =========================================================================
# Base NSFW Adapters (Unlocks adult generation on Wan Base Models)
# =========================================================================
NSFW_BASE_ADAPTERS = {
    "WAN General NSFW Model (Civitai 1307155 - Recomendado)": {
        "source": "https://civitai.com/models/1307155",
        "trigger_words": "nsfwsks, uncensored, naked, explicit",
        "default_strength": 1.0,
        "description": "Desbloqueia anatomia e geração explícita no modelo base Wan (Civitai Model 1307155).",
    },
    "HMNSFW AIO V2 Base (Civitai 2834417)": {
        "source": "https://civitai.com/api/download/models/3206518",
        "trigger_words": "hmmotion",
        "default_strength": 0.5,
        "description": "Adaptador anatômico geral HMNSFW V2 para Wan 2.1/2.2.",
    },
    "Desativado / Apenas Base SFW": {
        "source": "",
        "trigger_words": "",
        "default_strength": 0.0,
        "description": "Usa apenas o modelo base padrão sem adaptador NSFW adicional.",
    },
}

# =========================================================================
# Preset Catalog for Wan 2.2 LoRAs (Actions, Motions & Styles)
# =========================================================================
LORA_PRESETS = {
    "None / Desativado": {
        "type": "none",
        "trigger_words": "",
        "default_strength": 1.0,
        "description": "Nenhum LoRA selecionado neste slot.",
    },
    "WAN General NSFW model (Civitai 1307155)": {
        "type": "civitai",
        "source": "https://civitai.com/models/1307155",
        "trigger_words": "nsfwsks, uncensored, naked, explicit",
        "default_strength": 1.0,
        "description": "LoRA Geral NSFW para Wan 2.2 (Civitai 1307155).",
    },
    "HMNSFW AIO V2 / hmmotion (Wan 2.2)": {
        "type": "civitai",
        "source": "https://civitai.com/api/download/models/3206518",
        "trigger_words": "hmmotion, missionary, side, fast, third-person side view, medium shot.",
        "default_strength": 0.5,
        "description": "LoRA All-in-One de anatomia e movimento realista (Civitai 2834417 / 3206518). Use força <= 0.5 com prompts descritivos.",
    },
    "Icy Twerk Pro Max (Wan 2.2)": {
        "type": "civitai",
        "source": "https://civitai.com/api/download/models/3201584",
        "trigger_words": "icytw3rk, twerking, booty shake, rhythmic hip movement, dynamic motion, bouncing buttocks, high quality",
        "default_strength": 0.9,
        "description": "LoRA de animação e movimento de twerk / booty shake para Wan (Civitai 2836640 / 3201584).",
    },
    "Cumouf - Oral Creampie / CIM with Spasms": {
        "type": "civitai",
        "source": "https://civitai.com/api/download/models/3223411",
        "trigger_words": "cum in mouth, oral creampie, cum overflow, spasms, throat bulge, choking on cum, messy facial, open mouth",
        "default_strength": 0.85,
        "description": "Oral creampie com espasmos faciais e garganta (Civitai 2846978 / 3223411).",
    },
    "Epic Cumshots & Facials": {
        "type": "civitai",
        "source": "https://civitai.com/api/download/models/3052864",
        "trigger_words": "cumshot, thick semen, facial, climax, sticky ejaculation, messy dripping, high viscosity",
        "default_strength": 0.9,
        "description": "Ejaculação realista de alta viscosidade com respingos faciais e corporais.",
    },
    "Dynamic Cinematic Camera Motion": {
        "type": "prompt_only",
        "trigger_words": "dynamic cinematic camera, slow orbit shot, dramatic lighting, sweeping drone view, motion blur",
        "default_strength": 0.85,
        "description": "Movimento de câmera fluído e cinematográfico.",
    },
    "Cyberpunk / Sci-Fi Neon Realism": {
        "type": "prompt_only",
        "trigger_words": "cyberpunk, holographic HUD, volumetric neon reflections, cybernetic glow, futuristic city, 8k cinematic",
        "default_strength": 0.9,
        "description": "Estilo cyberpunk hiper-detalhado com iluminação volumétrica e neons.",
    },
}

# 300 Hours Civitai Guide: Strict multiples-of-16 resolutions
CANVASES = {
    # 16:9 Landscape
    "832x480 · 16:9 Landscape (Fast 480p)": (480, 832),
    "960x544 · 16:9 Landscape (Balanced)": (544, 960),
    "1280x720 · 16:9 Landscape (HD 720p)": (720, 1280),
    # 9:16 Portrait / Reels
    "480x832 · 9:16 Portrait (Fast 480p)": (832, 480),
    "544x960 · 9:16 Portrait (Balanced)": (960, 544),
    "720x1280 · 9:16 Portrait (HD 720p)": (1280, 720),
    # 1:1 Square
    "640x640 · 1:1 Square (Fast)": (640, 640),
    "768x768 · 1:1 Square (HD)": (768, 768),
    # 4:3 / 3:4
    "768x576 · 4:3 Standard": (576, 768),
    "576x768 · 3:4 Portrait": (768, 576),
    # 21:9 Ultrawide
    "1152x512 · 21:9 Ultrawide": (512, 1152),
}
DEFAULT_CANVAS = "832x480 · 16:9 Landscape (Fast 480p)"

PIPE = None
CURRENT_MODEL_REPO = None
LOAD_ERROR: str | None = None
LOADED_IN: float | None = None


def get_local_loras() -> list[str]:
    """Scans the local loras/ folder for .safetensors files."""
    if not os.path.exists(LOCAL_LORAS_DIR):
        try:
            os.makedirs(LOCAL_LORAS_DIR, exist_ok=True)
        except Exception:
            return []
    files = [f for f in os.listdir(LOCAL_LORAS_DIR) if f.endswith(".safetensors")]
    return sorted(files)


def normalize_civitai_url(url: str) -> str:
    """Extracts direct download link from any Civitai model or version URL."""
    url = url.strip()
    if not url:
        return url
    
    match_version = re.search(r"modelVersionId=(\d+)", url)
    if match_version:
        version_id = match_version.group(1)
        return f"https://civitai.com/api/download/models/{version_id}"
    
    if "api/download/models/" in url:
        return re.sub(r"https?://[^/]+", "https://civitai.com", url)

    match_model = re.search(r"civitai\.(?:com|red|org|blue|work)/models/(\d+)", url)
    if match_model:
        model_id = match_model.group(1)
        try:
            import requests
            r = requests.get(f"https://civitai.com/api/v1/models/{model_id}", timeout=8)
            if r.ok:
                data = r.json()
                versions = data.get("modelVersions", [])
                if versions and "id" in versions[0]:
                    return f"https://civitai.com/api/download/models/{versions[0]['id']}"
        except Exception as err:
            print(f"[civitai] failed to resolve model {model_id} metadata: {err}", flush=True)
    return url


def resolve_canvas(value: str) -> str:
    canvas = str(value).strip()
    if canvas in CANVASES:
        return canvas
    return DEFAULT_CANVAS


def _sha256(path: str) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _lora_cache_directory(source: str) -> str:
    root = os.path.join(tempfile.gettempdir(), "wan-lora-downloads")
    cache_key = hashlib.sha256(source.encode()).hexdigest()[:20]
    local_dir = os.path.join(root, cache_key)
    os.makedirs(local_dir, exist_ok=True)
    others = sorted(
        (entry for entry in os.scandir(root) if entry.is_dir() and entry.path != local_dir),
        key=lambda entry: entry.stat().st_mtime,
        reverse=True,
    )
    for stale in others[4:]:
        shutil.rmtree(stale.path, ignore_errors=True)
    return local_dir


def _validate_public_lora_url(url: str) -> str:
    if len(url) > 2048:
        raise ValueError("Direct LoRA URL is too long.")
    parsed = urlsplit(url)
    if parsed.scheme.lower() != "https" or not parsed.hostname:
        raise ValueError("Direct LoRA URLs must use public HTTPS.")
    if parsed.username or parsed.password or parsed.port not in (None, 443):
        raise ValueError("Direct LoRA URLs cannot contain credentials or non-standard ports.")
    try:
        addresses = {item[4][0] for item in socket.getaddrinfo(parsed.hostname, 443, type=socket.SOCK_STREAM)}
    except socket.gaierror as error:
        raise ValueError("Direct LoRA URL hostname could not be resolved.") from error
    for raw_address in addresses:
        address = ipaddress.ip_address(raw_address)
        if isinstance(address, ipaddress.IPv6Address) and address.ipv4_mapped is not None:
            address = address.ipv4_mapped
        if not address.is_global:
            raise ValueError("Direct LoRA URLs cannot access private or local networks.")
    return url


def _download_lora_url(url: str, civitai_token: str = "") -> tuple[str, str]:
    import requests

    token = (civitai_token or os.environ.get("CIVITAI_API_KEY", "") or os.environ.get("CIVITAI_TOKEN", "") or DEFAULT_CIVITAI_KEY).strip()

    if "civitai." in url and token and "token=" not in url:
        sep = "&" if "?" in url else "?"
        url = f"{url}{sep}token={token}"

    original = _validate_public_lora_url(url)
    local_dir = _lora_cache_directory(original)
    path = os.path.join(local_dir, "adapter.safetensors")
    if os.path.isfile(path) and 0 < os.path.getsize(path) <= LORA_MAX_BYTES:
        try:
            with safe_open(path, framework="pt", device="cpu") as handle:
                if handle.keys():
                    os.utime(local_dir, None)
                    parsed = urlsplit(original)
                    return path, f"{parsed.hostname}{parsed.path}"[:180]
        except Exception:
            os.unlink(path)

    temporary = path + ".download"
    current = original
    for _ in range(6):
        current = _validate_public_lora_url(current)
        parsed_current = urlsplit(current)
        req_headers = {"User-Agent": "Wan-Studio-Space/1.0"}
        if token and ("civitai.com" in parsed_current.netloc or "civitai.red" in parsed_current.netloc):
            req_headers["Authorization"] = f"Bearer {token}"

        try:
            with requests.get(
                current,
                stream=True,
                allow_redirects=False,
                timeout=(15, 300),
                headers=req_headers,
            ) as response:
                if response.is_redirect or response.is_permanent_redirect:
                    location = response.headers.get("location")
                    if not location:
                        raise ValueError("Direct LoRA URL returned an empty redirect.")
                    current = urljoin(current, location)
                    continue
                if response.status_code in (401, 403):
                    raise gr.Error(
                        "🔒 O Civitai bloqueou o download deste modelo (401 Unauthorized / NSFW). "
                        "Verifique sua Civitai API Key nas configurações."
                    )
                response.raise_for_status()
                total = 0
                with open(temporary, "wb") as output:
                    for chunk in response.iter_content(1024 * 1024):
                        if not chunk:
                            continue
                        total += len(chunk)
                        if total > LORA_MAX_BYTES:
                            raise ValueError("Direct LoRA exceeds the 2 GiB safety limit.")
                        output.write(chunk)
                with safe_open(temporary, framework="pt", device="cpu") as handle:
                    if not handle.keys():
                        raise ValueError("Direct LoRA contains no safetensors tensors.")
                os.replace(temporary, path)
                os.utime(local_dir, None)
                parsed = urlsplit(original)
                return path, f"{parsed.hostname}{parsed.path}"[:180]
        except requests.exceptions.HTTPError as err:
            if "response" in locals() and response.status_code in (401, 403):
                raise gr.Error(
                    "🔒 O Civitai bloqueou o download deste modelo (401 Unauthorized / NSFW). "
                    "Verifique sua Civitai API Key nas configurações."
                ) from err
            raise
    raise ValueError("Too many redirects downloading LoRA.")


def resolve_turbo_lora(turbo_choice: str, model_repo: str) -> tuple[str | None, float]:
    """Resolves and downloads the appropriate 4-step or 8-step distillation LoRA."""
    if turbo_choice not in TURBO_PRESETS or "Desativado" in turbo_choice:
        return None, 0.0

    spec = TURBO_PRESETS[turbo_choice]
    is_i2v = "I2V" in model_repo or "i2v" in model_repo or "TI2V" in model_repo
    is_wan22 = "Wan2.2" in model_repo or "wan2.2" in model_repo

    if is_wan22:
        key = "wan2.2_i2v" if is_i2v else "wan2.2_t2v"
    else:
        key = "wan2.1_i2v" if is_i2v else "wan2.1_t2v"

    if key not in spec:
        return None, 0.0

    repo_id, filename = spec[key]
    from huggingface_hub import hf_hub_download
    local_dir = _lora_cache_directory(f"hf://{repo_id}/{filename}")
    path = hf_hub_download(repo_id=repo_id, filename=filename, token=False, local_dir=local_dir)
    os.utime(local_dir, None)
    return path, 1.0


def resolve_single_lora(
    preset_type: str, custom_url: str, hf_repo: str, hf_file: str, local_file: str, strength: float, civitai_token: str = ""
) -> tuple[str | None, str, float]:
    """Resolves one LoRA file path, label and scale for Wan 2.2."""
    if float(strength) == 0.0 or preset_type in ("None / Desativado", "None", ""):
        return None, "None", 0.0

    if preset_type in LORA_PRESETS:
        spec = LORA_PRESETS[preset_type]
        if spec.get("type") == "hf":
            from huggingface_hub import hf_hub_download
            repo_id = spec["hf_repo"]
            filename = spec["hf_file"]
            local_dir = _lora_cache_directory(f"hf://{repo_id}/{filename}")
            path = hf_hub_download(repo_id=repo_id, filename=filename, token=False, local_dir=local_dir)
            os.utime(local_dir, None)
            return path, preset_type, strength

        if spec.get("source"):
            url = normalize_civitai_url(spec["source"])
            if url:
                path, label = _download_lora_url(url, civitai_token)
                return path, preset_type, strength
        return None, "None", 0.0

    if preset_type == "Custom URL / Civitai":
        url = normalize_civitai_url(custom_url)
        if not url:
            return None, "None", 0.0
        path, label = _download_lora_url(url, civitai_token)
        return path, f"URL: {label}", strength

    if preset_type == "Custom Hugging Face":
        repo_id = str(hf_repo or "").strip()
        filename = str(hf_file or "").strip()
        if not repo_id or not filename:
            return None, "None", 0.0
        if not re.fullmatch(r"[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+", repo_id):
            raise ValueError("Custom LoRA repo must be in `owner/repository` format.")
        if not filename.endswith(".safetensors"):
            raise ValueError("Custom LoRA file must be a `.safetensors` file.")

        from huggingface_hub import get_hf_file_metadata, hf_hub_download, hf_hub_url
        metadata = get_hf_file_metadata(hf_hub_url(repo_id, filename), token=False)
        if metadata.size is None or metadata.size > LORA_MAX_BYTES:
            raise ValueError("LoRA file exceeds the 2 GiB limit.")
        local_dir = _lora_cache_directory(f"hf://{repo_id}/{filename}")
        path = hf_hub_download(repo_id=repo_id, filename=filename, token=False, local_dir=local_dir)
        os.utime(local_dir, None)
        return path, f"{repo_id}/{filename}", strength

    if preset_type == "Local File (loras/ folder)":
        if not local_file:
            return None, "None", 0.0
        local_path = os.path.join(LOCAL_LORAS_DIR, local_file)
        if not os.path.isfile(local_path):
            raise ValueError(f"Arquivo local {local_file} não encontrado na pasta loras/.")
        return local_path, f"local:{local_file}", strength

    return None, "None", 0.0


def get_or_load_pipeline(model_repo: str = DEFAULT_MODEL_REPO):
    global PIPE, CURRENT_MODEL_REPO, LOAD_ERROR, LOADED_IN

    if PIPE is not None and CURRENT_MODEL_REPO == model_repo:
        return PIPE

    started = time.time()
    try:
        from diffusers import DiffusionPipeline, WanImageToVideoPipeline, WanPipeline

        print(f"[wan] loading pipeline from {model_repo} ...", flush=True)
        if "I2V" in model_repo or "i2v" in model_repo or "TI2V" in model_repo:
            try:
                pipe = WanImageToVideoPipeline.from_pretrained(
                    model_repo,
                    torch_dtype=torch.bfloat16,
                )
            except Exception:
                pipe = DiffusionPipeline.from_pretrained(
                    model_repo,
                    torch_dtype=torch.bfloat16,
                )
        else:
            try:
                pipe = WanPipeline.from_pretrained(
                    model_repo,
                    torch_dtype=torch.bfloat16,
                )
            except Exception:
                pipe = DiffusionPipeline.from_pretrained(
                    model_repo,
                    torch_dtype=torch.bfloat16,
                )

        PIPE = pipe
        CURRENT_MODEL_REPO = model_repo
        LOADED_IN = time.time() - started
        print(f"[wan] ready in {LOADED_IN:.0f}s on CPU", flush=True)
    except Exception as error:
        traceback.print_exc()
        LOAD_ERROR = f"**Loading `{model_repo}` failed**: `{type(error).__name__}: {error}`"
        raise RuntimeError(LOAD_ERROR) from error
    return PIPE


def _fit_keyframe(image_input, target_width: int, target_height: int) -> Image.Image:
    if isinstance(image_input, str):
        img = Image.open(image_input)
    elif isinstance(image_input, Image.Image):
        img = image_input
    else:
        raise ValueError("Invalid image input")

    img = ImageOps.exif_transpose(img).convert("RGB")
    target_aspect = target_width / target_height
    img_aspect = img.width / img.height

    if abs(img_aspect - target_aspect) > 1e-3:
        if img_aspect > target_aspect:
            new_w = int(img.height * target_aspect)
            left = (img.width - new_w) // 2
            img = img.crop((left, 0, left + new_w, img.height))
        else:
            new_h = int(img.width / target_aspect)
            top = (img.height - new_h) // 2
            img = img.crop((0, top, img.width, top + new_h))

    img = img.resize((target_width, target_height), Image.Resampling.LANCZOS)
    return img


def get_duration(*args, **kwargs):
    try:
        steps = kwargs.get("steps", args[6] if len(args) > 6 else 4)
        return max(180, min(MAX_GPU_DURATION, int(steps) * 8 + 60))
    except Exception:
        return 240


@spaces.GPU(duration=get_duration, size=GPU_SIZE)
def _generate_video_gpu(
    prompt: str,
    negative_prompt: str,
    image_input: Image.Image | None,
    height: int,
    width: int,
    num_frames: int,
    steps: int,
    guidance_scale: float,
    seed: int,
    lora_configs: list[tuple[str, float]],
    model_repo: str,
):
    pipe = get_or_load_pipeline(model_repo)
    pipe.to("cuda")

    active_lora_names = []
    if lora_configs:
        try:
            pipe.unload_lora_weights()
        except Exception:
            pass

        for lora_path, lora_scale in lora_configs:
            if lora_path and lora_scale > 0:
                adapter_name = f"lora_{len(active_lora_names)}"
                try:
                    pipe.load_lora_weights(lora_path, adapter_name=adapter_name)
                    active_lora_names.append(adapter_name)
                except Exception as lora_err:
                    print(f"[wan] warning: skipping incompatible LoRA `{lora_path}`: {lora_err}", flush=True)

        if active_lora_names:
            scales = [scale for _, scale in lora_configs if scale > 0][:len(active_lora_names)]
            pipe.set_adapters(active_lora_names, adapter_weights=scales)

    generator = torch.Generator("cuda").manual_seed(int(seed))

    try:
        with torch.inference_mode():
            if image_input is not None and ("I2V" in model_repo or "i2v" in model_repo or "TI2V" in model_repo):
                output = pipe(
                    image=image_input,
                    prompt=prompt,
                    negative_prompt=negative_prompt if negative_prompt else None,
                    height=height,
                    width=width,
                    num_frames=num_frames,
                    num_inference_steps=int(steps),
                    guidance_scale=float(guidance_scale),
                    generator=generator,
                )
            else:
                output = pipe(
                    prompt=prompt,
                    negative_prompt=negative_prompt if negative_prompt else None,
                    height=height,
                    width=width,
                    num_frames=num_frames,
                    num_inference_steps=int(steps),
                    guidance_scale=float(guidance_scale),
                    generator=generator,
                )
        frames = output.frames[0]
    finally:
        if active_lora_names:
            try:
                pipe.unload_lora_weights()
            except Exception:
                pass

    return frames


def generate_video(
    prompt: str,
    negative_prompt: str,
    input_image: Image.Image | str | None,
    turbo_choice: str,
    nsfw_base_choice: str,
    nsfw_base_strength: float,
    canvas: str,
    num_frames: int,
    fps: int,
    steps: int,
    guidance_scale: float,
    seed: int,
    randomize_seed: bool,
    model_choice: str,
    lora1_preset: str,
    lora1_custom_url: str,
    lora1_hf_repo: str,
    lora1_hf_file: str,
    lora1_local_file: str,
    lora1_strength: float,
    lora2_preset: str,
    lora2_custom_url: str,
    lora2_hf_repo: str,
    lora2_hf_file: str,
    lora2_local_file: str,
    lora2_strength: float,
    civitai_api_key: str,
    progress=gr.Progress(track_tqdm=True),
):
    if not prompt or not prompt.strip():
        raise gr.Error("Por favor, digite um prompt descrevendo a cena do vídeo.")

    if randomize_seed:
        seed = random.randint(0, 2147483647)

    canvas = resolve_canvas(canvas)
    height, width = CANVASES[canvas]

    processed_image = None
    if input_image is not None:
        processed_image = _fit_keyframe(input_image, width, height)

    token_to_use = (civitai_api_key or DEFAULT_CIVITAI_KEY).strip()
    lora_configs = []
    active_labels = []

    # 0. Turbo LoRA (4-Step or 8-Step Distill)
    t_path, t_scale = resolve_turbo_lora(turbo_choice, model_choice)
    if t_path and t_scale > 0:
        lora_configs.append((t_path, t_scale))
        active_labels.append(f"⚡ Turbo: {turbo_choice.split('(')[0].strip()}")

    # 1. Base NSFW Adapter
    if nsfw_base_choice in NSFW_BASE_ADAPTERS and nsfw_base_strength > 0:
        base_spec = NSFW_BASE_ADAPTERS[nsfw_base_choice]
        if base_spec.get("source"):
            b_url = normalize_civitai_url(base_spec["source"])
            if b_url:
                b_path, b_label = _download_lora_url(b_url, token_to_use)
                lora_configs.append((b_path, float(nsfw_base_strength)))
                active_labels.append(f"🔞 NSFW Base: {nsfw_base_choice.split('(')[0].strip()} (@ {nsfw_base_strength:g})")

    # 2. Slot 1 LoRA
    l1_path, l1_label, l1_scale = resolve_single_lora(
        lora1_preset, lora1_custom_url, lora1_hf_repo, lora1_hf_file, lora1_local_file, lora1_strength, token_to_use
    )
    if l1_path and l1_scale > 0:
        lora_configs.append((l1_path, l1_scale))
        active_labels.append(f"{l1_label} (@ {lora1_strength:g})")

    # 3. Slot 2 LoRA
    l2_path, l2_label, l2_scale = resolve_single_lora(
        lora2_preset, lora2_custom_url, lora2_hf_repo, lora2_hf_file, lora2_local_file, lora2_strength, token_to_use
    )
    if l2_path and l2_scale > 0:
        lora_configs.append((l2_path, l2_scale))
        active_labels.append(f"{l2_label} (@ {lora2_strength:g})")

    progress(0.1, desc=f"Gerando {steps} passos a {width}x{height} ({num_frames} frames)...")
    started = time.time()

    frames = _generate_video_gpu(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image_input=processed_image,
        height=height,
        width=width,
        num_frames=int(num_frames),
        steps=int(steps),
        guidance_scale=float(guidance_scale),
        seed=int(seed),
        lora_configs=lora_configs,
        model_repo=model_choice,
    )

    gen_time = time.time() - started
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    out_video_path = os.path.join(OUTPUT_DIR, f"wan_{int(time.time() * 1000)}.mp4")
    export_to_video(frames, out_video_path, fps=int(fps))

    loras_str = " + ".join(active_labels) if active_labels else "None (Base Model)"
    report = (
        f"**Modelo**: `{model_choice.split('/')[-1]}` | **Resolução**: `{width}x{height}` (Múltiplo de 16) | "
        f"**Frames**: {num_frames} ({num_frames / fps:.2f}s @ {fps}fps) | **Passos**: {steps} | **Seed**: {seed}\n\n"
        f"🎯 **Active LoRAs**: `{loras_str}`\n\n"
        f"⏱️ **Tempo de Renderização**: {gen_time:.1f}s"
    )

    return out_video_path, report, seed


# =========================================================================
# Gradio UI Interface
# =========================================================================

custom_css = """
.gradio-container {
    max-width: 1350px !important;
    margin: 0 auto !important;
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Helvetica, Arial, sans-serif;
}
.header-card {
    background: linear-gradient(135deg, #1e1b4b 0%, #3b0764 50%, #0f172a 100%);
    border: 1px solid rgba(168, 85, 247, 0.3);
    border-radius: 16px;
    padding: 24px 32px;
    margin-bottom: 20px;
    box-shadow: 0 10px 25px -5px rgba(0, 0, 0, 0.4);
}
.header-card h1 {
    font-size: 2.2rem;
    font-weight: 800;
    margin: 0 0 8px 0;
    background: linear-gradient(90deg, #c084fc, #38bdf8, #818cf8);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}
.badge {
    display: inline-block;
    padding: 4px 10px;
    border-radius: 9999px;
    font-size: 0.8rem;
    font-weight: 600;
    margin-right: 6px;
}
.badge-turbo {
    background: rgba(245, 158, 11, 0.2);
    color: #fbbf24;
    border: 1px solid rgba(245, 158, 11, 0.4);
}
.badge-zerogpu {
    background: rgba(16, 185, 129, 0.2);
    color: #34d399;
    border: 1px solid rgba(16, 185, 129, 0.4);
}
.badge-model {
    background: rgba(99, 102, 241, 0.2);
    color: #a5b4fc;
    border: 1px solid rgba(99, 102, 241, 0.4);
}
.badge-multilora {
    background: rgba(236, 72, 153, 0.2);
    color: #f472b6;
    border: 1px solid rgba(236, 72, 153, 0.4);
}
.trigger-btn {
    background: rgba(168, 85, 247, 0.2) !important;
    border: 1px solid rgba(168, 85, 247, 0.5) !important;
    color: #e9d5ff !important;
    font-size: 0.85rem !important;
    font-weight: 600 !important;
    padding: 4px 12px !important;
    border-radius: 6px !important;
}
.generate-btn {
    background: linear-gradient(135deg, #9333ea 0%, #4f46e5 100%) !important;
    color: white !important;
    font-weight: 700 !important;
    font-size: 1.15rem !important;
    border-radius: 12px !important;
    padding: 12px 24px !important;
    box-shadow: 0 4px 15px rgba(147, 51, 234, 0.4) !important;
    border: none !important;
    transition: all 0.2s ease !important;
}
.generate-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(147, 51, 234, 0.6) !important;
}
"""

all_preset_choices = (
    list(LORA_PRESETS.keys())
    + ["Custom URL / Civitai", "Custom Hugging Face"]
    + (["Local File (loras/ folder)"] if get_local_loras() else [])
)

with gr.Blocks(css=custom_css, title="Wan 2.2 Turbo Video Studio") as app:
    gr.HTML(
        """
        <div class="header-card">
            <div style="margin-bottom: 12px;">
                <span class="badge badge-turbo">⚡ Turbo 4-Step Distill</span>
                <span class="badge badge-zerogpu">⚡ ZeroGPU</span>
                <span class="badge badge-model">🎬 Wan 2.2 (MoE Diffusers)</span>
                <span class="badge badge-multilora">🧩 Base NSFW + Multi-LoRA Engine</span>
            </div>
            <h1>Wan 2.2 Turbo Video Studio</h1>
            <p>Aceleração Turbo de 4 a 8 Passos com Arquitetura Wan 2.2 MoE, Adaptador Base NSFW e Multi-LoRA.</p>
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=6):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Descreva a cena detalhada do vídeo...",
                lines=4,
                value="nsfwsks, a cinematic shot of a beautiful woman with glowing eyes, dramatic lighting, 8k masterpiece",
            )

            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="Elementos indesejados (distorções, membros extras, baixa qualidade)...",
                lines=2,
                value="deformed, bad anatomy, extra limbs, blurry, low resolution, bad quality, distortion, censorship",
            )

            with gr.Accordion("📖 Guia & Modelos de Prompt HMNSFW / hmmotion (Wan 2.2)", open=False):
                gr.Markdown(
                    """
                    **Estrutura Recomendada pelo Guia de 300 Horas / Autor do LoRA (`hmmotion`)**:
                    - **Força recomendada**: `0.3` a `0.5` (use `<= 0.5`).
                    - **Cabeçalho Obrigatório**: `hmmotion, <class>, <viewpoint>, <pace>, <shot>.`
                      - *Class*: `missionary` / `cowgirl` / `blowjob` / `doggy` / `handjob` / `insertion`
                      - *Viewpoint*: `pov` ou `side`
                      - *Pace*: `fast` ou `slow`
                      - *Shot*: `close-up` / `medium shot` / `third-person side view` / `high-angle downward shot`
                    - **Dica**: Escreva um parágrafo contínuo descritivo de 180 a 260 palavras detalhando a pose, anatomia, movimento (*"The motion is..."*), superfícies (*"sheen"*) e áudio (*"The audio consists of..."*).
                    """
                )
                hmnsfw_example_btn = gr.Button("💡 Inserir Exemplo Completo de Prompt HMNSFW", elem_classes=["trigger-btn"])

            with gr.Tabs():
                with gr.TabItem("⚡ Aceleração Turbo (4-8 Passos)"):
                    gr.Markdown("Acelera a geração em até **8x** aplicando destilação LightX2V. Permite renderizar vídeos em ~15 segundos!")
                    turbo_choice = gr.Dropdown(
                        label="Modo Turbo",
                        choices=list(TURBO_PRESETS.keys()),
                        value="⚡ Turbo 4-Passos (LightX2V Distill - 4 Steps Ultra Rápido)",
                    )

                with gr.TabItem("🔞 Adaptador Base NSFW (Essencial)"):
                    gr.Markdown("Como o modelo base do Wan é treinado sem conteúdo explícito, este adaptador é carregado como base para desbloquear anatomia sem gastar seus slots de LoRA!")
                    with gr.Row():
                        nsfw_base_choice = gr.Dropdown(
                            label="Adaptador Base NSFW",
                            choices=list(NSFW_BASE_ADAPTERS.keys()),
                            value="WAN General NSFW Model (Civitai 1307155 - Recomendado)",
                        )
                        nsfw_base_strength = gr.Slider(
                            label="Força do Adaptador Base NSFW",
                            minimum=0.0,
                            maximum=1.5,
                            step=0.05,
                            value=1.0,
                        )
                    nsfw_base_trigger_btn = gr.Button("📋 Inserir Trigger `nsfwsks` no Prompt", elem_classes=["trigger-btn"])

                with gr.TabItem("🖼️ Image-to-Video (I2V)"):
                    gr.Markdown("Faça upload de uma imagem inicial (`First Frame`). O sistema ajustará o corte automaticamente para a proporção múltipla de 16 selecionada.")
                    input_image = gr.Image(label="First Frame (Imagem Inicial)", type="pil")

                with gr.TabItem("🎨 LoRA Slot 1 (Ação / Pose)"):
                    lora1_preset = gr.Dropdown(
                        label="LoRA Slot 1 Preset / Fonte",
                        choices=all_preset_choices,
                        value="HMNSFW AIO V2 / hmmotion (Wan 2.2)",
                    )
                    with gr.Group(visible=False) as lora1_custom_url_grp:
                        lora1_custom_url = gr.Textbox(
                            label="URL de Download do Civitai / SafeTensor",
                            placeholder="https://civitai.red/models/... ou https://civitai.com/api/download/models/...",
                        )
                    with gr.Group(visible=False) as lora1_hf_grp:
                        with gr.Row():
                            lora1_hf_repo = gr.Textbox(label="HF Repo ID", placeholder="owner/repo")
                            lora1_hf_file = gr.Textbox(label="HF Filename", placeholder="model.safetensors")
                    with gr.Group(visible=False) as lora1_local_grp:
                        lora1_local_file = gr.Dropdown(label="Arquivo Local (pasta loras/)", choices=get_local_loras())

                    lora1_trigger_display = gr.Textbox(
                        label="Trigger Words",
                        value=LORA_PRESETS.get("HMNSFW AIO V2 / hmmotion (Wan 2.2)", {}).get("trigger_words", ""),
                        interactive=False,
                    )
                    add_lora1_triggers_btn = gr.Button("📋 Inserir Trigger Words no Prompt", elem_classes=["trigger-btn"])

                    lora1_strength = gr.Slider(
                        label="LoRA 1 Força (Scale)",
                        minimum=0.0,
                        maximum=2.0,
                        step=0.05,
                        value=0.5,
                    )

                with gr.TabItem("🎬 LoRA Slot 2 (Movimento / Câmera)"):
                    lora2_preset = gr.Dropdown(
                        label="LoRA Slot 2 Preset / Fonte",
                        choices=all_preset_choices,
                        value="None / Desativado",
                    )
                    with gr.Group(visible=False) as lora2_custom_url_grp:
                        lora2_custom_url = gr.Textbox(
                            label="URL de Download do Civitai / SafeTensor",
                            placeholder="https://civitai.com/api/download/models/...",
                        )
                    with gr.Group(visible=False) as lora2_hf_grp:
                        with gr.Row():
                            lora2_hf_repo = gr.Textbox(label="HF Repo ID", placeholder="owner/repo")
                            lora2_hf_file = gr.Textbox(label="HF Filename", placeholder="model.safetensors")
                    with gr.Group(visible=False) as lora2_local_grp:
                        lora2_local_file = gr.Dropdown(label="Arquivo Local (pasta loras/)", choices=get_local_loras())

                    lora2_trigger_display = gr.Textbox(
                        label="Trigger Words",
                        value="",
                        interactive=False,
                    )
                    add_lora2_triggers_btn = gr.Button("📋 Inserir Trigger Words no Prompt", elem_classes=["trigger-btn"])

                    lora2_strength = gr.Slider(
                        label="LoRA 2 Força (Scale)",
                        minimum=0.0,
                        maximum=2.0,
                        step=0.05,
                        value=0.85,
                    )

            with gr.Accordion("⚙️ Configurações de Resolução & Renderização (Guia 300h)", open=True):
                with gr.Row():
                    model_choice = gr.Dropdown(
                        label="Modelo Base Wan",
                        choices=[
                            ("Wan 2.2 I2V A14B (MoE Image-to-Video - Recomendado)", "Wan-AI/Wan2.2-I2V-A14B-Diffusers"),
                            ("Wan 2.2 TI2V 5B (Híbrido Texto & Imagem)", "Wan-AI/Wan2.2-TI2V-5B-Diffusers"),
                            ("Wan 2.2 T2V A14B (MoE Texto para Vídeo)", "Wan-AI/Wan2.2-T2V-A14B-Diffusers"),
                            ("Wan 2.1 I2V 14B 480P (Versão Anterior)", "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"),
                            ("Wan 2.1 I2V 14B 720P (Versão Anterior HD)", "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers"),
                            ("Wan 2.1 T2V 14B (Versão Anterior T2V)", "Wan-AI/Wan2.1-T2V-14B-Diffusers"),
                        ],
                        value="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
                    )
                    canvas = gr.Dropdown(
                        label="Proporção & Resolução (Múltiplos de 16)",
                        choices=list(CANVASES.keys()),
                        value=DEFAULT_CANVAS,
                    )

                with gr.Row():
                    num_frames = gr.Slider(
                        label="Número de Frames",
                        minimum=17,
                        maximum=81,
                        step=16,
                        value=81,
                        info="81 frames = ~5 segundos a 16 fps; 49 frames = ~3 segundos.",
                    )
                    fps = gr.Slider(
                        label="FPS do Vídeo",
                        minimum=8,
                        maximum=30,
                        step=1,
                        value=16,
                    )
                    steps = gr.Slider(
                        label="Passos de Inferência (Steps)",
                        minimum=2,
                        maximum=40,
                        step=1,
                        value=4,
                        info="Com Turbo 4-Passos ativado, 4 passos fornecem render ultrarrápido.",
                    )

                with gr.Row():
                    guidance_scale = gr.Slider(
                        label="Guidance Scale (CFG)",
                        minimum=1.0,
                        maximum=10.0,
                        step=0.5,
                        value=1.5,
                        info="No modo Turbo, recomenda-se CFG baixo (1.0 a 2.0).",
                    )
                    seed = gr.Number(label="Seed", value=42, precision=0)
                    randomize_seed = gr.Checkbox(label="🎲 Randomizar Seed", value=True)

                civitai_api_key = gr.Textbox(
                    label="🔑 Civitai API Key",
                    value=DEFAULT_CIVITAI_KEY,
                    placeholder="Chave de API do Civitai configurada",
                    type="password",
                )

            generate_btn = gr.Button("🚀 Gerar Vídeo Wan 2.2 Turbo", variant="primary", elem_classes=["generate-btn"])

        with gr.Column(scale=6):
            output_video = gr.Video(label="Vídeo Gerado", autoplay=True, loop=True)
            output_report = gr.Markdown(label="Detalhes da Geração", value="Pronto para renderizar vídeo.")

    # Event handlers
    def on_turbo_change(turbo_val):
        spec = TURBO_PRESETS.get(turbo_val, {})
        return spec.get("default_steps", 4), spec.get("default_cfg", 1.5)

    turbo_choice.change(
        fn=on_turbo_change,
        inputs=[turbo_choice],
        outputs=[steps, guidance_scale],
    )

    def on_lora1_change(preset_val):
        spec = LORA_PRESETS.get(preset_val, {})
        triggers = spec.get("trigger_words", "")
        default_s = spec.get("default_strength", 1.0)
        return (
            gr.update(visible=preset_val == "Custom URL / Civitai"),
            gr.update(visible=preset_val == "Custom Hugging Face"),
            gr.update(visible=preset_val == "Local File (loras/ folder)"),
            triggers,
            default_s,
        )

    lora1_preset.change(
        fn=on_lora1_change,
        inputs=[lora1_preset],
        outputs=[lora1_custom_url_grp, lora1_hf_grp, lora1_local_grp, lora1_trigger_display, lora1_strength],
    )

    def on_lora2_change(preset_val):
        spec = LORA_PRESETS.get(preset_val, {})
        triggers = spec.get("trigger_words", "")
        default_s = spec.get("default_strength", 1.0)
        return (
            gr.update(visible=preset_val == "Custom URL / Civitai"),
            gr.update(visible=preset_val == "Custom Hugging Face"),
            gr.update(visible=preset_val == "Local File (loras/ folder)"),
            triggers,
            default_s,
        )

    lora2_preset.change(
        fn=on_lora2_change,
        inputs=[lora2_preset],
        outputs=[lora2_custom_url_grp, lora2_hf_grp, lora2_local_grp, lora2_trigger_display, lora2_strength],
    )

    def append_triggers(curr_prompt, triggers):
        if not triggers:
            return curr_prompt
        curr = curr_prompt.strip()
        if not curr:
            return triggers
        if triggers.lower() in curr.lower():
            return curr
        return f"{curr}, {triggers}"

    def load_hmnsfw_example():
        return (
            "hmmotion, missionary, side, fast, third-person side view, medium shot. "
            "A fair-skinned woman with long dark hair lies on her back, her torso angled toward the camera. "
            "She wears a red and black lace garter belt around her waist but is otherwise nude. "
            "Her left leg is raised and bent while her right leg is spread wide. "
            "The man is positioned above her, his torso and arms visible as he thrusts. "
            "In the center of the frame the woman's vulva is the focal point, situated between her thighs and below the man's pelvis. "
            "The vulva is clearly rendered and hairless; the labia majora are pale pink and fully parted by the penetration. "
            "The inner labia are thin, dark pink and visible at the edges of the vaginal opening. "
            "The clitoral hood is visible and flushed. "
            "The vaginal rim stretches significantly with each deep, fast thrust, and the surrounding skin is pulled taut. "
            "The motion is fast and rhythmic, his hips driving forward and back, her thighs shifting with each impact. "
            "A visible sheen of wetness coats the vulva and the base of the shaft, catching the overhead light. "
            "His hands grip her raised thigh, holding her leg open. Her head is tilted back with her mouth open. "
            "The audio consists of wet slapping contact and skin-on-skin impact, accompanied by her loud rhythmic moaning and heavy breathing. "
            "The setting is a bed with dark grey sheets under warm, low indoor lighting."
        )

    hmnsfw_example_btn.click(fn=load_hmnsfw_example, outputs=[prompt])
    nsfw_base_trigger_btn.click(fn=lambda p: append_triggers(p, "nsfwsks, uncensored, explicit"), inputs=[prompt], outputs=[prompt])
    add_lora1_triggers_btn.click(fn=append_triggers, inputs=[prompt, lora1_trigger_display], outputs=[prompt])
    add_lora2_triggers_btn.click(fn=append_triggers, inputs=[prompt, lora2_trigger_display], outputs=[prompt])

    generate_btn.click(
        fn=generate_video,
        inputs=[
            prompt,
            negative_prompt,
            input_image,
            turbo_choice,
            nsfw_base_choice,
            nsfw_base_strength,
            canvas,
            num_frames,
            fps,
            steps,
            guidance_scale,
            seed,
            randomize_seed,
            model_choice,
            lora1_preset,
            lora1_custom_url,
            lora1_hf_repo,
            lora1_hf_file,
            lora1_local_file,
            lora1_strength,
            lora2_preset,
            lora2_custom_url,
            lora2_hf_repo,
            lora2_hf_file,
            lora2_local_file,
            lora2_strength,
            civitai_api_key,
        ],
        outputs=[output_video, output_report, seed],
    )

if __name__ == "__main__":
    try:
        print("[wan] pre-initializing default pipeline on CPU...", flush=True)
        get_or_load_pipeline(DEFAULT_MODEL_REPO)
    except Exception as err:
        print(f"[wan] initial preload deferred to first call: {err}", flush=True)
    app.launch(show_error=True, allowed_paths=[OUTPUT_DIR])