app.py

import os; os.system('pip install --upgrade --no-deps spaces')
import spaces
import shutil
import subprocess
import sys
import copy
import random
import tempfile
import warnings
import time
import gc
import uuid
from tqdm import tqdm
import cv2
import numpy as np
import torch
import torch._dynamo
from torch.nn import functional as F
from PIL import Image

import gradio as gr
from diffusers import (
    FlowMatchEulerDiscreteScheduler,
    SASolverScheduler,
    DEISMultistepScheduler,
    DPMSolverMultistepInverseScheduler,
    UniPCMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
)
from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
from diffusers.utils.export_utils import export_to_video

from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig, Int8WeightOnlyConfig
import aoti
import lora_loader

os.environ["TOKENIZERS_PARALLELISM"] = "true"
warnings.filterwarnings("ignore")
IS_ZERO_GPU = bool(os.getenv("SPACES_ZERO_GPU"))

# --- FRAME EXTRACTION JS & LOGIC ---

get_timestamp_js = """
function() {
    const video = document.querySelector('#generated-video video');

    if (video) {
        console.log("Video found! Time: " + video.currentTime);
        return video.currentTime;
    } else {
        console.log("No video element found.");
        return 0;
    }
}
"""


get_timestamp_js_last = get_timestamp_js


def extract_frame(video_path, timestamp):
    if not video_path:
        return None

    print(f"Extracting frame at timestamp: {timestamp}")

    cap = cv2.VideoCapture(video_path)

    if not cap.isOpened():
        return None

    fps = cap.get(cv2.CAP_PROP_FPS)
    target_frame_num = int(float(timestamp) * fps)

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if target_frame_num >= total_frames:
        target_frame_num = total_frames - 1

    cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame_num)
    ret, frame = cap.read()
    cap.release()

    if ret:
        return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    return None

# --- END FRAME EXTRACTION LOGIC ---


def clear_vram():
    gc.collect()
    torch.cuda.empty_cache()


# RIFE
if not os.path.exists("RIFEv4.26_0921.zip"):
    print("Downloading RIFE Model...")
    subprocess.run([
        "wget", "-q",
        "https://huggingface.co/thornmaze/RIFE/resolve/main/RIFEv4.26_0921.zip",
        "-O", "RIFEv4.26_0921.zip"
    ], check=True)
    subprocess.run(["unzip", "-o", "RIFEv4.26_0921.zip"], check=True)

from train_log.RIFE_HDv3 import Model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
rife_model = Model()
rife_model.load_model("train_log", -1)
rife_model.eval()


