"""
Multi-Model Image Upscaler — ZeroGPU Space.

Загружает модели лениво (по требованию) через spandrel, поэтому старт Space
быстрый, а квота ZeroGPU тратится только на само инференс-время.
"""

import gc
import os
import sys
import tempfile

import time
import uuid

# --------------------------------------------------------------------------
# Подавление шумных трейсбеков ZeroGPU.
#
# ZeroGPU форкает процессы для выделения GPU, и дочерние процессы наследуют
# asyncio event loop с уже закрытыми файловыми дескрипторами. При сборке мусора
# BaseEventLoop.__del__ пытается закрыть fd = -1 и печатает многострочный
# "Exception ignored in: <function BaseEventLoop.__del__>".
#
# Такие исключения идут не через except, а через sys.unraisablehook, поэтому
# ловим именно его. Всё, что не относится к этому конкретному шуму, пробрасываем
# в стандартный обработчик, чтобы не прятать настоящие ошибки.
# --------------------------------------------------------------------------
_default_unraisablehook = sys.unraisablehook

_NOISY_MESSAGES = ("Invalid file descriptor", "Event loop is closed")


def _quiet_unraisablehook(unraisable):
    exc = unraisable.exc_value
    obj_repr = repr(unraisable.object)

    is_loop_teardown = "BaseEventLoop.__del__" in obj_repr or "eventloop" in obj_repr.lower()
    is_noisy_msg = isinstance(exc, ValueError) and any(
        m in str(exc) for m in _NOISY_MESSAGES
    )

    if is_loop_teardown and is_noisy_msg:
        return  # молча глотаем — это безвредный шум ZeroGPU

    _default_unraisablehook(unraisable)


sys.unraisablehook = _quiet_unraisablehook

import gradio as gr
import numpy as np
import spaces
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from spandrel import ImageModelDescriptor, ModelLoader

# --------------------------------------------------------------------------
# Каталог моделей: имя -> (HF repo, файл весов, множитель, описание)
# Добавить свою модель = дописать строчку сюда.
# --------------------------------------------------------------------------
MODELS = {
    "4xNomosWebPhoto_RealPLKSR": (
        "Phips/4xNomosWebPhoto_RealPLKSR",
        "4xNomosWebPhoto_RealPLKSR.safetensors",
        4,
        "📷 Фото из интернета — лучший универсальный выбор",
    ),
    "4xNomos2_realplksr_dysample": (
        "Phips/4xNomos2_realplksr_dysample",
        "4xNomos2_realplksr_dysample.safetensors",
        4,
        "📷 Фото, чуть резче, быстрый",
    ),
    "4xNomos8kSCHAT-L": (
        "Phips/4xNomos8kSCHAT-L",
        "4xNomos8kSCHAT-L.safetensors",
        4,
        "🎯 Максимальное качество (HAT-L, медленный)",
    ),
    "2xNomosUni_span_multijpg": (
        "Phips/2xNomosUni_span_multijpg",
        "2xNomosUni_span_multijpg.safetensors",
        2,
        "⚡ Универсальный 2x, очень быстрый",
    ),
    "4xNomosUni_span_multijpg": (
        "Phips/4xNomosUni_span_multijpg",
        "4xNomosUni_span_multijpg.safetensors",
        4,
        "⚡ Универсальный 4x, быстрый",
    ),
    "2xHFA2kCompact": (
        "Phips/2xHFA2kCompact",
        "2xHFA2kCompact.safetensors",
        2,
        "🎨 Аниме / иллюстрации 2x",
    ),
    "4xFaceUpDAT": (
        "Phips/4xFaceUpDAT",
        "4xFaceUpDAT.safetensors",
        4,
        "🙂 Портреты и лица",
    ),
    "1xDeJPG_realplksr_otf": (
        "Phips/1xDeJPG_realplksr_otf",
        "1xDeJPG_realplksr_otf.safetensors",
        1,
        "🧹 Только чистка JPEG-артефактов (без увеличения)",
    ),
}

DEFAULT_MODEL = "4xNomosWebPhoto_RealPLKSR"

