"""UniSpace — unified text-to-image generation and instruction-based image editing.

Paper: UniSpace: Unified Visual Representation and Scalable Multimodal Modeling
Model: https://huggingface.co/yjb6/UniSpace
Code:  https://github.com/yjb6/UniSpace
"""

import os

# Must be set before torch / the UniSpace modeling code is imported.
os.environ.setdefault("ATTENTION_BACKEND", "torch_sdpa")
os.environ.setdefault("USE_FLEX_ATTENTION", "0")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  (must precede torch / CUDA-touching imports)

import gc  # noqa: E402
import json  # noqa: E402
import random  # noqa: E402
import sys  # noqa: E402
import time  # noqa: E402

import gradio as gr  # noqa: E402
import torch  # noqa: E402
from huggingface_hub import hf_hub_download  # noqa: E402
from omegaconf import OmegaConf  # noqa: E402
from safetensors import safe_open  # noqa: E402

ROOT = os.path.dirname(os.path.abspath(__file__))
PR_ROOT = os.path.join(ROOT, "patch-reparameterization")
sys.path.insert(0, os.path.join(PR_ROOT, "src"))
sys.path.insert(0, os.path.join(ROOT, "unispace"))

MODEL_REPO = "yjb6/UniSpace"
CKPT_DIR = "unispace-sft-0012000"
QWEN_LLM = "Qwen/Qwen3-8B"
QWEN_VL = "Qwen/Qwen3-VL-8B-Instruct"

# Official inference settings (unispace/scripts/eval/run_release_generation.sh +
# unispace/eval/gen/eval_sft_0012000.example.json).
MAX_LATENT_SIZE = 96
TIMESTEP_SHIFT = 0.112
DEFAULT_STEPS = 50
DEFAULT_CFG = 10.0
GEN_CFG_RENORM_MIN = 0.9   # gen_images_qwen3_unified_mot.py main()
EDIT_CFG_RENORM_MIN = 0.0  # gen_images_imgedit_qwen3_unified_mot.py argparse default
EDIT_RESOLUTION = 1024

DEVICE = "cuda"
DEVICE_TYPE = "cuda"


def _log(msg):
    print(f"[unispace] {msg}", flush=True)


def _free_cached_file(path):
    """Delete a hub-cached file (symlink + blob) to keep the Space disk small."""
    try:
        real = os.path.realpath(path)
        if os.path.islink(path):
            os.unlink(path)
        if os.path.exists(real):
            os.remove(real)
    except OSError as exc:  # pragma: no cover
        _log(f"could not free {path}: {exc}")


# ---------------------------------------------------------------------------
# 1. Vision-tower-only loader for the PatchReparam encoder.
#
# stage1/encoders/qwen3_unified.py loads the *full* Qwen3-VL-8B-Instruct
# (17.5 GB) only to keep `.visual`, and does so with device_map="cuda:N",
# which is incompatible with the ZeroGPU hijack. We swap in a loader that
# builds the vision tower on CPU and pulls its weights from the single shard
# that holds them (2.7 GB).
# ---------------------------------------------------------------------------
import transformers.models.qwen3_vl.modeling_qwen3_vl as _q3vl_mod  # noqa: E402
from transformers import AutoConfig, AutoTokenizer  # noqa: E402
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionModel  # noqa: E402

try:
    from transformers.modeling_utils import no_init_weights
except ImportError:  # pragma: no cover
    from contextlib import nullcontext as no_init_weights


class _VisionOnlyQwen3VL:
    """Minimal stand-in exposing the two attributes Qwen3Unified touches."""

    def __init__(self, visual):
        self.visual = visual
        self.model = None

    @classmethod
    def from_pretrained(cls, model_name, torch_dtype=None, device_map=None, **kwargs):
        cfg = AutoConfig.from_pretrained(model_name)
        vis_cfg = cfg.vision_config
        vis_cfg._attn_implementation = "sdpa"
        t0 = time.time()
        with no_init_weights():
            visual = Qwen3VLVisionModel(vis_cfg)
        visual = visual.to(torch.float32)
        shard = hf_hub_download(model_name, "model-00004-of-00004.safetensors")
        prefix = "model.visual."
        state = {}
        with safe_open(shard, framework="pt", device="cpu") as f:
            for key in f.keys():
                if key.startswith(prefix):
                    state[key[len(prefix):]] = f.get_tensor(key).to(torch.float32)
        missing, unexpected = visual.load_state_dict(state, strict=False)
        _log(
            f"vision tower: {len(state)} tensors, missing={len(missing)}, "
            f"unexpected={len(unexpected)} ({time.time() - t0:.1f}s)"
        )
        if missing:
            _log(f"vision tower missing keys (first 10): {missing[:10]}")
        del state
        _free_cached_file(shard)
        gc.collect()
        return cls(visual)


