"""OpenAI-style /v1/images/generations server for MiniMax-H3-Image on one GB10 (m64 capsule).

Written for this capsule (gx10 has no start script for this model). It follows diffusers'
MiniMax-H3-Image docs (v0.43.0, the first release with the pipeline) with one change for the
GB10's 119 GB of unified memory: the transformer (~62 GB) and the Qwen3-VL-32B text encoder
(~62 GB) do not both fit in bf16, and CPU offload frees nothing when CPU and GPU share one
memory. So the text encoder is cut down to what the pipeline reads and quantized:
- the pipeline conditions on hidden_states[50] of its 64 decoder layers and never runs the
  language-model head, so only the first 51 layers are loaded (diffusers' encoder step accepts
  any stack of more than 50 layers; hidden_states[50] is unchanged) and the head is dropped;
- those layers load in 4-bit (bitsandbytes nf4, bf16 compute), about 15 GB in all.
The transformer and the VAE stay bf16, as in the docs. All components stay on the GPU, and the
whole pipeline peaks around 85 GB.

The checkpoint is guidance-distilled and samples on a fixed 6-step grid with one output head
per step: there is no num_inference_steps, guidance_scale or negative_prompt. The size is a
baked constant (the capsule's [params]); a request carries only `prompt` and `seed`.
Progress is reported as JSON lines appended to M64_STATUS.
"""

import base64
import io
import json
import os
import secrets
import threading
import time
import traceback
from typing import Optional

import torch
import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field

SERVED_MODEL_ID = "minimax-h3-image"
WEIGHTS = "/weights/minimax-h3-image"
PORT = int(os.environ["M64_PORT"])
STATUS = os.environ["M64_STATUS"]

# Baked generation settings. Keep in sync with [params] in capsule.toml.
SIZE = "2048x2048"
STEPS = 6  # fixed by the checkpoint (its sample_sigmas); reported, not passed
TEXT_ENCODER_LAYERS = 51  # of 64: hidden_states[50] is the conditioning
TEXT_ENCODER_QUANT = "nf4 (bitsandbytes), first 51 of 64 layers"
WIDTH, HEIGHT = (int(x) for x in SIZE.split("x"))
MAX_SEED = 2**63 - 1

app = FastAPI()
pipe = None
lock = threading.Lock()  # one generation at a time on the single pipeline


def status(phase: str, **extra) -> None:
    with open(STATUS, "a") as f:
        f.write(json.dumps({"phase": phase, **extra}) + "\n")


def run(prompt: str, seed: int, height: int = HEIGHT, width: int = WIDTH):
    # The docs seed a CPU generator.
    return pipe(prompt=prompt, height=height, width=width,
                generator=torch.Generator().manual_seed(seed), output="images")