@torch.no_grad()
def interpolate_bits(frames_np, multiplier=2, scale=1.0):
    if isinstance(frames_np, list):
        T = len(frames_np)
        H, W, C = frames_np[0].shape
    else:
        T, H, W, C = frames_np.shape

    if multiplier < 2:
        if isinstance(frames_np, np.ndarray):
            return list(frames_np)
        return frames_np

    n_interp = multiplier - 1

    tmp = max(128, int(128 / scale))
    ph = ((H - 1) // tmp + 1) * tmp
    pw = ((W - 1) // tmp + 1) * tmp
    padding = (0, pw - W, 0, ph - H)

    def to_tensor(frame_np):
        t = torch.from_numpy(frame_np).to(device)
        t = t.permute(2, 0, 1).unsqueeze(0)
        return F.pad(t, padding).half()

    def from_tensor(tensor):
        t = tensor[0, :, :H, :W]
        t = t.permute(1, 2, 0)
        return t.float().cpu().numpy()

    def make_inference(I0, I1, n):
        if rife_model.version >= 3.9:
            res = []
            for i in range(n):
                res.append(rife_model.inference(I0, I1, (i+1) * 1. / (n+1), scale))
            return res
        else:
            middle = rife_model.inference(I0, I1, scale)
            if n == 1:
                return [middle]
            first_half = make_inference(I0, middle, n=n//2)
            second_half = make_inference(middle, I1, n=n//2)
            if n % 2:
                return [*first_half, middle, *second_half]
            else:
                return [*first_half, *second_half]

    output_frames = []
    I1 = to_tensor(frames_np[0])
    total_steps = T - 1

    with tqdm(total=total_steps, desc="Interpolating", unit="frame") as pbar:
        for i in range(total_steps):
            I0 = I1
            output_frames.append(from_tensor(I0))
            I1 = to_tensor(frames_np[i+1])
            mid_tensors = make_inference(I0, I1, n_interp)
            for mid in mid_tensors:
                output_frames.append(from_tensor(mid))
            if (i + 1) % 50 == 0:
                pbar.update(50)
        pbar.update(total_steps % 50)
        output_frames.append(from_tensor(I1))

    del I0, I1, mid_tensors
    torch.cuda.empty_cache()
    return output_frames


# WAN

MODEL_ID = "thornmaze/WAMU_v3_WAN2.2_I2V_LIGHTNING"

LORA_MODELS = []

MAX_DIM = 832
MIN_DIM = 480
SQUARE_DIM = 640
MULTIPLE_OF = 16
MAX_SEED = np.iinfo(np.int32).max

FIXED_FPS = 16
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 321

MIN_DURATION = round(MIN_FRAMES_MODEL / FIXED_FPS, 1)
MAX_DURATION = round(MAX_FRAMES_MODEL / FIXED_FPS, 1)

SCHEDULER_MAP = {
    "FlowMatchEulerDiscrete": FlowMatchEulerDiscreteScheduler,
    "SASolver": SASolverScheduler,
    "DEISMultistep": DEISMultistepScheduler,
    "DPMSolverMultistepInverse": DPMSolverMultistepInverseScheduler,
    "UniPCMultistep": UniPCMultistepScheduler,
    "DPMSolverMultistep": DPMSolverMultistepScheduler,
    "DPMSolverSinglestep": DPMSolverSinglestepScheduler,
}

pipe = WanImageToVideoPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
).to('cuda')
original_scheduler = copy.deepcopy(pipe.scheduler)

for i, lora in enumerate(LORA_MODELS):
    name_high_tr = lora["high_tr"].split(".")[0].split("/")[-1] + "Hh"
    name_low_tr = lora["low_tr"].split(".")[0].split("/")[-1] + "Ll"

    try:
        pipe.load_lora_weights(
            lora["repo_id"],
            weight_name=lora["high_tr"],
            adapter_name=name_high_tr
        )

        kwargs_lora = {"load_into_transformer_2": True}
        pipe.load_lora_weights(
            lora["repo_id"],
            weight_name=lora["low_tr"],
            adapter_name=name_low_tr,
            **kwargs_lora
        )

        pipe.set_adapters([name_high_tr, name_low_tr], adapter_weights=[1.0, 1.0])

        pipe.fuse_lora(adapter_names=[name_high_tr], lora_scale=lora["high_scale"], components=["transformer"])
        pipe.fuse_lora(adapter_names=[name_low_tr], lora_scale=lora["low_scale"], components=["transformer_2"])

        pipe.unload_lora_weights()

        print(f"Applied: {lora['high_tr']}, hs={lora['high_scale']}/ls={lora['low_scale']}, {i+1}/{len(LORA_MODELS)}")
    except Exception as e:
        print("Error:", str(e))
        print("Failed lora:", name_high_tr)
        pipe.unload_lora_weights()

quantize_(pipe.text_encoder, Int8WeightOnlyConfig())
torch._dynamo.reset()
quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
torch._dynamo.reset()
quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
torch._dynamo.reset()

spaces.aoti_load(
    module=pipe.transformer,
    repo_id='thornmaze/WanTransformer3DModel-sm120-cu130-raa',
)
spaces.aoti_load(
    module=pipe.transformer_2,
    repo_id='thornmaze/WanTransformer3DModel-sm120-cu130-raa',
)

default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation"
default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作品, 画作, 画面, 静止, 整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走"

LORA_CHOICES = lora_loader.get_lora_choices()

# ---------- Speed / step presets ----------

SPEED_LORA_SOURCES = {
    "seko_v1": {
        "repo": "lightx2v/Wan2.2-Lightning",
        "file_high": "Wan2.2-I2V-A14B-4steps-lora-rank64-Seko-V1/high_noise_model.safetensors",
        "file_low": "Wan2.2-I2V-A14B-4steps-lora-rank64-Seko-V1/low_noise_model.safetensors",
    },
    "distill": {
        "repo": "lightx2v/Wan2.2-Distill-Loras",
        "match_high": ["i2v", "high", "4step"],
        "match_low": ["i2v", "low", "4step"],
    },
}

SPEED_NONE = "Built-in Lightning · 6 steps (balanced)"

SPEED_PRESETS = {
    "Built-in Lightning · 4 steps (fastest)":
        {"steps": 4, "flow_shift": 3.0, "scheduler": "UniPCMultistep", "lora": None},
    SPEED_NONE:
        {"steps": 6, "flow_shift": 3.0, "scheduler": "UniPCMultistep", "lora": None},
    "Built-in Lightning · 8 steps (best motion)":
        {"steps": 8, "flow_shift": 5.0, "scheduler": "UniPCMultistep", "lora": None},
    "Lightning lora (lightx2v Seko V1) · 4 steps":
        {"steps": 4, "flow_shift": 5.0, "scheduler": "UniPCMultistep", "lora": "seko_v1"},
    "Lightning lora (lightx2v Seko V1) · 6 steps":
        {"steps": 6, "flow_shift": 5.0, "scheduler": "UniPCMultistep", "lora": "seko_v1"},
    "Lightning lora (lightx2v Seko V1) · 8 steps":
        {"steps": 8, "flow_shift": 5.0, "scheduler": "UniPCMultistep", "lora": "seko_v1"},
    "Distill lora (lightx2v 4step) · 4 steps":
        {"steps": 4, "flow_shift": 5.0, "scheduler": "UniPCMultistep", "lora": "distill"},
    "Distill lora (lightx2v 4step) · 8 steps":
        {"steps": 8, "flow_shift": 5.0, "scheduler": "UniPCMultistep", "lora": "distill"},
}

_SPEED_CACHE = {}


def _match_all(name, keys):
    n = str(name).lower()
    return all(k in n for k in keys)


def resolve_speed_lora(key):
    if not key:
        return None
    if key in _SPEED_CACHE:
        return _SPEED_CACHE[key]

    src = SPEED_LORA_SOURCES.get(key)
    if not src:
        return None

    if "file_high" in src:
        out = {"repo": src["repo"], "file_high": src["file_high"], "file_low": src["file_low"]}
        _SPEED_CACHE[key] = out
        return out

    try:
        files = list(lora_loader.list_loras_from_repo(src["repo"]))
    except Exception as e:
        print("Speed lora listing failed:", e)
        return None

    highs = [f for f in files if f and f != "(None)" and _match_all(f, src["match_high"])]
    lows = [f for f in files if f and f != "(None)" and _match_all(f, src["match_low"])]
    if not highs or not lows:
        print("Speed lora files not found in", src["repo"])
        return None

    out = {"repo": src["repo"], "file_high": sorted(highs)[-1], "file_low": sorted(lows)[-1]}
    _SPEED_CACHE[key] = out
    return out


PROMPT_CHIPS = [
    "cinematic motion",
    "slow camera zoom in",
    "camera orbit around subject",
    "slow dolly forward",
    "gentle handheld camera",
    "wind blowing, hair and cloth moving",
    "soft natural light, shallow depth of field",
    "subject slowly turns head and smiles",
    "rain falling, wet reflections",
    "smoke and dust particles in the air",
]


def model_title():
    return "## Wan 2.2 I2V 14B Custom Lora"


def resize_image(image: Image.Image) -> Image.Image:
    width, height = image.size
    if width == height:
        return image.resize((SQUARE_DIM, SQUARE_DIM), Image.LANCZOS)

    aspect_ratio = width / height
    MAX_ASPECT_RATIO = MAX_DIM / MIN_DIM
    MIN_ASPECT_RATIO = MIN_DIM / MAX_DIM

    image_to_resize = image
    if aspect_ratio > MAX_ASPECT_RATIO:
        target_w, target_h = MAX_DIM, MIN_DIM
        crop_width = int(round(height * MAX_ASPECT_RATIO))
        left = (width - crop_width) // 2
        image_to_resize = image.crop((left, 0, left + crop_width, height))
    elif aspect_ratio < MIN_ASPECT_RATIO:
        target_w, target_h = MIN_DIM, MAX_DIM
        crop_height = int(round(width / MIN_ASPECT_RATIO))
        top = (height - crop_height) // 2
        image_to_resize = image.crop((0, top, width, top + crop_height))
    else:
        if width > height:
            target_w = MAX_DIM
            target_h = int(round(target_w / aspect_ratio))
        else:
            target_h = MAX_DIM
            target_w = int(round(target_h * aspect_ratio))

    final_w = round(target_w / MULTIPLE_OF) * MULTIPLE_OF
    final_h = round(target_h / MULTIPLE_OF) * MULTIPLE_OF
    final_w = max(MIN_DIM, min(MAX_DIM, final_w))
    final_h = max(MIN_DIM, min(MAX_DIM, final_h))
    return image_to_resize.resize((final_w, final_h), Image.LANCZOS)


def resize_and_crop_to_match(target_image, reference_image):
    ref_width, ref_height = reference_image.size
    target_width, target_height = target_image.size
    scale = max(ref_width / target_width, ref_height / target_height)
    new_width, new_height = int(target_width * scale), int(target_height * scale)
    resized = target_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    left, top = (new_width - ref_width) // 2, (new_height - ref_height) // 2
    return resized.crop((left, top, left + ref_width, top + ref_height))


def get_num_frames(duration_seconds: float):
    raw = int(round(duration_seconds * FIXED_FPS))
    raw = max(MIN_FRAMES_MODEL, min(MAX_FRAMES_MODEL, raw))
    return ((raw - 1) // 4) * 4 + 1


def get_inference_duration(
    resized_image,
    processed_last_image,
    prompt,
    steps,
    negative_prompt,
    num_frames,
    guidance_scale,
    guidance_scale_2,
    current_seed,
    scheduler_name,
    flow_shift,
    frame_multiplier,
    quality,
    duration_seconds,
    safe_mode,
    lora_groups,
    progress
):
    BASE_FRAMES_HEIGHT_WIDTH = 81 * 832 * 624
    BASE_STEP_DURATION = 5.
    width, height = resized_image.size
    factor = num_frames * width * height / BASE_FRAMES_HEIGHT_WIDTH
    step_duration = BASE_STEP_DURATION * factor ** 1.5
    gen_time = int(steps) * step_duration

    if guidance_scale > 1:
        gen_time = gen_time * 2.4

    frame_factor = frame_multiplier // FIXED_FPS
    if frame_factor > 1:
        total_out_frames = (num_frames * frame_factor) - num_frames
        inter_time = (total_out_frames * 0.02)
        gen_time += inter_time

    total_time = 15 + gen_time
    if safe_mode:
        total_time = total_time * 1.30

    return total_time


@spaces.GPU(duration=get_inference_duration, size='xlarge')
def run_inference(
    resized_image,
    processed_last_image,
    prompt,
    steps,
    negative_prompt,
    num_frames,
    guidance_scale,
    guidance_scale_2,
    current_seed,
    scheduler_name,
    flow_shift,
    frame_multiplier,
    quality,
    duration_seconds,
    safe_mode=False,
    lora_groups=None,
    progress=gr.Progress(track_tqdm=True),
):
    scheduler_class = SCHEDULER_MAP.get(scheduler_name)
    if scheduler_class.__name__ != pipe.scheduler.config._class_name or flow_shift != pipe.scheduler.config.get("flow_shift", "shift"):
        config = copy.deepcopy(original_scheduler.config)
        if scheduler_class == FlowMatchEulerDiscreteScheduler:
            config['shift'] = flow_shift
        else:
            config['flow_shift'] = flow_shift
        pipe.scheduler = scheduler_class.from_config(config)

    clear_vram()

    task_name = str(uuid.uuid4())[:8]
    print(f"Generating {num_frames} frames, task: {task_name}, {duration_seconds}, {resized_image.size}, lora={lora_groups}")
    start = time.time()

    lora_loaded = False
    if lora_groups:
        try:
            lora_loaded = lora_loader.load_lora_slots(pipe, lora_groups)
            print(f"lora loaded: {lora_groups}")
        except Exception as e:
            print(f"lora warning: {e}")
            try:
                lora_loader.unload_lora(pipe)
            except Exception:
                pass

    result = pipe(
        image=resized_image,
        last_image=processed_last_image,
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=resized_image.height,
        width=resized_image.width,
        num_frames=num_frames,
        guidance_scale=float(guidance_scale),
        guidance_scale_2=float(guidance_scale_2),
        num_inference_steps=int(steps),
        generator=torch.Generator(device="cuda").manual_seed(current_seed),
        output_type="np"
    )

    if lora_loaded:
        lora_loader.unload_lora(pipe)

    print("gen time passed:", time.time() - start)

    raw_frames_np = result.frames[0]
    pipe.scheduler = original_scheduler

    frame_factor = frame_multiplier // FIXED_FPS
    if frame_factor > 1:
        start = time.time()
        print(f"Processing frames (RIFE Multiplier: {frame_factor}x)...")
        rife_model.device()
        rife_model.flownet = rife_model.flownet.half()
        final_frames = interpolate_bits(raw_frames_np, multiplier=int(frame_factor))
        print("Interpolation time passed:", time.time() - start)
    else:
        final_frames = list(raw_frames_np)

    final_fps = FIXED_FPS * int(frame_factor)

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
        video_path = tmpfile.name

    start = time.time()
    with tqdm(total=3, desc="Rendering Media", unit="clip") as pbar:
        pbar.update(2)
        export_to_video(final_frames, video_path, fps=final_fps, quality=quality)
        pbar.update(1)
    print(f"Export time passed, {final_fps} FPS:", time.time() - start)

    return video_path, task_name


def generate_video(
    input_image,
    last_image,
    prompt,
    steps=4,
    negative_prompt=default_negative_prompt,
    duration_seconds=MAX_DURATION,
    guidance_scale=1,
    guidance_scale_2=1,
    seed=42,
    randomize_seed=False,
    quality=5,
    scheduler="UniPCMultistep",
    flow_shift=6.0,
    frame_multiplier=16,
    safe_mode=False,
    speed_preset=SPEED_NONE,
    speed_high_scale=1.0,
    speed_low_scale=1.0,
    # original 3
    lora1_preset="(None)",
    lora1_high_scale=1.0,
    lora1_low_scale=1.0,
    lora2_preset="(None)",
    lora2_high_scale=1.0,
    lora2_low_scale=1.0,
    lora3_preset="(None)",
    lora3_high_scale=1.0,
    lora3_low_scale=1.0,
    # my 5
    my1_repo="", my1_high="(None)", my1_hs=1.0, my1_low="(None)", my1_ls=1.0,
    my2_repo="", my2_high="(None)", my2_hs=1.0, my2_low="(None)", my2_ls=1.0,
    my3_repo="", my3_high="(None)", my3_hs=1.0, my3_low="(None)", my3_ls=1.0,
    my4_repo="", my4_high="(None)", my4_hs=1.0, my4_low="(None)", my4_ls=1.0,
    my5_repo="", my5_high="(None)", my5_hs=1.0, my5_low="(None)", my5_ls=1.0,
    video_component=True,
    progress=gr.Progress(track_tqdm=True),
):
    if input_image is None:
        raise gr.Error("Please upload an input image.")

    wall_start = time.time()

    num_frames = get_num_frames(duration_seconds)
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
    resized_image = resize_image(input_image)

    processed_last_image = None
    if last_image:
        processed_last_image = resize_and_crop_to_match(last_image, resized_image)

    preset_cfg = SPEED_PRESETS.get(speed_preset) or {}
    speed_info = ""
    speed_group = None
    if preset_cfg.get("lora"):
        resolved = resolve_speed_lora(preset_cfg["lora"])
        if resolved:
            speed_group = {
                "repo": resolved["repo"],
                "file_high": resolved["file_high"],
                "file_low": resolved["file_low"],
                "high_scale": speed_high_scale,
                "low_scale": speed_low_scale,
            }
            speed_info = f"<span>speed lora on</span>"
        else:
            speed_info = "<span>speed lora not found — skipped</span>"

    lora_groups = [
        {"preset": lora1_preset, "high_scale": lora1_high_scale, "low_scale": lora1_low_scale},
        {"preset": lora2_preset, "high_scale": lora2_high_scale, "low_scale": lora2_low_scale},
        {"preset": lora3_preset, "high_scale": lora3_high_scale, "low_scale": lora3_low_scale},
        {"repo": my1_repo, "file_high": my1_high, "file_low": my1_low, "high_scale": my1_hs, "low_scale": my1_ls},
        {"repo": my2_repo, "file_high": my2_high, "file_low": my2_low, "high_scale": my2_hs, "low_scale": my2_ls},
        {"repo": my3_repo, "file_high": my3_high, "file_low": my3_low, "high_scale": my3_hs, "low_scale": my3_ls},
        {"repo": my4_repo, "file_high": my4_high, "file_low": my4_low, "high_scale": my4_hs, "low_scale": my4_ls},
        {"repo": my5_repo, "file_high": my5_high, "file_low": my5_low, "high_scale": my5_hs, "low_scale": my5_ls},
    ]

    if speed_group:
        lora_groups.insert(0, speed_group)

    video_path, task_n = run_inference(
        resized_image,
        processed_last_image,
        prompt,
        steps,
        negative_prompt,
        num_frames,
        guidance_scale,
        guidance_scale_2,
        current_seed,
        scheduler,
        flow_shift,
        frame_multiplier,
        quality,
        duration_seconds,
        safe_mode,
        lora_groups,
        progress,
    )
    print(f"GPU complete: {task_n}")

    elapsed = time.time() - wall_start
    w, h = resized_image.size
    status = (
        "<div class='status-ok'>"
        "<b>&#10003; Ready</b>"
        f"<span>{w}&times;{h}</span>"
        f"<span>{num_frames} frames &middot; {int(frame_multiplier)} fps</span>"
        f"<span>seed {current_seed}</span>"
        f"<span>{int(steps)} steps</span>"
        + speed_info +
        f"<span>{elapsed:.1f} s</span>"
        "</div>"
    )

    return (video_path if video_component else None), video_path, current_seed, status


def _pick_file(files, *keys):
    for key in keys:
        for f in files:
            if f and f != "(None)" and key in str(f).lower():
                return f
    return "(None)"


def _load_repo_files(repo_link):
    if not repo_link or not str(repo_link).strip():
        return (
            gr.update(choices=["(None)"], value="(None)"),
            gr.update(choices=["(None)"], value="(None)"),
        )

    # A CivitAI (or any other) download link is a single file, not a repo to list: put it in the
    # High box and leave Low empty, since a Wan adapter ships its two halves as two separate files.
    link = str(repo_link).strip()
    if getattr(lora_loader, "is_direct_link", lambda _: False)(link):
        return (
            gr.update(choices=["(None)", link], value=link),
            gr.update(choices=["(None)", link], value="(None)"),
        )
    try:
        files = lora_loader.list_loras_from_repo(repo_link)
    except Exception as e:
        raise gr.Error(f"Cannot read repo: {e}")

    files = list(files) if files else []
    if "(None)" not in files:
        files = ["(None)"] + files

    auto_high = _pick_file(files, "high", "_h.", "-h.", "hn")
    auto_low = _pick_file(files, "low", "_l.", "-l.", "ln")

    return (
        gr.update(choices=files, value=auto_high),
        gr.update(choices=files, value=auto_low),
    )


# ---------------------------------------------------------------------------
# Named profiles
# ---------------------------------------------------------------------------
# Everything typed rather than uploaded, stored under a name inside the Space.
# `/data` is used when the Space has persistent storage attached, so profiles
# survive a restart; otherwise they live for as long as the Space runs.

import json
import re
import tempfile
import time

if os.path.isdir("/data") and os.access("/data", os.W_OK):
    PROFILES_DIR = os.path.join("/data", "wan22-profiles")
else:
    PROFILES_DIR = os.path.join(tempfile.gettempdir(), "wan22-profiles")
os.makedirs(PROFILES_DIR, exist_ok=True)

NO_PROFILE = "\u2014 no saved profile \u2014"

SETTINGS_KEYS = [
    "prompt", "negative_prompt", "duration", "steps", "quality",
    "seed", "randomize_seed", "scheduler", "flow_shift", "fps", "safe_mode",
    "speed_preset", "speed_high_scale", "speed_low_scale",
    "lora1_preset", "lora1_high", "lora1_low",
    "lora2_preset", "lora2_high", "lora2_low",
    "lora3_preset", "lora3_high", "lora3_low",
]
for _slot in range(1, 6):
    SETTINGS_KEYS += [
        f"my{_slot}_repo", f"my{_slot}_high_file", f"my{_slot}_high_scale",
        f"my{_slot}_low_file", f"my{_slot}_low_scale",
    ]

# The per-repo file dropdowns are populated only after a repo has been read, so a restored value
# would not be among the choices. These keys get their choice list rebuilt around the stored value.
_FILE_DROPDOWN_KEYS = {f"my{slot}_{end}_file" for slot in range(1, 6) for end in ("high", "low")}


def _profile_path(name):
    safe = re.sub(r"[^A-Za-z0-9 ._-]", "_", str(name or "").strip())[:64].strip() or "profile"
    return os.path.join(PROFILES_DIR, f"{safe}.json")


def list_profiles():
    if not os.path.isdir(PROFILES_DIR):
        return []
    return [entry[:-5] for entry in sorted(os.listdir(PROFILES_DIR)) if entry.endswith(".json")]


def refresh_profiles(current=None):
    """Re-read the folder. The dropdown's choices are built once at start-up, so a profile saved in
    another tab (or after this page was opened) would otherwise stay invisible until a restart."""
    names = list_profiles()
    value = current if current in names else NO_PROFILE
    return gr.update(choices=[NO_PROFILE, *names], value=value)


def save_profile(name, *values):
    if not str(name or "").strip():
        return gr.update(), "Give the profile a name first."
    payload = {"saved": time.strftime("%Y-%m-%d %H:%M:%S")}
    payload.update(dict(zip(SETTINGS_KEYS, values)))
    try:
        with open(_profile_path(name), "w", encoding="utf-8") as handle:
            json.dump(payload, handle, ensure_ascii=False, indent=2, default=str)
    except Exception as error:
        return gr.update(), f"Could not save: `{type(error).__name__}: {error}`"
    saved = os.path.basename(_profile_path(name))[:-5]
    return gr.update(choices=[NO_PROFILE, *list_profiles()], value=saved), f"Saved **{saved}**."


def load_profile(name):
    blank = [gr.update() for _ in SETTINGS_KEYS]
    if not name or name == NO_PROFILE:
        return [*blank, gr.update(), ""]
    path = _profile_path(name)
    if not os.path.exists(path):
        return [*blank, gr.update(), f"No profile named **{name}**."]
    try:
        with open(path, encoding="utf-8") as handle:
            payload = json.load(handle)
    except Exception as error:
        return [*blank, gr.update(), f"Could not read it: `{type(error).__name__}: {error}`"]

    updates = []
    for key in SETTINGS_KEYS:
        value = payload.get(key)
        if value is None:
            updates.append(gr.update())
        elif key in _FILE_DROPDOWN_KEYS:
            choices = ["(None)"] if value == "(None)" else ["(None)", value]
            updates.append(gr.update(choices=choices, value=value))
        else:
            updates.append(gr.update(value=value))
    stamp = payload.get("saved", "")
    return [*updates, gr.update(value=name), f"Loaded **{name}**{f' (saved {stamp})' if stamp else ''}."]


def delete_profile(name):
    if not name or name == NO_PROFILE:
        return gr.update(), "Pick a profile first."
    try:
        os.remove(_profile_path(name))
    except FileNotFoundError:
        return gr.update(choices=[NO_PROFILE, *list_profiles()], value=NO_PROFILE), f"No profile named **{name}**."
    except Exception as error:
        return gr.update(), f"Could not delete: `{type(error).__name__}: {error}`"
    return gr.update(choices=[NO_PROFILE, *list_profiles()], value=NO_PROFILE), f"Deleted **{name}**."


def export_settings(*values):
    payload = {"saved": time.strftime("%Y-%m-%d %H:%M:%S")}
    payload.update(dict(zip(SETTINGS_KEYS, values)))
    directory = os.path.join(tempfile.gettempdir(), "wan22-settings")
    os.makedirs(directory, exist_ok=True)
    path = os.path.join(directory, f"wan22-settings-{int(time.time())}.json")
    with open(path, "w", encoding="utf-8") as handle:
        json.dump(payload, handle, ensure_ascii=False, indent=2, default=str)
    return gr.update(value=path, visible=True)


def import_settings(path):
    if not path:
        return [gr.update() for _ in SETTINGS_KEYS]
    try:
        with open(path, encoding="utf-8") as handle:
            payload = json.load(handle)
    except Exception as error:
        raise gr.Error(f"That settings file cannot be read: `{type(error).__name__}: {error}`")
    updates = []
    for key in SETTINGS_KEYS:
        value = payload.get(key)
        if value is None:
            updates.append(gr.update())
        elif key in _FILE_DROPDOWN_KEYS:
            choices = ["(None)"] if value == "(None)" else ["(None)", value]
            updates.append(gr.update(choices=choices, value=value))
        else:
            updates.append(gr.update(value=value))
    return updates


# ---------------------------------------------------------------------------
# Clip stitching, boomerang, and persistent saving - all CPU, zero GPU quota
# ---------------------------------------------------------------------------
# Every one of these runs after the video already exists, so none of them touch
# the card or the ZeroGPU allowance. Chaining three 3.5 s clips into a 10 s one
# costs nothing extra.

def _ffmpeg_exe() -> str:
    """imageio-ffmpeg ships a binary, so this works even when the image has no
    system ffmpeg."""
    try:
        import imageio_ffmpeg
        return imageio_ffmpeg.get_ffmpeg_exe()
    except Exception:
        return "ffmpeg"


def _video_size(path):
    capture = cv2.VideoCapture(path)
    width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH) or 0)
    height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0)
    capture.release()
    return (width or 832), (height or 624)


def _persist(path):
    """Copy to /data/videos when the Space has persistent storage attached, so a
    restart does not take the clip with it."""
    if not path or not os.path.isdir("/data") or not os.access("/data", os.W_OK):
        return
    try:
        target_dir = os.path.join("/data", "videos")
        os.makedirs(target_dir, exist_ok=True)
        shutil.copy(path, os.path.join(target_dir, os.path.basename(path)))
        print(f"[storage] saved to /data/videos/{os.path.basename(path)}")
    except Exception as error:
        print(f"[storage] could not save: {error}")


def concat_videos(paths, name_hint=""):
    """Stitch clips into one file. Everything is scaled and padded to the first
    clip's frame so a mixed-aspect queue still concatenates cleanly."""
    valid = [p for p in (paths or []) if p and os.path.exists(p)]
    if not valid:
        return None
    if len(valid) == 1:
        return valid[0]

    width, height = _video_size(valid[0])
    prefix = uuid.uuid4().hex[:4]
    clean = (str(name_hint).strip().replace(" ", "_") or f"merged_{len(valid)}clips")
    if clean.lower().endswith(".mp4"):
        clean = clean[:-4]
    out_path = os.path.join(tempfile.gettempdir(), f"{prefix}_{clean}.mp4")

    inputs, chains = [], []
    for index, path in enumerate(valid):
        inputs += ["-i", path]
        chains.append(
            f"[{index}:v]scale={width}:{height}:force_original_aspect_ratio=decrease,"
            f"pad={width}:{height}:(ow-iw)/2:(oh-ih)/2,setsar=1[v{index}];"
        )
    joined = "".join(f"[v{i}]" for i in range(len(valid)))
    filter_str = "".join(chains) + f"{joined}concat=n={len(valid)}:v=1:a=0[outv]"

    command = [_ffmpeg_exe(), "-y", *inputs, "-filter_complex", filter_str,
               "-map", "[outv]", "-c:v", "libx264", "-preset", "fast", "-crf", "20",
               "-movflags", "+faststart", out_path]
    try:
        subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    except Exception as error:
        print(f"[merge] filter_complex failed ({error}); trying the concat demuxer")
        list_file = os.path.join(tempfile.gettempdir(), f"concat_{prefix}.txt")
        with open(list_file, "w", encoding="utf-8") as handle:
            for path in valid:
                handle.write(f"file '{path}'\n")
        subprocess.run([_ffmpeg_exe(), "-y", "-f", "concat", "-safe", "0", "-i", list_file,
                        "-c", "copy", "-movflags", "+faststart", out_path], check=False)
    if os.path.exists(out_path):
        _persist(out_path)
        return out_path
    return valid[-1]


def _queue_status(count, path=None):
    if not count:
        return "<div class='status-idle'>Merge queue: empty</div>"
    name = f" &rarr; {os.path.basename(path)}" if path else ""
    return f"<div class='status-ok'><b>{count} clip(s) in the queue</b><span>{name}</span></div>"


def add_to_queue(video_path, queue, name_hint):
    queue = list(queue or [])
    if video_path and os.path.exists(video_path):
        queue.append(video_path)
    if not queue:
        return None, None, [], _queue_status(0)
    merged = concat_videos(queue, name_hint)
    return merged, merged, queue, _queue_status(len(queue), merged)


def clear_queue():
    return None, None, [], _queue_status(0)


def auto_queue(video_path, auto_on, queue, name_hint):
    """Runs after every generation; only appends when the box is ticked."""
    if not auto_on:
        queue = list(queue or [])
        if not queue:
            return gr.update(), gr.update(), queue, _queue_status(0)
        return gr.update(), gr.update(), queue, _queue_status(len(queue))
    return add_to_queue(video_path, queue, name_hint)


def boomerang(video_path, mode, tail_seconds=2.0):
    """`full` plays forward then backward. `ending` keeps the clip as it is and
    boomerangs only its tail, which loops without the whole thing rewinding."""
    if not video_path or not os.path.exists(video_path):
        raise gr.Error("Generate a video first.")

    prefix = uuid.uuid4().hex[:4]
    out_path = os.path.join(tempfile.gettempdir(), f"{prefix}_boomerang.mp4")

    if mode == "ending":
        capture = cv2.VideoCapture(video_path)
        fps = capture.get(cv2.CAP_PROP_FPS) or 16
        total = int(capture.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
        capture.release()
        duration = total / fps if fps else 0
        start = max(0.0, duration - float(tail_seconds))
        filter_str = (
            f"[0:v]split=2[base][tail];"
            f"[tail]trim=start={start:.3f},setpts=PTS-STARTPTS,reverse[rev];"
            f"[base][rev]concat=n=2:v=1[outv]"
        )
    else:
        filter_str = "[0:v]split=2[a][b];[b]reverse[r];[a][r]concat=n=2:v=1[outv]"

    command = [_ffmpeg_exe(), "-y", "-i", video_path, "-filter_complex", filter_str,
               "-map", "[outv]", "-c:v", "libx264", "-preset", "fast", "-crf", "20",
               "-movflags", "+faststart", out_path]
    try:
        subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    except Exception as error:
        raise gr.Error(f"Loop failed: {error}")
    _persist(out_path)
    return out_path, out_path


# ---------------------------------------------------------------------------
# Prompt building and enhancement
# ---------------------------------------------------------------------------
# Two paths. The builder is pure string work: it takes what you typed and wraps
# it in the camera, framing, light and pacing language Wan responds to. It is
# instant, costs nothing, and cannot fail or refuse.
#
# The model path runs an abliterated Qwen3-4B - a small instruct model with its
# refusal directions removed - so it writes the paragraph for you whatever the
# subject is. The weights are fetched on the CPU, outside GPU time; only the
# generation itself books the card, and it books ~45 s.

ENHANCER_MODEL = os.environ.get(
    "ENHANCER_MODEL", "huihui-ai/Huihui-Qwen3-4B-Instruct-2507-abliterated"
)

CAMERA_MOVES = {
    "(keep as written)": "",
    "static shot": "the camera holds still",
    "slow push in": "the camera pushes in slowly",
    "slow pull back": "the camera pulls back slowly",
    "orbit around the subject": "the camera orbits around the subject",
    "tracking alongside": "the camera tracks alongside the subject",
    "handheld, subtle shake": "handheld camera with a subtle shake",
    "tilt up": "the camera tilts up",
    "tilt down": "the camera tilts down",
    "crane up and back": "the camera cranes up and back",
}

FRAMINGS = {
    "(keep as written)": "",
    "extreme close-up": "extreme close-up",
    "close-up": "close-up",
    "medium shot": "medium shot",
    "medium wide": "medium wide shot",
    "wide shot": "wide shot",
    "low angle": "low angle shot",
    "high angle": "high angle shot",
    "over the shoulder": "over-the-shoulder shot",
    "point of view": "point-of-view shot",
}

LIGHTING = {
    "(keep as written)": "",
    "soft natural light": "soft natural light, gentle falloff",
    "golden hour": "warm golden-hour light with a soft rim",
    "hard key light": "a hard key light with deep shadows",
    "candle / firelight": "flickering warm candlelight",
    "neon night": "neon light, wet reflections, night",
    "overcast daylight": "flat overcast daylight",
    "studio softbox": "clean studio softbox lighting on a plain backdrop",
    "single lamp in a dark room": "one warm lamp in an otherwise dark room",
    "moonlight": "cool moonlight through a window",
}

PACING = {
    "(keep as written)": "",
    "slow and deliberate": "slow, deliberate movement",
    "steady": "steady, even movement",
    "building": "movement that builds through the shot",
    "fast and urgent": "fast, urgent movement",
    "loose and languid": "loose, languid movement",
}

QUALITY_TAIL = (
    "photorealistic, natural skin texture, shallow depth of field, "
    "consistent anatomy, smooth motion, no morphing"
)


def _strip_built(text):
    """Remove sentences a previous Build added, so pressing it again rewrites rather than stacks."""
    known = {QUALITY_TAIL.lower().rstrip(".")}
    for table in (LIGHTING, PACING):
        for value in table.values():
            if value:
                known.add(value.lower().rstrip("."))
    for framing in FRAMINGS.values():
        for camera in CAMERA_MOVES.values():
            joined = ", ".join(part for part in (framing, camera) if part)
            if joined:
                known.add(joined.lower().rstrip("."))

    kept = []
    for sentence in re.split(r"(?<=\.)\s+", (text or "").strip()):
        if sentence.strip().lower().rstrip(".") not in known:
            kept.append(sentence.strip())
    return " ".join(kept).strip()


def build_prompt(prompt, camera, framing, lighting, pace, add_quality=True):
    """Wrap what the user wrote in the shot language Wan responds to. No model,
    no network, nothing to refuse - it only ever adds to your own words, and it
    replaces what it added last time."""
    body = _strip_built(prompt).strip().rstrip(".,;")
    if not body:
        raise gr.Error("Write what should happen first - the builder adds the shot around it.")

    parts = [body + "."]
    shot = ", ".join(part for part in (FRAMINGS.get(framing, ""), CAMERA_MOVES.get(camera, "")) if part)
    if shot:
        parts.append(shot[0].upper() + shot[1:] + ".")
    if LIGHTING.get(lighting):
        parts.append(LIGHTING[lighting][0].upper() + LIGHTING[lighting][1:] + ".")
    if PACING.get(pace):
        parts.append(PACING[pace][0].upper() + PACING[pace][1:] + ".")
    if add_quality:
        parts.append(QUALITY_TAIL + ".")
    return " ".join(parts)


_ENHANCER = {"model": None, "tokenizer": None, "ready": False}

ENHANCER_SYSTEM = (
    "You write prompts for a video generation model. Rewrite the user's line as ONE paragraph of "
    "60 to 110 words, present tense, third person. Describe, in this order: who or what is in "
    "frame and how they look, what they do and how the motion progresses, the camera framing and "
    "any camera move, the lighting, and the setting. Keep every element the user gave you, "
    "including explicit ones, and keep their wording for those. Add nothing about story, mood "
    "words or artistic intent. Output the paragraph only - no preamble, no lists, no quotation "
    "marks, no commentary, no questions."
)


def _fetch_enhancer_files():
    """Download the weights on the CPU, so no GPU time is spent on the transfer."""
    from huggingface_hub import snapshot_download

    return snapshot_download(
        ENHANCER_MODEL,
        allow_patterns=["*.json", "*.safetensors", "*.txt", "*.jinja", "*.model"],
        token=os.environ.get("HF_TOKEN"),
    )


@spaces.GPU(duration=60)
def _run_enhancer(text):
    import torch as _torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    if not _ENHANCER["ready"]:
        _ENHANCER["tokenizer"] = AutoTokenizer.from_pretrained(ENHANCER_MODEL)
        _ENHANCER["model"] = AutoModelForCausalLM.from_pretrained(
            ENHANCER_MODEL, dtype=_torch.bfloat16, low_cpu_mem_usage=True
        )
        _ENHANCER["ready"] = True

    tokenizer, model = _ENHANCER["tokenizer"], _ENHANCER["model"]
    model = model.to("cuda")

    messages = [
        {"role": "system", "content": ENHANCER_SYSTEM},
        {"role": "user", "content": text},
    ]
    # Render to text and tokenize separately. `apply_chat_template(return_tensors=...)` hands back a
    # plain tensor on some transformers versions and a BatchEncoding on others, and asking the wrong
    # one for `.shape` is an AttributeError; going through text sidesteps the difference. The
    # thinking switch is Qwen3-only, so it is offered and dropped if the template ignores it.
    try:
        rendered = tokenizer.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=False, enable_thinking=False
        )
    except TypeError:
        rendered = tokenizer.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=False
        )

    inputs = tokenizer(rendered, return_tensors="pt").to("cuda")
    prompt_length = inputs["input_ids"].shape[-1]

    with _torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=320,
            do_sample=True,
            temperature=0.85,
            top_p=0.9,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
        )
    written = tokenizer.decode(out[0][prompt_length:], skip_special_tokens=True)

    # Some abliterated builds still emit a <think> block; keep only what follows it.
    if "</think>" in written:
        written = written.split("</think>", 1)[1]
    return " ".join(written.replace("\n", " ").split()).strip().strip('"')