_q3vl_mod.Qwen3VLForConditionalGeneration = _VisionOnlyQwen3VL


# ---------------------------------------------------------------------------
# 2. PatchReparam ("RAE") tokenizer / decoder
# ---------------------------------------------------------------------------
from modeling.autoencoder_unified import load_unified_vae  # noqa: E402
from modeling.qwen3.configuration_qwen3 import Qwen3Config  # noqa: E402
from modeling.unimm.qwen3_mot import Qwen3MoTForConditionalGeneration  # noqa: E402
from modeling.unimm.unimm_mot import UnimmConfig, UnimmMoT  # noqa: E402

import unispace_infer as ui  # noqa: E402


def _build_vae_config_yaml():
    src = os.path.join(PR_ROOT, "configs", "vae", "qwen3-temp-run-mar1024_eval.yaml")
    cfg = OmegaConf.load(src)
    cfg.stage_1.params.encoder_config_path = QWEN_VL
    cfg.stage_1.params.encoder_params.model_name = QWEN_VL
    cfg.stage_1.params.normalization_stat_path = hf_hub_download(
        MODEL_REPO, "stats/pr-qwen-vit-normalization-stats.pt"
    )
    cfg.stage_1.checkpoint.path = hf_hub_download(
        MODEL_REPO, "encoders/pr-qwen-vit-tokenizer.pt"
    )
    out = os.path.join(ROOT, "_vae_eval_resolved.yaml")
    OmegaConf.save(cfg, out, resolve=False)
    return out, str(cfg.stage_1.checkpoint.path)


_log("loading PatchReparam tokenizer/decoder ...")
_t0 = time.time()
_vae_yaml, _vae_ckpt = _build_vae_config_yaml()
vae_model, vae_config = load_unified_vae(local_path=_vae_yaml, patch_reparam_root=PR_ROOT)
vae_model = vae_model.eval()
_free_cached_file(_vae_ckpt)
gc.collect()
_log(
    f"PatchReparam ready in {time.time() - _t0:.1f}s "
    f"(downsample={vae_config.downsample}, z_channels={vae_config.z_channels})"
)


# ---------------------------------------------------------------------------
# 3. Qwen3-8B Mixture-of-Transformers backbone
#
# Every `language_model.*` tensor lives in the UniSpace checkpoint, so the base
# Qwen3-8B weights are never needed — only its config + tokenizer. The 17-shard
# fp32 checkpoint (61 GB) is streamed one shard at a time, cast to bf16 and
# released, so peak disk stays around 4 GB.
# ---------------------------------------------------------------------------
_log("building Qwen3-8B MoT backbone from config ...")
_t0 = time.time()
llm_config = Qwen3Config.from_pretrained(QWEN_LLM)
llm_config.qk_norm = True
llm_config.tie_word_embeddings = False
llm_config._attn_implementation = "sdpa"

_orig_linear_reset = torch.nn.Linear.reset_parameters
_orig_embedding_reset = torch.nn.Embedding.reset_parameters
torch.nn.Linear.reset_parameters = lambda self: None
torch.nn.Embedding.reset_parameters = lambda self: None
torch.set_default_dtype(torch.bfloat16)
try:
    with no_init_weights():
        language_model = Qwen3MoTForConditionalGeneration(llm_config)
finally:
    torch.set_default_dtype(torch.float32)
    torch.nn.Linear.reset_parameters = _orig_linear_reset
    torch.nn.Embedding.reset_parameters = _orig_embedding_reset
_log(f"backbone skeleton built in {time.time() - _t0:.1f}s")

unimm_config = UnimmConfig(
    visual_gen=True,
    visual_und=True,
    llm_config=llm_config,
    vit_config=None,
    vae_config=vae_config,
    vit_max_num_patch_per_side=MAX_LATENT_SIZE,
    connector_act="gelu_pytorch_tanh",
    latent_patch_size=1,
    max_latent_size=MAX_LATENT_SIZE,
    use_qwen_vit=False,
    use_moe=True,
    use_qwen3_unified=True,
    share_unified2llm=False,
    use_spatial_merge=False,
    use_spatial_merge_gen=False,
    use_spatial_merge_und=True,
    use_mrope=False,
)

tokenizer = AutoTokenizer.from_pretrained(QWEN_LLM)
model = UnimmMoT(language_model, None, tokenizer, unimm_config)
del language_model
gc.collect()