def load_model() -> None:
    global pipe
    from diffusers import (AutoencoderKLMiniMaxH3Image, MiniMaxH3ImageTransformer2DModel,
                           MiniMaxH3Scheduler, ModularPipeline)
    from transformers import (AutoConfig, BitsAndBytesConfig, Qwen2TokenizerFast,
                              Qwen3VLForConditionalGeneration, Qwen3VLProcessor)

    t0 = time.time()
    # The pipeline's config (its sample_sigmas) comes from modular_model_index.json; its component
    # specs point at the Hub, so every component is loaded here from the mount and handed over.
    p = ModularPipeline.from_pretrained(WEIGHTS)

    status("loading", progress=0.0, detail="loading the Qwen3-VL-32B text encoder (51 layers, nf4, ~15 GB)")
    print("Loading the text encoder (first 51 layers, nf4)...", flush=True)
    te_dir = os.path.join(WEIGHTS, "text_encoder")
    cfg = AutoConfig.from_pretrained(te_dir)
    full_layers = cfg.text_config.num_hidden_layers
    cfg.text_config.num_hidden_layers = TEXT_ENCODER_LAYERS  # the weights of the rest are skipped
    text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
        te_dir, config=cfg, dtype=torch.bfloat16, device_map="cuda",
        quantization_config=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                                               bnb_4bit_compute_dtype=torch.bfloat16))
    text_encoder.lm_head = torch.nn.Identity()  # never run: the pipeline reads hidden states only
    text_encoder.eval()
    torch.cuda.empty_cache()
    print(f"  {TEXT_ENCODER_LAYERS} of {full_layers} layers; GPU memory "
          f"{torch.cuda.memory_allocated() / 2**30:.1f} GiB", flush=True)
    print(f"  text encoder in {time.time() - t0:.0f}s", flush=True)

    status("loading", progress=0.35, detail="loading the 33B transformer in bf16 (~62 GB)")
    print("Loading the transformer (bf16)...", flush=True)
    transformer = MiniMaxH3ImageTransformer2DModel.from_pretrained(
        WEIGHTS, subfolder="transformer_diffusers", torch_dtype=torch.bfloat16).to("cuda")
    print(f"  transformer done at {time.time() - t0:.0f}s", flush=True)

    status("loading", progress=0.9, detail="loading the VAE, tokenizer and scheduler")
    vae = AutoencoderKLMiniMaxH3Image.from_pretrained(
        WEIGHTS, subfolder="vae_diffusers", torch_dtype=torch.bfloat16).to("cuda")
    p.update_components(
        text_encoder=text_encoder, transformer=transformer, vae=vae,
        tokenizer=Qwen2TokenizerFast.from_pretrained(os.path.join(WEIGHTS, "tokenizer")),
        processor=Qwen3VLProcessor.from_pretrained(os.path.join(WEIGHTS, "processor")),
        scheduler=MiniMaxH3Scheduler.from_pretrained(WEIGHTS, subfolder="scheduler"),
    )
    pipe = p
    print(f"Model loaded in {time.time() - t0:.1f}s | GPU memory allocated "
          f"{torch.cuda.memory_allocated() / 2**30:.1f} GiB", flush=True)

    status("warming", detail="one small warm-up generation")
    t0 = time.time()
    with lock:
        # Generated at the 4 MP minimum canvas whatever is asked; this only warms the kernels.
        run("warm-up", 0)
    print(f"Warm-up done in {time.time() - t0:.1f}s.", flush=True)


class ImageRequest(BaseModel):
    prompt: str = Field(..., min_length=1, description="Text prompt")
    seed: Optional[int] = Field(None, ge=0, le=MAX_SEED, description="RNG seed. null = random; the one used is returned.")


def baked() -> dict:
    return {"size": SIZE, "num_inference_steps": STEPS, "guidance": "distilled (none)",
            "text_encoder": TEXT_ENCODER_QUANT}


@app.get("/health")
def health():
    if pipe is None:
        raise HTTPException(503, "Model not loaded")
    return {"status": "ok", "model": SERVED_MODEL_ID}


@app.get("/v1/config")
def config():
    return {
        "model": SERVED_MODEL_ID,
        "params": baked(),
        "request_fields": {
            "prompt": {"required": True, "description": "Text prompt, used verbatim (Qwen3-VL-32B encoder, "
                                                        "up to 4096 tokens); long, complete descriptions work best"},
            "seed": {"required": False, "default": None, "description": "RNG seed; null = random"},
        },
    }


@app.post("/v1/images/generations")
def generate(req: ImageRequest):
    if pipe is None:
        raise HTTPException(503, "Model not loaded")
    seed = req.seed if req.seed is not None else secrets.randbelow(2**31)
    t0 = time.time()
    with lock:
        images = run(req.prompt, seed)
    elapsed = time.time() - t0
    print(f"Generated {len(images)} image(s) in {elapsed:.1f}s (seed {seed})", flush=True)
    data = []
    for img in images:
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        data.append({"b64_json": base64.b64encode(buf.getvalue()).decode()})
    return JSONResponse({
        "created": int(time.time()),
        "model": SERVED_MODEL_ID,
        "seed": seed,
        "elapsed_seconds": round(elapsed, 2),
        "data": data,
        "params_used": {**baked(), "seed": seed},
    })


def main() -> None:
    try:
        load_model()
    except Exception:
        traceback.print_exc()
        status("error", detail="model load failed; see container logs")
        raise
    status("ready")
    uvicorn.run(app, host="0.0.0.0", port=PORT)  # the gate reaches us over the capsule network


if __name__ == "__main__":
    main()