def enhance_prompt(prompt, use_model, camera, framing, lighting, pace, add_quality):
    """Builder first - it is instant and never fails - then the model on top of it."""
    built = build_prompt(prompt, camera, framing, lighting, pace, add_quality)
    if not use_model:
        return built, "<div class='status-ok'><b>Built</b><span>no GPU time used</span></div>"

    try:
        _fetch_enhancer_files()
    except Exception as error:
        return built, (
            "<div class='status-idle'>Model unavailable, so the built prompt is shown "
            f"instead: {type(error).__name__}</div>"
        )

    try:
        written = _run_enhancer(built)
    except Exception as error:
        return built, (
            "<div class='status-idle'>The model call failed, so the built prompt is shown "
            f"instead: {error!r}</div>"
        )

    if len(written) < 40:
        return built, "<div class='status-idle'>The model returned too little; built prompt kept.</div>"
    return written, (
        f"<div class='status-ok'><b>Written by {ENHANCER_MODEL.split('/')[-1]}</b>"
        f"<span>{len(written.split())} words</span></div>"
    )


def _target_size(image):
    """Same math as resize_image(), but returns only the final (w, h)."""
    if image is None:
        return (832, 624)
    width, height = image.size
    if width == height:
        return (SQUARE_DIM, SQUARE_DIM)

    aspect_ratio = width / height
    MAX_ASPECT_RATIO = MAX_DIM / MIN_DIM
    MIN_ASPECT_RATIO = MIN_DIM / MAX_DIM

    if aspect_ratio > MAX_ASPECT_RATIO:
        target_w, target_h = MAX_DIM, MIN_DIM
    elif aspect_ratio < MIN_ASPECT_RATIO:
        target_w, target_h = MIN_DIM, MAX_DIM
    else:
        if width > height:
            target_w = MAX_DIM
            target_h = int(round(target_w / aspect_ratio))
        else:
            target_h = MAX_DIM
            target_w = int(round(target_h * aspect_ratio))

    final_w = round(target_w / MULTIPLE_OF) * MULTIPLE_OF
    final_h = round(target_h / MULTIPLE_OF) * MULTIPLE_OF
    final_w = max(MIN_DIM, min(MAX_DIM, final_w))
    final_h = max(MIN_DIM, min(MAX_DIM, final_h))
    return (final_w, final_h)