new_token_ids = {
    "bos_token_id": tokenizer.encode("<|im_start|>")[0],
    "eos_token_id": tokenizer.encode("<|im_end|>")[0],
    "pad_token_id": tokenizer.pad_token_id or tokenizer.eos_token_id,
    "start_of_image": tokenizer.encode("<|vision_start|>")[0],
    "end_of_image": tokenizer.encode("<|vision_end|>")[0],
}

# The two positional embeddings are frozen 2-D sincos tables rebuilt at
# construction time — the official edit script drops them from the checkpoint.
SKIP_KEYS = {"latent_pos_embed.pos_embed", "vit_pos_embed.pos_embed"}

_log("streaming UniSpace SFT checkpoint ...")
_t0 = time.time()
_index = json.load(
    open(hf_hub_download(MODEL_REPO, f"{CKPT_DIR}/model.safetensors.index.json"))
)
_shards = sorted(set(_index["weight_map"].values()))
_loaded = set()
for _i, _shard in enumerate(_shards, 1):
    _path = hf_hub_download(MODEL_REPO, f"{CKPT_DIR}/{_shard}")
    _sd = {}
    with safe_open(_path, framework="pt", device="cpu") as f:
        for _k in f.keys():
            if _k in SKIP_KEYS:
                continue
            _sd[_k] = f.get_tensor(_k).to(torch.bfloat16)
    _msg = model.load_state_dict(_sd, strict=False)
    if _msg.unexpected_keys:
        _log(f"  shard {_i}: unexpected {_msg.unexpected_keys[:5]}")
    _loaded.update(_sd.keys())
    del _sd, _msg
    gc.collect()
    _free_cached_file(_path)
    _log(f"  shard {_i}/{len(_shards)} loaded ({time.time() - _t0:.0f}s elapsed)")

_expected = set(model.state_dict().keys())
_missing = sorted(_expected - _loaded - SKIP_KEYS)
_log(f"checkpoint loaded: {len(_loaded)} tensors, missing={len(_missing)}")
if _missing:
    _log(f"MISSING KEYS (first 20): {_missing[:20]}")

model = model.to(torch.bfloat16).eval()
for _p in model.parameters():
    _p.requires_grad_(False)
for _p in vae_model.parameters():
    _p.requires_grad_(False)
gc.collect()

_log("moving weights to GPU ...")
model = model.to(DEVICE)
vae_model = vae_model.to(DEVICE)
_log("model ready")


# ---------------------------------------------------------------------------
# 4. Gradio handlers
# ---------------------------------------------------------------------------
MAX_SEED = 2**31 - 1


# Durations are sized from measured runs on this Space (ZeroGPU large), not
# guessed. Sampling dominates and is linear in `steps` at ~1.13 s/step at
# 1024x1024, with essentially no fixed cost:
#     text-to-image  16 steps -> 17.4s | 50 steps -> 56.2s | 60 steps -> 67.0s
# Editing adds ~2s for the reference-image encode (Qwen3-VL tower + VAE) and
# runs at the same per-step cost (mar_1024 buckets keep the area ~1024^2).
# ~10s covers GPU allocation + weight streaming, then a 15% safety margin.
# Kept deliberately tight: `duration` is what a visitor's quota is charged
# against, and lower values also rank higher in the ZeroGPU queue.
_STEP_SECONDS = 1.13
_ALLOC_SECONDS = 10.0
_MARGIN = 1.15


def _duration_generate(prompt="", steps=DEFAULT_STEPS, *args, **kwargs):
    try:
        steps = int(steps)
    except (TypeError, ValueError):
        steps = DEFAULT_STEPS
    return min(100, int((_ALLOC_SECONDS + steps * _STEP_SECONDS) * _MARGIN))


def _duration_edit(image=None, instruction="", steps=DEFAULT_STEPS, *args, **kwargs):
    try:
        steps = int(steps)
    except (TypeError, ValueError):
        steps = DEFAULT_STEPS
    return min(105, int((_ALLOC_SECONDS + 2.0 + steps * _STEP_SECONDS) * _MARGIN))