# --------------------------------------------------------------------------
# Qwen-Image-Edit-2509 + Upscale LoRA (генеративный апскейл).
#
# Это не spandrel-модель, а LoRA-адаптер к 20B диффузионке, поэтому живёт
# отдельно от MODELS. Веса ~37 ГБ, грузятся в фоновом потоке при старте:
# так восемь быстрых моделей доступны сразу, а Qwen подключается, когда готов.
# --------------------------------------------------------------------------
QWEN_BASE = "Qwen/Qwen-Image-Edit-2509"
QWEN_TRANSFORMER = "linoyts/Qwen-Image-Edit-Rapid-AIO"  # дистиллят, 8 шагов
QWEN_LORA_REPO = "vafipas663/Qwen-Edit-2509-Upscale-LoRA"
# Внимание: официальное демо падает с 404, потому что просит
# qwen-edit-enhance_64-v3_000001500.safetensors — автор перенёс этот чекпоинт
# в archive/. Ниже — актуальное имя, проверено что отдаёт 200.
QWEN_LORA_FILE = "qwen-edit-enhance_64-v3_000001000.safetensors"

QWEN_LABEL = "🧠 Qwen-Edit Upscale LoRA (генеративный)"

_qwen_pipe = None
_qwen_status = "not_started"  # not_started | loading | ready | failed | disabled
_qwen_error = ""


def _load_qwen():
    """
    Загружает Qwen-пайплайн НА УРОВНЕ МОДУЛЯ, В ГЛАВНОМ ПОТОКЕ.

    Так требует ZeroGPU: вне @spaces.GPU работает эмуляция CUDA, которая
    перехватывает операции и позволяет "положить" модель на cuda без реального
    GPU. Эмуляция активна только в главном потоке при импорте — попытка
    загрузить модель в фоновом потоке падает с
    "Low-level CUDA init (torch._C._cuda_init) reached".

    Загрузка блокирует старт примерно на 2-4 минуты (Xet качает ~37 ГБ
    на скорости ~700 МБ/с). Ошибка здесь не роняет приложение: быстрые
    spandrel-модели продолжат работать, вкладка Qwen покажет причину.
    """
    global _qwen_pipe, _qwen_status, _qwen_error

    if os.environ.get("ENABLE_QWEN", "1") == "0":
        _qwen_status = "disabled"
        print("[qwen] отключён через ENABLE_QWEN=0", flush=True)
        return

    _qwen_status = "loading"
    print("[qwen] загрузка (~37 ГБ), это займёт несколько минут...", flush=True)
    started = time.time()

    try:
        from diffusers import QwenImageEditPlusPipeline
        from diffusers.models import QwenImageTransformer2DModel

        token = os.environ.get("HF_TOKEN") or None
        dtype = torch.bfloat16

        transformer = QwenImageTransformer2DModel.from_pretrained(
            QWEN_TRANSFORMER,
            subfolder="transformer",
            dtype=dtype,
            device_map="cuda" if torch.cuda.is_available() else None,
            token=token,
        )
        pipe = QwenImageEditPlusPipeline.from_pretrained(
            QWEN_BASE,
            transformer=transformer,
            dtype=dtype,
            token=token,
        )
        if torch.cuda.is_available():
            pipe = pipe.to("cuda")

        # Вплавляем LoRA в веса: инференс быстрее, памяти меньше
        pipe.load_lora_weights(
            QWEN_LORA_REPO,
            weight_name=QWEN_LORA_FILE,
            adapter_name="upscale",
            token=token,
        )
        pipe.set_adapters(["upscale"], adapter_weights=[1.0])
        pipe.fuse_lora(adapter_names=["upscale"], lora_scale=1.0)
        pipe.unload_lora_weights()

        _qwen_pipe = pipe
        _qwen_status = "ready"
        print(f"[qwen] готов за {time.time() - started:.0f} с", flush=True)

    except Exception as exc:  # noqa: BLE001
        _qwen_status = "failed"
        _qwen_error = f"{type(exc).__name__}: {exc}"
        print(f"[qwen] загрузка не удалась: {_qwen_error}", flush=True)

# Ограничения, чтобы не ловить OOM и не жечь квоту
MAX_INPUT_PIXELS = 1_200_000  # ~1400x850 на входе
MAX_OUTPUT_PIXELS = 20_000_000  # ~5000x4000 на выходе
QWEN_MAX_INPUT_PIXELS = 1_048_576  # ~1024x1024 — больше не влезает в VRAM
TILE_SIZE = 512
TILE_OVERLAP = 32

_cache: dict[str, ImageModelDescriptor] = {}