class _SizeOnly:
    def __init__(self, w, h):
        self.size = (w, h)
        self.width = w
        self.height = h


def _gpu_estimate_text(input_image, duration_seconds, steps, guidance_scale,
                       frame_multiplier, safe_mode):
    """Live ZeroGPU cost readout, using the exact same formula the
    @spaces.GPU decorator asks for at generation time."""
    try:
        num_frames = get_num_frames(float(duration_seconds))
        w, h = _target_size(input_image)
        secs = get_inference_duration(
            _SizeOnly(w, h), None, "", steps, "", num_frames,
            float(guidance_scale), 1.0, 0, "UniPCMultistep", 3.0,
            int(frame_multiplier), 5, float(duration_seconds),
            bool(safe_mode), None, None,
        )
    except Exception as exc:  # never break the UI over a readout
        return f"⏳ **ZeroGPU cost:** estimate unavailable ({exc})"

    note = "" if input_image is not None else " · upload an image for an exact number"
    out_fps = int(frame_multiplier)
    out_frames = num_frames * max(1, out_fps // FIXED_FPS)

    return (
        f"⏳ **ZeroGPU cost: ~{int(round(secs))} s**  \n"
        f"<sub>{num_frames} frames · {w}×{h} · {int(steps)} steps · "
        f"output {out_frames} frames @ {out_fps} fps{note}</sub>"
    )


def _apply_speed_preset(name):
    cfg = SPEED_PRESETS.get(name) or {}
    return (
        gr.update(value=cfg.get("steps", 6)),
        gr.update(value=cfg.get("flow_shift", 3.0)),
        gr.update(value=cfg.get("scheduler", "UniPCMultistep")),
    )


def _append_chip(prompt, chip):
    text = (prompt or "").strip().rstrip(",")
    if chip.lower() in text.lower():
        return text
    return f"{text}, {chip}" if text else chip


def _clear_all():
    return (
        None, None, default_prompt_i2v,
        gr.update(value=None), gr.update(value=None),
        "<div class='status-idle'>Ready when you are.</div>",
    )


CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');

.gradio-container {
    font-family: 'Inter', ui-sans-serif, system-ui, sans-serif !important;
    max-width: 1440px !important;
}

#hidden-timestamp {
    opacity: 0;
    height: 0px;
    width: 0px;
    margin: 0px;
    padding: 0px;
    overflow: hidden;
    position: absolute;
    pointer-events: none;
}