@spaces.GPU(duration=_duration_generate)
def generate(
    prompt: str,
    steps: int = DEFAULT_STEPS,
    cfg_scale: float = DEFAULT_CFG,
    resolution: str = "1024",
    seed: int = 42,
    randomize_seed: bool = False,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a square image from a text prompt with UniSpace.

    Args:
        prompt: Text description of the image to generate.
        steps: Number of flow-matching sampling steps.
        cfg_scale: Classifier-free guidance strength.
        resolution: Output width/height in pixels, one of "512", "768", "1024".
        seed: Random seed.
        randomize_seed: Draw a fresh random seed instead of using `seed`.

    Returns:
        The generated PIL image, the seed that was used, and a timing note.
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    resolution = int(resolution)
    t0 = time.perf_counter()
    images = ui.generate_image(
        prompt=prompt.strip(),
        num_timesteps=int(steps),
        cfg_scale=float(cfg_scale),
        cfg_interval=[0, 1.0],
        cfg_renorm_min=GEN_CFG_RENORM_MIN,
        timestep_shift=TIMESTEP_SHIFT,
        max_t=1.0,
        num_images=1,
        resolution=resolution,
        inference_mode="flash",
        device=DEVICE,
        device_type=DEVICE_TYPE,
        gen_model=model,
        tokenizer=tokenizer,
        new_token_ids=new_token_ids,
        vae_model=vae_model,
        seed=seed,
    )
    elapsed = time.perf_counter() - t0
    return images[0], seed, f"{resolution}×{resolution} · {steps} steps · {elapsed:.1f}s"


@spaces.GPU(duration=_duration_edit)
def edit(
    image,
    instruction: str,
    steps: int = DEFAULT_STEPS,
    cfg_scale: float = DEFAULT_CFG,
    seed: int = 42,
    randomize_seed: bool = False,
    progress=gr.Progress(track_tqdm=True),
):
    """Edit an image by following a natural-language instruction.

    Args:
        image: The source image to edit.
        instruction: What to change, e.g. "Change the background to a forest.".
        steps: Number of flow-matching sampling steps.
        cfg_scale: Classifier-free guidance strength on the instruction.
        seed: Random seed.
        randomize_seed: Draw a fresh random seed instead of using `seed`.

    Returns:
        The edited PIL image, the seed that was used, and a timing note.
    """
    if image is None:
        raise gr.Error("Please upload an image to edit.")
    if not instruction or not instruction.strip():
        raise gr.Error("Please enter an editing instruction.")
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    ref = image.convert("RGB")
    target_size = ui.compute_target_size_from_ref(ref, EDIT_RESOLUTION)
    t0 = time.perf_counter()
    out = ui.editing_image(
        gen_model=model,
        tokenizer=tokenizer,
        new_token_ids=new_token_ids,
        vae_model=vae_model,
        ref_image=ref,
        prompt=instruction.strip(),
        ref_image_size=448,
        ref_max_size=None,
        ref_min_size=256,
        ref_use_mar=True,
        ref_mar_resolution=EDIT_RESOLUTION,
        target_size=target_size,
        num_timesteps=int(steps),
        cfg_text_scale=float(cfg_scale),
        cfg_interval=[0, 1.0],
        cfg_renorm_min=EDIT_CFG_RENORM_MIN,
        timestep_shift=TIMESTEP_SHIFT,
        device=DEVICE,
        device_type=DEVICE_TYPE,
        seed=seed,
    )
    elapsed = time.perf_counter() - t0
    return out, seed, f"{target_size[1]}×{target_size[0]} · {steps} steps · {elapsed:.1f}s"


# Thin wrappers for gr.Examples. They deliberately take no `gr.Progress`
# parameter: Gradio injects the Progress object at the index of that parameter
# and would otherwise land it in `steps` when an example row supplies fewer
# values than the full signature.
def generate_example(prompt: str):
    """Run the text-to-image example row with default settings."""
    return generate(prompt)


def edit_example(image, instruction: str):
    """Run the image-editing example row with default settings."""
    return edit(image, instruction)


# ---------------------------------------------------------------------------
# 5. UI
# ---------------------------------------------------------------------------
GEN_EXAMPLES = [
    ["A photo of a press conference with microphones and a blurred crowd in the background, realistic news style"],
    ["A watercolor painting of a lighthouse during a winter sunrise."],
    ["An ink wash painting style image illustrating a Chinese poem. A solitary boat floats on a misty river near mountains. Vertical traditional Chinese text on the right reads: '孤舟蓑笠翁，独钓寒江雪'. The mood is serene and melancholic. A red seal stamp (chop) is visible in the bottom left corner."],
    ["A realistic photo of a modern airport terminal sign hanging from the ceiling. The sign has a yellow background with black text. It shows directions with arrows. The text is in three languages: English 'Departures', Chinese '出发', and Japanese '出発'. An icon of an airplane taking off is next to the text. The background shows a blurred terminal hall with travelers."],
    ["A professional product photo of a sneaker on a white background, studio lighting, advertising style"],
]

EDIT_EXAMPLES = [
    ["examples/background_change.png", "Change the background to a forest."],
    ["examples/color_alter.png", "Alter the color of the mirror frame to orange."],
    ["examples/material_alter.png", "Change the cup in hand to ceramic."],
    ["examples/subject-add.png", "Add a balloon decoration strip below the airplane."],
]

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="UniSpace") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🛰️ UniSpace
            **Unified Visual Representation and Scalable Multimodal Modeling** — an 8B
            Mixture-of-Transformers that generates and edits images on top of a single
            patch-reparameterized visual tokenizer.

            [Paper](https://huggingface.co/papers/2608.08676) ·
            [Model](https://huggingface.co/yjb6/UniSpace) ·
            [Code](https://github.com/yjb6/UniSpace)
            """
        )

        with gr.Tabs():
            with gr.Tab("Text to image"):
                with gr.Row():
                    with gr.Column():
                        g_prompt = gr.Textbox(
                            label="Prompt",
                            lines=3,
                            placeholder="A watercolor painting of a lighthouse during a winter sunrise.",
                        )
                        g_run = gr.Button("Generate", variant="primary")
                        with gr.Accordion("Advanced settings", open=False):
                            # String choices so the generated API/MCP schema
                            # (Literal['512','768','1024']) matches what the
                            # component actually accepts; handler casts to int.
                            g_resolution = gr.Radio(
                                choices=["512", "768", "1024"],
                                value="1024",
                                label="Resolution",
                            )
                            g_steps = gr.Slider(
                                8, 60, value=DEFAULT_STEPS, step=1, label="Sampling steps"
                            )
                            g_cfg = gr.Slider(
                                1.0, 15.0, value=DEFAULT_CFG, step=0.5,
                                label="Guidance scale (CFG)",
                            )
                            g_seed = gr.Slider(
                                0, MAX_SEED, value=42, step=1, label="Seed"
                            )
                            g_rand = gr.Checkbox(value=False, label="Randomize seed")
                    with gr.Column():
                        g_out = gr.Image(label="Result", type="pil", height=520)
                        g_used_seed = gr.Number(label="Seed used", interactive=False)
                        g_info = gr.Textbox(label="Run info", interactive=False)

                gr.Examples(
                    examples=GEN_EXAMPLES,
                    inputs=[g_prompt],
                    outputs=[g_out, g_used_seed, g_info],
                    fn=generate_example,
                    cache_examples=True,
                    cache_mode="lazy",
                    label="Prompts from the UniSpace release",
                )

            with gr.Tab("Instruction editing"):
                with gr.Row():
                    with gr.Column():
                        e_image = gr.Image(label="Input image", type="pil", height=340)
                        e_prompt = gr.Textbox(
                            label="Editing instruction",
                            lines=2,
                            placeholder="Change the background to a forest.",
                        )
                        e_run = gr.Button("Edit", variant="primary")
                        with gr.Accordion("Advanced settings", open=False):
                            e_steps = gr.Slider(
                                8, 60, value=DEFAULT_STEPS, step=1, label="Sampling steps"
                            )
                            e_cfg = gr.Slider(
                                1.0, 15.0, value=DEFAULT_CFG, step=0.5,
                                label="Guidance scale (CFG)",
                            )
                            e_seed = gr.Slider(
                                0, MAX_SEED, value=42, step=1, label="Seed"
                            )
                            e_rand = gr.Checkbox(value=False, label="Randomize seed")
                    with gr.Column():
                        e_out = gr.Image(label="Result", type="pil", height=520)
                        e_used_seed = gr.Number(label="Seed used", interactive=False)
                        e_info = gr.Textbox(label="Run info", interactive=False)

                gr.Examples(
                    examples=EDIT_EXAMPLES,
                    inputs=[e_image, e_prompt],
                    outputs=[e_out, e_used_seed, e_info],
                    fn=edit_example,
                    cache_examples=True,
                    cache_mode="lazy",
                    label="Examples from the UniSpace project page (GEdit-Bench, MIT)",
                )

    g_run.click(
        fn=generate,
        inputs=[g_prompt, g_steps, g_cfg, g_resolution, g_seed, g_rand],
        outputs=[g_out, g_used_seed, g_info],
        api_name="generate",
    )
    e_run.click(
        fn=edit,
        inputs=[e_image, e_prompt, e_steps, e_cfg, e_seed, e_rand],
        outputs=[e_out, e_used_seed, e_info],
        api_name="edit",
    )

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