def _device() -> torch.device:
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def load_model(name: str) -> ImageModelDescriptor:
    """Скачивает (с кешем на диске Space) и загружает модель."""
    if name in _cache:
        return _cache[name]

    repo, filename, _, _ = MODELS[name]
    # HF_TOKEN подхватывается из секретов Space, если он там задан:
    # с ним выше rate-limit на скачивание весов. Без него тоже работает.
    path = hf_hub_download(
        repo_id=repo, filename=filename, token=os.environ.get("HF_TOKEN") or None
    )
    model = ModelLoader().load_from_file(path)

    if not isinstance(model, ImageModelDescriptor):
        raise gr.Error(f"Модель {name} не является image-to-image моделью")

    model.eval()
    # держим в кеше максимум 2 модели, чтобы не забить VRAM
    if len(_cache) >= 2:
        old = next(iter(_cache))
        del _cache[old]
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    _cache[name] = model
    return model


def _to_tensor(img: Image.Image) -> torch.Tensor:
    arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
    return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)


def _to_image(t: torch.Tensor) -> Image.Image:
    arr = t.squeeze(0).permute(1, 2, 0).clamp(0, 1).cpu().numpy()
    return Image.fromarray((arr * 255.0 + 0.5).astype(np.uint8))


@torch.inference_mode()
def _run_tiled(model, tensor: torch.Tensor, scale: int) -> torch.Tensor:
    """Прогон по тайлам с перекрытием — иначе большие картинки дают OOM."""
    _, _, h, w = tensor.shape
    if h <= TILE_SIZE and w <= TILE_SIZE:
        return model(tensor)

    out = torch.zeros(1, 3, h * scale, w * scale, device=tensor.device)
    weight = torch.zeros_like(out)
    step = TILE_SIZE - TILE_OVERLAP

    for y in range(0, h, step):
        for x in range(0, w, step):
            y2, x2 = min(y + TILE_SIZE, h), min(x + TILE_SIZE, w)
            y1, x1 = max(0, y2 - TILE_SIZE), max(0, x2 - TILE_SIZE)
            tile = model(tensor[:, :, y1:y2, x1:x2])
            out[:, :, y1 * scale : y2 * scale, x1 * scale : x2 * scale] += tile
            weight[:, :, y1 * scale : y2 * scale, x1 * scale : x2 * scale] += 1

    return out / weight.clamp(min=1)


def _qwen_duration(image, prompt, steps, *_):
    """
    Бюджет GPU-времени под запрос.

    Важно: первый вызов после старта Space дополнительно переносит ~40 ГБ весов
    из RAM в VRAM, и это занимает больше минуты. Если запросить слишком мало,
    ZeroGPU обрывает задачу с "GPU task aborted", поэтому база щедрая.
    """
    px = (image.size[0] * image.size[1]) if image is not None else 500_000
    return int(min(300, 150 + steps * 5 + px / 20_000))


@spaces.GPU(duration=_qwen_duration)
def upscale_qwen(
    image: Image.Image,
    prompt: str,
    steps: int,
    guidance: float,
    seed: int,
    randomize_seed: bool,
):
    if image is None:
        raise gr.Error("Загрузите изображение")

    if _qwen_status == "disabled":
        raise gr.Error("Qwen отключён (ENABLE_QWEN=0). Используйте быстрые модели.")
    if _qwen_status == "failed":
        raise gr.Error(f"Qwen не загрузился: {_qwen_error}")
    if _qwen_pipe is None:
        raise gr.Error("Qwen недоступен. Попробуйте перезапустить Space.")

    src = image.convert("RGB")

    # Qwen работает в latent-пространстве: слишком большой вход = OOM и слив квоты
    px = src.size[0] * src.size[1]
    if px > QWEN_MAX_INPUT_PIXELS:
        ratio = (QWEN_MAX_INPUT_PIXELS / px) ** 0.5
        new = (max(1, int(src.width * ratio)), max(1, int(src.height * ratio)))
        gr.Info(f"Вход уменьшен до {new[0]}×{new[1]} — иначе не хватит видеопамяти")
        src = src.resize(new, Image.LANCZOS)

    if randomize_seed:
        seed = int(torch.randint(0, 2**31 - 1, (1,)).item())

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

    full_prompt = (
        prompt.strip()
        or "Enhance image quality, sharp focus, fine natural detail, high resolution photo"
    )

    started = time.time()
    try:
        result = _qwen_pipe(
            image=[src],
            prompt=full_prompt,
            num_inference_steps=int(steps),
            true_cfg_scale=float(guidance),
            generator=generator,
            num_images_per_prompt=1,
        ).images[0]
    except torch.cuda.OutOfMemoryError:
        torch.cuda.empty_cache()
        raise gr.Error("Не хватило видеопамяти. Возьмите картинку поменьше.")
    except Exception as exc:  # noqa: BLE001
        msg = str(exc)
        if "quota" in msg.lower() or "aborted" in msg.lower():
            raise gr.Error(
                "Дневная квота ZeroGPU исчерпана. Один запуск Qwen стоит "
                "1–3 минуты GPU, а у бесплатного аккаунта всего 5 мин в сутки. "
                "Квота сбрасывается через 24 ч после первого использования. "
                "Быстрые модели на первой вкладке стоят ~2 с и почти не тратят квоту."
            ) from exc
        raise

    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    out_path = os.path.join(
        tempfile.gettempdir(), f"qwen_{uuid.uuid4().hex[:8]}.png"
    )
    result.save(out_path, "PNG")

    info = (
        f"**{src.width}×{src.height}** → **{result.width}×{result.height}** "
        f"· Qwen-Edit Upscale LoRA · {int(steps)} шагов · seed `{seed}` "
        f"· {time.time() - started:.1f} с"
    )
    return (src, result), out_path, info, seed