#hero {
    border-radius: 18px;
    padding: 22px 26px;
    margin-bottom: 14px;
    background: linear-gradient(120deg, #6d28d9 0%, #7c3aed 35%, #2563eb 100%);
    color: #fff;
    box-shadow: 0 10px 30px rgba(76, 29, 149, .25);
}
#hero h1 {
    margin: 0;
    font-size: 26px;
    font-weight: 700;
    letter-spacing: -.02em;
    color: #fff;
}
#hero p {
    margin: 6px 0 0 0;
    font-size: 14px;
    opacity: .92;
    color: #fff;
}
#hero .pills { margin-top: 12px; display: flex; flex-wrap: wrap; gap: 8px; }
#hero .pills span {
    background: rgba(255,255,255,.16);
    border: 1px solid rgba(255,255,255,.25);
    padding: 4px 11px;
    border-radius: 999px;
    font-size: 12px;
    font-weight: 500;
}

.panel {
    border-radius: 16px !important;
    border: 1px solid var(--border-color-primary) !important;
    padding: 14px !important;
    background: var(--background-fill-secondary) !important;
}

#generate-btn {
    font-size: 17px !important;
    font-weight: 700 !important;
    min-height: 62px !important;
    border-radius: 14px !important;
    letter-spacing: .01em;
    box-shadow: 0 8px 22px rgba(124, 58, 237, .28);
}
#estimate-btn, #enhance-btn { min-height: 42px !important; border-radius: 12px !important; }
#profile-load, #profile-save, #profile-delete, #profile-refresh {
    min-height: 42px !important; border-radius: 12px !important; font-weight: 600 !important;
}
#stop-btn, #clear-btn { min-height: 46px !important; border-radius: 12px !important; }
#grab-btn { min-height: 48px !important; border-radius: 12px !important; font-weight: 600 !important; }