def _duration(image, model_name, *_):
    """Динамическая длительность: короткие задачи = выше приоритет в очереди."""
    if image is None:
        return 30
    px = image.size[0] * image.size[1]
    base = 20 if "HAT" in model_name or "DAT" in model_name else 8
    return int(min(120, base + px / 20_000))


@spaces.GPU(duration=_duration)
def upscale(image: Image.Image, model_name: str, downscale_after: bool):
    if image is None:
        raise gr.Error("Загрузите изображение")

    src = image.convert("RGB")
    px = src.size[0] * src.size[1]
    if px > MAX_INPUT_PIXELS:
        ratio = (MAX_INPUT_PIXELS / px) ** 0.5
        new = (max(1, int(src.width * ratio)), max(1, int(src.height * ratio)))
        gr.Info(f"Изображение уменьшено до {new[0]}x{new[1]} перед апскейлом (лимит входа)")
        src = src.resize(new, Image.LANCZOS)

    scale = MODELS[model_name][2]
    if px * scale * scale > MAX_OUTPUT_PIXELS:
        raise gr.Error(
            f"Результат был бы слишком большим. Возьмите модель с меньшим "
            f"множителем (сейчас {scale}x) или уменьшите вход."
        )

    started = time.time()
    model = load_model(model_name).to(_device())
    tensor = _to_tensor(src).to(_device())

    try:
        result = _run_tiled(model, tensor, model.scale)
    except torch.cuda.OutOfMemoryError:
        torch.cuda.empty_cache()
        raise gr.Error("Не хватило видеопамяти. Попробуйте картинку поменьше или модель полегче.")

    out = _to_image(result)

    if downscale_after:
        out = out.resize(src.size, Image.LANCZOS)

    del tensor, result
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # gr.File отдаёт файл по пути, поэтому сохраняем PNG без потерь на диск
    out_path = os.path.join(
        tempfile.gettempdir(), f"upscaled_{uuid.uuid4().hex[:8]}.png"
    )
    out.save(out_path, "PNG", optimize=True)

    info = (
        f"**{src.width}×{src.height}** → **{out.width}×{out.height}** "
        f"· модель `{model_name}` · {time.time() - started:.1f} с"
    )
    return (src, out), out_path, info


# --------------------------------------------------------------------------
# Загрузка Qwen на уровне модуля — обязательное требование ZeroGPU.
# Выполняется до старта UI и блокирует его на несколько минут.
# --------------------------------------------------------------------------
_load_qwen()


# --------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------
choices = [(f"{k} — {v[3]}", k) for k, v in MODELS.items()]