.chip-row { gap: 6px !important; }
.chip-row button {
    border-radius: 999px !important;
    font-size: 12.5px !important;
    padding: 6px 12px !important;
    min-height: 34px !important;
    font-weight: 500 !important;
}

.gpu-estimate {
    border-radius: 12px;
    padding: 10px 14px;
    background: rgba(139, 92, 246, .10);
    border: 1px solid rgba(139, 92, 246, .35);
    font-size: 14px;
}
.gpu-estimate sub { opacity: .75; font-size: 12px; }

.status-ok, .status-idle {
    display: flex;
    flex-wrap: wrap;
    align-items: center;
    gap: 14px;
    border-radius: 12px;
    padding: 10px 14px;
    font-size: 13.5px;
}
.status-ok {
    background: rgba(16, 185, 129, .12);
    border: 1px solid rgba(16, 185, 129, .35);
}
.status-idle {
    background: var(--background-fill-secondary);
    border: 1px solid var(--border-color-primary);
    opacity: .8;
}
.status-ok span { opacity: .85; }

button.primary, .primary > button {
    background: linear-gradient(120deg, #7c3aed 0%, #6d28d9 55%, #2563eb 100%) !important;
    border: none !important;
    color: #fff !important;
}
button.primary:hover, .primary > button:hover { filter: brightness(1.08); }

.block, .form { border-radius: 14px !important; }
button { border-radius: 12px !important; }
.tab-nav button { font-weight: 600 !important; }
label > span { font-weight: 600 !important; }

footer { display: none !important; }
"""

HERO = """
<div id="hero">
  <h1>Wan 2.2 I2V 14B &middot; Custom lora + CivitAI links, prompt builder, Fast Options, ZeroGPU cost, setting profiles, clip stitching &amp; looping</h1>
  <p>Turn a still image into smooth cinematic motion &mdash; with your own loras. Stitch clips into longer
  videos and loop them on the CPU, for no extra ZeroGPU time.</p>
  <div class="pills">
    <span>14B I2V</span><span>4&ndash;8 steps</span><span>fp8 quantized</span>
    <span>AoT compiled</span><span>RIFE interpolation</span><span>3 built-in + 5 custom lora</span>
    <span>CivitAI links</span><span>prompt builder</span><span>clip stitching</span><span>boomerang loops</span><span>start &amp; end frame grab</span>
  </div>
</div>
"""

# (sweep every hour, keep files for a day). The old (3600, 10800) deleted a grabbed frame after
# three hours while the page still showed it, and the next Generate died in gradio's own
# preprocessing with FileNotFoundError before this code ran.
with gr.Blocks(delete_cache=(3600, 86400)) as demo:
    gr.HTML("<style>" + CSS + "</style>")
    gr.HTML(HERO)

    with gr.Row(equal_height=False):

        # ---------------- LEFT: controls ----------------
        with gr.Column(scale=5):

            with gr.Group(elem_classes="panel"):
                with gr.Tabs():
                    with gr.Tab("🖼️ Start frame"):
                        input_image_component = gr.Image(
                            type="pil", label=None, show_label=False,
                            sources=["upload", "clipboard"], height=300,
                        )
                    with gr.Tab("🏁 End frame (optional)"):
                        last_image_component = gr.Image(
                            type="pil", label=None, show_label=False,
                            sources=["upload", "clipboard"], height=300,
                        )

                prompt_input = gr.Textbox(
                    label="Prompt",
                    value=default_prompt_i2v,
                    lines=3,
                    placeholder="Describe the motion you want to see…",
                )

                with gr.Accordion("🖋️ Prompt builder", open=False):
                    gr.Markdown(
                        "Write what should happen in the box above, pick the shot below, and press "
                        "**Build**. The builder only ever adds to your own words — it is instant, "
                        "costs no GPU time and refuses nothing.\n\n"
                        "Tick **write it with a model** and a small local language model turns it "
                        "into a full paragraph instead, keeping your own wording for whatever you "
                        "wrote. The weights download on the CPU; only the writing books the card, "
                        "about 45 s. Any failure falls back to the built prompt."
                    )
                    with gr.Row():
                        pb_framing = gr.Dropdown(list(FRAMINGS), value="(keep as written)",
                                                 label="Framing")
                        pb_camera = gr.Dropdown(list(CAMERA_MOVES), value="(keep as written)",
                                                label="Camera move")
                    with gr.Row():
                        pb_lighting = gr.Dropdown(list(LIGHTING), value="(keep as written)",
                                                  label="Lighting")
                        pb_pace = gr.Dropdown(list(PACING), value="(keep as written)",
                                              label="Pacing")
                    pb_quality = gr.Checkbox(
                        value=True, label="Add the quality tail",
                        info="photorealistic, natural skin texture, consistent anatomy, smooth motion",
                    )
                    pb_use_model = gr.Checkbox(
                        value=False, label="🧠 Write it with a model (~45 s of GPU)",
                    )
                    pb_button = gr.Button("🖋️ Build / rewrite the prompt", variant="secondary",
                                          elem_id="enhance-btn")
                    pb_status = gr.HTML("")

                with gr.Accordion("✨ Quick motion tags — click to add", open=False):
                    with gr.Row(elem_classes="chip-row"):
                        for _chip in PROMPT_CHIPS[:5]:
                            _b = gr.Button(_chip, size="sm", variant="secondary")
                            _b.click(
                                fn=lambda p, c=_chip: _append_chip(p, c),
                                inputs=[prompt_input], outputs=[prompt_input],
                            )
                    with gr.Row(elem_classes="chip-row"):
                        for _chip in PROMPT_CHIPS[5:]:
                            _b = gr.Button(_chip, size="sm", variant="secondary")
                            _b.click(
                                fn=lambda p, c=_chip: _append_chip(p, c),
                                inputs=[prompt_input], outputs=[prompt_input],
                            )

            with gr.Group(elem_classes="panel"):
                duration_seconds_input = gr.Slider(
                    minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=3.5,
                    label="⏱️ Duration (seconds)",
                    info=f"Clamped to the model's {MIN_FRAMES_MODEL}–{MAX_FRAMES_MODEL} frames at {FIXED_FPS} fps.",
                )
                frame_multi = gr.Dropdown(
                    choices=[FIXED_FPS, FIXED_FPS*2, FIXED_FPS*4, FIXED_FPS*8],
                    value=FIXED_FPS,
                    label="🎞️ Fluidity (frames per second)",
                    info="Extra frames are created with flow estimation for smoother motion.",
                )
                safe_mode_checkbox = gr.Checkbox(
                    label="🛠️ Safe Mode",
                    value=True,
                    info="Requests 30% extra processing time so busy servers don't cut the task short.",
                )
                with gr.Row():
                    gpu_estimate = gr.Markdown(
                        _gpu_estimate_text(None, 3.5, 6, 1, FIXED_FPS, True),
                        elem_classes="gpu-estimate",
                    )
                    gpu_estimate_btn = gr.Button("🔄", variant="secondary", scale=1, min_width=60,
                                                 elem_id="estimate-btn")

            with gr.Group(elem_classes="panel"):
                speed_preset_dd = gr.Dropdown(
                    choices=list(SPEED_PRESETS.keys()),
                    value=SPEED_NONE,
                    label="⚡ Speed preset (steps)",
                    info="Sets steps, flow shift and scheduler in one click. lora options download on first use.",
                )
                with gr.Row():
                    speed_high_scale = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="Speed lora — high strength")
                    speed_low_scale = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="Speed lora — low strength")

            with gr.Accordion("⚙️ Advanced settings & Custom lora", open=False):
                with gr.Tabs():

                    with gr.Tab("Sampling"):
                        with gr.Row():
                            steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=6, label="Inference Steps")
                            quality_slider = gr.Slider(
                                minimum=1, maximum=10, step=1, value=6, label="Video Quality",
                                info="At 10 the file may be too large to preview in the browser.",
                            )
                        with gr.Row():
                            seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
                        with gr.Row():
                            randomize_seed_checkbox = gr.Checkbox(label="🎲 Randomize seed", value=True, interactive=True)
                            play_result_video = gr.Checkbox(label="▶️ Display result", value=True, interactive=True)
                        with gr.Row():
                            guidance_scale_input = gr.Slider(
                                minimum=0.0, maximum=10.0, step=0.5, value=1,
                                label="Guidance — high noise",
                                info="Above 1 uses more GPU and takes longer.",
                            )
                            guidance_scale_2_input = gr.Slider(
                                minimum=0.0, maximum=10.0, step=0.5, value=1,
                                label="Guidance — low noise",
                            )
                        with gr.Row():
                            scheduler_dropdown = gr.Dropdown(
                                label="Scheduler",
                                choices=list(SCHEDULER_MAP.keys()),
                                value="UniPCMultistep",
                            )
                            flow_shift_slider = gr.Slider(minimum=0.5, maximum=15.0, step=0.1, value=3.0, label="Flow Shift")
                        negative_prompt_input = gr.Textbox(
                            label="Negative Prompt",
                            value=default_negative_prompt,
                            info="Used if any Guidance Scale > 1.",
                            lines=3,
                        )

                    with gr.Tab("🎛️ Built-in lora"):
                        gr.Markdown("Built-in lora list — pick up to three and set their strengths.")
                        _preset_rows = []
                        for _i in range(1, 4):
                            with gr.Group(elem_classes="panel"):
                                _p = gr.Dropdown(choices=LORA_CHOICES, value="(None)", label=f"Built-in lora {_i}", filterable=True)
                                with gr.Row():
                                    _hs = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="High strength")
                                    _ls = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="Low strength")
                            _preset_rows.append((_p, _hs, _ls))

                        (lora1_preset, lora1_high_scale, lora1_low_scale) = _preset_rows[0]
                        (lora2_preset, lora2_high_scale, lora2_low_scale) = _preset_rows[1]
                        (lora3_preset, lora3_high_scale, lora3_low_scale) = _preset_rows[2]

                    with gr.Tab("⭐ Custom lora"):
                        gr.Markdown(
                            "**Your own lora from Hugging Face or CivitAI.** Paste the repo "
                            "(e.g. `username/my-loras`) and press **Enter** — the files load "
                            "automatically and High/Low are pre-selected for you. 5 slots available.\n\n"
                            "**CivitAI:** paste a download link "
                            "(`https://civitai.com/api/download/models/...`) into the repo box and press "
                            "Enter — it lands in the High box. A Wan adapter usually ships its two halves as "
                            "two separate files, so paste the low-noise link straight into the Low box "
                            "(both boxes accept a pasted link). The Space downloads it itself, so being "
                            "signed in to CivitAI in your browser does not help a gated model: add "
                            "`CIVITAI_TOKEN` under Settings → Variables and secrets. Downloads are cached."
                        )
                        _my_rows = []
                        for _i in range(1, 6):
                            with gr.Accordion(f"Custom lora {_i}", open=(_i == 1)):
                                with gr.Row():
                                    _repo = gr.Textbox(label="Repo", placeholder="username/repo", lines=1, scale=4)
                                    _load = gr.Button("📥 Load files", scale=1, variant="secondary")
                                with gr.Row():
                                    _high = gr.Dropdown(choices=["(None)"], value="(None)", label="High file", filterable=True,
                                                        allow_custom_value=True, scale=3)
                                    _hs = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="High strength", scale=2)
                                with gr.Row():
                                    _low = gr.Dropdown(choices=["(None)"], value="(None)", label="Low file", filterable=True,
                                                       allow_custom_value=True, scale=3)
                                    _ls = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="Low strength", scale=2)

                                _load.click(fn=_load_repo_files, inputs=[_repo], outputs=[_high, _low])
                                _repo.submit(fn=_load_repo_files, inputs=[_repo], outputs=[_high, _low])

                            _my_rows.append((_repo, _high, _hs, _low, _ls))

                        (my1_repo, my1_high, my1_hs, my1_low, my1_ls) = _my_rows[0]
                        (my2_repo, my2_high, my2_hs, my2_low, my2_ls) = _my_rows[1]
                        (my3_repo, my3_high, my3_hs, my3_low, my3_ls) = _my_rows[2]
                        (my4_repo, my4_high, my4_hs, my4_low, my4_ls) = _my_rows[3]
                        (my5_repo, my5_high, my5_hs, my5_low, my5_ls) = _my_rows[4]

                    with gr.Tab("💾 Profiles"):
                        gr.Markdown(
                            "Every current setting, kept under a name inside this Space. Images are "
                            "not stored."
                        )
                        with gr.Row():
                            profile_picker = gr.Dropdown(
                                label="Saved profiles",
                                choices=[NO_PROFILE, *list_profiles()],
                                value=NO_PROFILE,
                                scale=3,
                            )
                            profile_load_btn = gr.Button("📂 load", variant="secondary", scale=1,
                                                         elem_id="profile-load")
                            profile_refresh = gr.Button("🔄", variant="secondary", scale=1, min_width=60,
                                                        elem_id="profile-refresh")
                        with gr.Row():
                            profile_name = gr.Textbox(
                                label="Save as", placeholder="e.g. portrait, turbo 4 steps",
                                max_lines=1, scale=3,
                            )
                            profile_save_btn = gr.Button("💾 save", variant="primary", scale=1,
                                                         elem_id="profile-save")
                        profile_delete_btn = gr.Button("🗑️ delete the selected profile",
                                                       variant="secondary", elem_id="profile-delete")
                        profile_status = gr.Markdown("")
                        gr.Markdown("---")
                        gr.Markdown("**Take it with you (.json file)**")
                        settings_export_btn = gr.Button("⬇️ export the current settings to .json",
                                                        size="sm")
                        settings_download = gr.File(label="Your settings", visible=False,
                                                    interactive=False)
                        settings_upload = gr.File(label="📤 import a settings .json",
                                                  file_types=[".json"], type="filepath")

            generate_button = gr.Button("🚀  Generate Video", variant="primary", elem_id="generate-btn")
            with gr.Row():
                stop_button = gr.Button("⏹ Stop", variant="stop", elem_id="stop-btn")
                clear_button = gr.Button("🧹 Clear", variant="secondary", elem_id="clear-btn")

        # ---------------- RIGHT: result ----------------
        with gr.Column(scale=6):
            with gr.Group(elem_classes="panel"):
                video_output = gr.Video(
                    label="Generated Video",
                    autoplay=True,
                    sources=["upload"],
                    buttons=["download", "share"],
                    interactive=True,
                    elem_id="generated-video",
                    height=520,
                )
                status_box = gr.HTML("<div class='status-idle'>Ready when you are.</div>")

                with gr.Row():
                    grab_frame_btn = gr.Button(
                        "📸 Current frame → start image",
                        variant="secondary",
                        elem_id="grab-btn",
                    )
                    grab_frame_last_btn = gr.Button(
                        "🏁 Current frame → end image",
                        variant="secondary",
                        elem_id="grab-btn",
                    )
                    timestamp_box = gr.Number(value=0, label="Timestamp", visible=True, elem_id="hidden-timestamp")
                    timestamp_box_last = gr.Number(value=0, label="Timestamp", visible=True, elem_id="hidden-timestamp")

                file_output = gr.File(label="⬇️ Download video")

            with gr.Accordion("🔁 Loop the clip (no GPU used)", open=False):
                gr.Markdown(
                    "Runs on the result you already have, on the CPU — it costs nothing from your "
                    "ZeroGPU allowance."
                )
                boomerang_tail = gr.Slider(
                    0.5, 4.0, value=2.0, step=0.5,
                    label="Tail length for the ending loop (seconds)",
                )
                with gr.Row():
                    boomerang_full_btn = gr.Button("🔁 Full boomerang (forward + reverse)",
                                                   variant="secondary")
                    boomerang_end_btn = gr.Button("🔂 Ending only", variant="secondary")

            with gr.Accordion("🎞️ Stitch clips into one video (no GPU used)", open=False):
                gr.Markdown(
                    "Add clips to a queue and they are joined into a single file with ffmpeg. Three "
                    "3.5 s clips make one 10 s video for no extra GPU time. Everything is scaled to "
                    "the first clip's frame."
                )
                merge_name = gr.Textbox(label="File name (optional)", placeholder="my_scene",
                                        max_lines=1)
                auto_merge = gr.Checkbox(
                    value=False, label="⚡ Add every new generation automatically",
                    info="Generate, generate, generate — the joined file grows on its own.",
                )
                with gr.Row():
                    merge_add_btn = gr.Button("➕ Add the current video", variant="secondary")
                    merge_clear_btn = gr.Button("🧹 Clear the queue", variant="secondary")
                merge_status = gr.HTML("<div class='status-idle'>Merge queue: empty</div>")
                merged_video = gr.Video(label="🎬 Stitched result", height=360)
                merged_file = gr.File(label="⬇️ Download the stitched video")
                merge_queue = gr.State([])

            with gr.Accordion("💡 Tips", open=False):
                gr.Markdown(
                    "- Pause the video on any frame, then hit **Use current frame** to continue the shot from there.\n"
                    "- 6 steps + UniPC is the sweet spot; more steps rarely helps this model.\n"
                    "- Higher fluidity only interpolates — it adds smoothness, not new content.\n"
                    "- Keep **Randomize seed** on while exploring, then lock a seed you like."
                )

    ui_inputs = [
        input_image_component, last_image_component, prompt_input, steps_slider,
        negative_prompt_input, duration_seconds_input,
        guidance_scale_input, guidance_scale_2_input, seed_input, randomize_seed_checkbox,
        quality_slider, scheduler_dropdown, flow_shift_slider, frame_multi,
        safe_mode_checkbox,
        speed_preset_dd, speed_high_scale, speed_low_scale,
        # original 3
        lora1_preset, lora1_high_scale, lora1_low_scale,
        lora2_preset, lora2_high_scale, lora2_low_scale,
        lora3_preset, lora3_high_scale, lora3_low_scale,
        # my 5
        my1_repo, my1_high, my1_hs, my1_low, my1_ls,
        my2_repo, my2_high, my2_hs, my2_low, my2_ls,
        my3_repo, my3_high, my3_hs, my3_low, my3_ls,
        my4_repo, my4_high, my4_hs, my4_low, my4_ls,
        my5_repo, my5_high, my5_hs, my5_low, my5_ls,
        play_result_video,
    ]

    speed_preset_dd.change(
        fn=_apply_speed_preset,
        inputs=[speed_preset_dd],
        outputs=[steps_slider, flow_shift_slider, scheduler_dropdown],
    )

    gen_event = generate_button.click(
        fn=generate_video,
        inputs=ui_inputs,
        outputs=[video_output, file_output, seed_input, status_box],
    ).then(
        # Keeps the clip when the Space has persistent storage, then grows the
        # stitched file if auto-merge is on. Both are CPU-only.
        fn=lambda path: (_persist(path), gr.update())[1],
        inputs=[file_output],
        outputs=[status_box],
        show_progress="hidden",
    ).then(
        fn=auto_queue,
        inputs=[file_output, auto_merge, merge_queue, merge_name],
        outputs=[merged_video, merged_file, merge_queue, merge_status],
        show_progress="hidden",
    )

    stop_button.click(fn=None, inputs=None, outputs=None, cancels=[gen_event])

    clear_button.click(
        fn=_clear_all,
        inputs=None,
        outputs=[input_image_component, last_image_component, prompt_input,
                 video_output, file_output, status_box],
    )

    grab_frame_btn.click(
        fn=None,
        inputs=None,
        outputs=[timestamp_box],
        js=get_timestamp_js
    )

    timestamp_box.change(
        fn=extract_frame,
        inputs=[video_output, timestamp_box],
        outputs=[input_image_component]
    )

    grab_frame_last_btn.click(
        fn=None,
        inputs=None,
        outputs=[timestamp_box_last],
        js=get_timestamp_js_last,
    )

    timestamp_box_last.change(
        fn=extract_frame,
        inputs=[video_output, timestamp_box_last],
        outputs=[last_image_component],
    )

    boomerang_full_btn.click(
        fn=lambda path: boomerang(path, "full"),
        inputs=[video_output],
        outputs=[video_output, file_output],
    )
    boomerang_end_btn.click(
        fn=lambda path, tail: boomerang(path, "ending", tail),
        inputs=[video_output, boomerang_tail],
        outputs=[video_output, file_output],
    )

    merge_add_btn.click(
        fn=add_to_queue,
        inputs=[video_output, merge_queue, merge_name],
        outputs=[merged_video, merged_file, merge_queue, merge_status],
    )
    merge_clear_btn.click(
        fn=clear_queue,
        inputs=None,
        outputs=[merged_video, merged_file, merge_queue, merge_status],
    )

    _SETTINGS_FIELDS = [
        prompt_input, negative_prompt_input, duration_seconds_input, steps_slider, quality_slider,
        seed_input, randomize_seed_checkbox, scheduler_dropdown, flow_shift_slider, frame_multi,
        safe_mode_checkbox,
        speed_preset_dd, speed_high_scale, speed_low_scale,
        lora1_preset, lora1_high_scale, lora1_low_scale,
        lora2_preset, lora2_high_scale, lora2_low_scale,
        lora3_preset, lora3_high_scale, lora3_low_scale,
        my1_repo, my1_high, my1_hs, my1_low, my1_ls,
        my2_repo, my2_high, my2_hs, my2_low, my2_ls,
        my3_repo, my3_high, my3_hs, my3_low, my3_ls,
        my4_repo, my4_high, my4_hs, my4_low, my4_ls,
        my5_repo, my5_high, my5_hs, my5_low, my5_ls,
    ]

    profile_save_btn.click(
        save_profile, [profile_name, *_SETTINGS_FIELDS], [profile_picker, profile_status]
    )
    profile_load_btn.click(
        load_profile, profile_picker, [*_SETTINGS_FIELDS, profile_name, profile_status]
    )
    profile_delete_btn.click(delete_profile, profile_picker, [profile_picker, profile_status])
    profile_refresh.click(refresh_profiles, profile_picker, profile_picker)
    settings_export_btn.click(export_settings, _SETTINGS_FIELDS, settings_download)
    settings_upload.upload(import_settings, settings_upload, _SETTINGS_FIELDS)
    # Repopulate on every page open, so profiles saved elsewhere show up without a restart.
    demo.load(refresh_profiles, profile_picker, profile_picker)

    _ESTIMATE_INPUTS = [
        input_image_component, duration_seconds_input, steps_slider,
        guidance_scale_input, frame_multi, safe_mode_checkbox,
    ]
    # Recalculated on demand only. An automatic readout means an event on every slider
    # move, and a browser that never gets to rest.
    pb_button.click(
        fn=enhance_prompt,
        inputs=[prompt_input, pb_use_model, pb_camera, pb_framing, pb_lighting, pb_pace, pb_quality],
        outputs=[prompt_input, pb_status],
    )

    gpu_estimate_btn.click(
        fn=_gpu_estimate_text,
        inputs=_ESTIMATE_INPUTS,
        outputs=[gpu_estimate],
        show_progress="hidden",
    )

if __name__ == "__main__":
    demo.queue().launch(
        mcp_server=True,
        show_error=True,
    )