with gr.Blocks(title="Multi-Model Image Upscaler") as demo:
    gr.Markdown(
        """
        # 🔍 Multi-Model Image Upscaler
        Апскейл на ZeroGPU. Несколько моделей под разные типы картинок — фото,
        аниме, лица, чистка JPEG. Бесплатно, доступно и как API.
        """
    )

    with gr.Tabs():
        # ------------------------- Быстрые модели -------------------------
        with gr.Tab("⚡ Быстрые модели (2–3 с)"):
            with gr.Row():
                with gr.Column(scale=1):
                    inp = gr.Image(type="pil", label="Исходное изображение", height=340)
                    model_dd = gr.Dropdown(
                        choices=choices,
                        value=DEFAULT_MODEL,
                        label="Модель",
                        info="Начните с 4xNomosWebPhoto_RealPLKSR",
                    )
                    downscale = gr.Checkbox(
                        False,
                        label="Вернуть исходный размер",
                        info="Апскейлит и уменьшает обратно — работает как шумодав/деблюр",
                    )
                    btn = gr.Button("Upscale", variant="primary")

                with gr.Column(scale=2):
                    slider = gr.ImageSlider(label="До / После", type="pil", height=460)
                    info_md = gr.Markdown()
                    download = gr.File(label="Скачать PNG без потерь")

        # --------------------------- Qwen LoRA ---------------------------
        with gr.Tab(QWEN_LABEL):
            gr.Markdown(
                """
                Генеративный апскейл на **Qwen-Image-Edit-2509** с
                [Upscale LoRA](https://huggingface.co/vafipas663/Qwen-Edit-2509-Upscale-LoRA)
                от vafipas663. В отличие от моделей на первой вкладке, он не просто
                восстанавливает пиксели, а **дорисовывает** правдоподобные детали
                по текстовому описанию.

                ⚠️ Один запуск съедает **30–60 с** GPU-квоты
                (у анонима её всего 2 мин в сутки, у залогиненного — 5 мин).

                💡 Модель обучена на фотографиях. Для 2D-артов и иллюстраций
                берите `2xHFA2kCompact` с первой вкладки.
                """
            )

            qwen_state = gr.Markdown()

            with gr.Row():
                with gr.Column(scale=1):
                    q_inp = gr.Image(type="pil", label="Исходное изображение", height=300)
                    q_prompt = gr.Textbox(
                        label="Описание сцены",
                        placeholder="a close-up photo of purple petunia flowers, green leaves",
                        info="Чем детальнее описание, тем лучше результат",
                        lines=2,
                    )
                    with gr.Accordion("Параметры", open=False):
                        q_steps = gr.Slider(
                            4, 30, value=8, step=1,
                            label="Шаги",
                            info="8 — оптимум для дистиллированного трансформера",
                        )
                        q_guidance = gr.Slider(
                            1.0, 7.0, value=1.0, step=0.1,
                            label="True CFG scale",
                            info="1.0 для 8-шагового режима",
                        )
                        q_seed = gr.Number(value=0, label="Seed", precision=0)
                        q_random = gr.Checkbox(True, label="Случайный seed")
                    q_btn = gr.Button("Сгенерировать апскейл", variant="primary")

                with gr.Column(scale=2):
                    q_slider = gr.ImageSlider(label="До / После", type="pil", height=460)
                    q_info = gr.Markdown()
                    q_download = gr.File(label="Скачать PNG")

            def _qwen_state_text():
                return {
                    "ready": "✅ **Qwen готов к работе.**",
                    "disabled": "⚪ Qwen отключён (`ENABLE_QWEN=0`).",
                    "failed": f"❌ **Qwen не загрузился:** {_qwen_error}",
                }.get(_qwen_status, "⏳ Состояние неизвестно.")

            demo.load(_qwen_state_text, outputs=qwen_state)

            q_btn.click(
                upscale_qwen,
                inputs=[q_inp, q_prompt, q_steps, q_guidance, q_seed, q_random],
                outputs=[q_slider, q_download, q_info, q_seed],
                api_name="upscale_qwen",
            )

    gr.Markdown(
        """
        ### Использование через API
        ```python
        from gradio_client import Client, handle_file
        client = Client("ЗАМЕНИ_НА_СВОЙ/upscaler")
        result = client.predict(
            image=handle_file("low-res.jpg"),
            model_name="4xNomosWebPhoto_RealPLKSR",
            downscale_after=False,
            api_name="/upscale",
        )
        print(result)
        ```
        **Квота ZeroGPU:** аноним 2 мин GPU/сутки, бесплатный аккаунт 5 мин,
        PRO 40 мин. Передайте `hf_token="hf_..."` в `Client(...)`, чтобы поднять лимит.
        """
    )

    btn.click(
        upscale,
        inputs=[inp, model_dd, downscale],
        outputs=[slider, download, info_md],
        api_name="upscale",
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch(ssr_mode=False, theme=gr.themes.Soft())