1"""OpenAI-style /v1/images/generations server for MiniMax-H3-Image on one GB10 (m64 capsule).23Written for this capsule (gx10 has no start script for this model). It follows diffusers'4MiniMax-H3-Image docs (v0.43.0, the first release with the pipeline) with one change for the5GB10's 119 GB of unified memory: the transformer (~62 GB) and the Qwen3-VL-32B text encoder6(~62 GB) do not both fit in bf16, and CPU offload frees nothing when CPU and GPU share one7memory. So the text encoder is cut down to what the pipeline reads and quantized:8- the pipeline conditions on hidden_states[50] of its 64 decoder layers and never runs the9  language-model head, so only the first 51 layers are loaded (diffusers' encoder step accepts10  any stack of more than 50 layers; hidden_states[50] is unchanged) and the head is dropped;11- those layers load in 4-bit (bitsandbytes nf4, bf16 compute), about 15 GB in all.12The transformer and the VAE stay bf16, as in the docs. All components stay on the GPU, and the13whole pipeline peaks around 85 GB.1415The checkpoint is guidance-distilled and samples on a fixed 6-step grid with one output head16per step: there is no num_inference_steps, guidance_scale or negative_prompt. The size is a17baked constant (the capsule's [params]); a request carries only `prompt` and `seed`.18Progress is reported as JSON lines appended to M64_STATUS.19"""2021import base6422import io23import json24import os25import secrets26import threading27import time28import traceback29from typing import Optional3031import torch32import uvicorn33from fastapi import FastAPI, HTTPException34from fastapi.responses import JSONResponse35from pydantic import BaseModel, Field3637SERVED_MODEL_ID = "minimax-h3-image"38WEIGHTS = "/weights/minimax-h3-image"39PORT = int(os.environ["M64_PORT"])40STATUS = os.environ["M64_STATUS"]4142# Baked generation settings. Keep in sync with [params] in capsule.toml.43SIZE = "2048x2048"44STEPS = 6  # fixed by the checkpoint (its sample_sigmas); reported, not passed45TEXT_ENCODER_LAYERS = 51  # of 64: hidden_states[50] is the conditioning46TEXT_ENCODER_QUANT = "nf4 (bitsandbytes), first 51 of 64 layers"47WIDTH, HEIGHT = (int(x) for x in SIZE.split("x"))48MAX_SEED = 2**63 - 14950app = FastAPI()51pipe = None52lock = threading.Lock()  # one generation at a time on the single pipeline535455def status(phase: str, **extra) -> None:56    with open(STATUS, "a") as f:57        f.write(json.dumps({"phase": phase, **extra}) + "\n")585960def run(prompt: str, seed: int, height: int = HEIGHT, width: int = WIDTH):61    # The docs seed a CPU generator.62    return pipe(prompt=prompt, height=height, width=width,63                generator=torch.Generator().manual_seed(seed), output="images")646566def load_model() -> None:67    global pipe68    from diffusers import (AutoencoderKLMiniMaxH3Image, MiniMaxH3ImageTransformer2DModel,69                           MiniMaxH3Scheduler, ModularPipeline)70    from transformers import (AutoConfig, BitsAndBytesConfig, Qwen2TokenizerFast,71                              Qwen3VLForConditionalGeneration, Qwen3VLProcessor)7273    t0 = time.time()74    # The pipeline's config (its sample_sigmas) comes from modular_model_index.json; its component75    # specs point at the Hub, so every component is loaded here from the mount and handed over.76    p = ModularPipeline.from_pretrained(WEIGHTS)7778    status("loading", progress=0.0, detail="loading the Qwen3-VL-32B text encoder (51 layers, nf4, ~15 GB)")79    print("Loading the text encoder (first 51 layers, nf4)...", flush=True)80    te_dir = os.path.join(WEIGHTS, "text_encoder")81    cfg = AutoConfig.from_pretrained(te_dir)82    full_layers = cfg.text_config.num_hidden_layers83    cfg.text_config.num_hidden_layers = TEXT_ENCODER_LAYERS  # the weights of the rest are skipped84    text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(85        te_dir, config=cfg, dtype=torch.bfloat16, device_map="cuda",86        quantization_config=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",87                                               bnb_4bit_compute_dtype=torch.bfloat16))88    text_encoder.lm_head = torch.nn.Identity()  # never run: the pipeline reads hidden states only89    text_encoder.eval()90    torch.cuda.empty_cache()91    print(f"  {TEXT_ENCODER_LAYERS} of {full_layers} layers; GPU memory "92          f"{torch.cuda.memory_allocated() / 2**30:.1f} GiB", flush=True)93    print(f"  text encoder in {time.time() - t0:.0f}s", flush=True)9495    status("loading", progress=0.35, detail="loading the 33B transformer in bf16 (~62 GB)")96    print("Loading the transformer (bf16)...", flush=True)97    transformer = MiniMaxH3ImageTransformer2DModel.from_pretrained(98        WEIGHTS, subfolder="transformer_diffusers", torch_dtype=torch.bfloat16).to("cuda")99    print(f"  transformer done at {time.time() - t0:.0f}s", flush=True)100101    status("loading", progress=0.9, detail="loading the VAE, tokenizer and scheduler")102    vae = AutoencoderKLMiniMaxH3Image.from_pretrained(103        WEIGHTS, subfolder="vae_diffusers", torch_dtype=torch.bfloat16).to("cuda")104    p.update_components(105        text_encoder=text_encoder, transformer=transformer, vae=vae,106        tokenizer=Qwen2TokenizerFast.from_pretrained(os.path.join(WEIGHTS, "tokenizer")),107        processor=Qwen3VLProcessor.from_pretrained(os.path.join(WEIGHTS, "processor")),108        scheduler=MiniMaxH3Scheduler.from_pretrained(WEIGHTS, subfolder="scheduler"),109    )110    pipe = p111    print(f"Model loaded in {time.time() - t0:.1f}s | GPU memory allocated "112          f"{torch.cuda.memory_allocated() / 2**30:.1f} GiB", flush=True)113114    status("warming", detail="one small warm-up generation")115    t0 = time.time()116    with lock:117        # Generated at the 4 MP minimum canvas whatever is asked; this only warms the kernels.118        run("warm-up", 0)119    print(f"Warm-up done in {time.time() - t0:.1f}s.", flush=True)120121122class ImageRequest(BaseModel):123    prompt: str = Field(..., min_length=1, description="Text prompt")124    seed: Optional[int] = Field(None, ge=0, le=MAX_SEED, description="RNG seed. null = random; the one used is returned.")125126127def baked() -> dict:128    return {"size": SIZE, "num_inference_steps": STEPS, "guidance": "distilled (none)",129            "text_encoder": TEXT_ENCODER_QUANT}130131132@app.get("/health")133def health():134    if pipe is None:135        raise HTTPException(503, "Model not loaded")136    return {"status": "ok", "model": SERVED_MODEL_ID}137138139@app.get("/v1/config")140def config():141    return {142        "model": SERVED_MODEL_ID,143        "params": baked(),144        "request_fields": {145            "prompt": {"required": True, "description": "Text prompt, used verbatim (Qwen3-VL-32B encoder, "146                                                        "up to 4096 tokens); long, complete descriptions work best"},147            "seed": {"required": False, "default": None, "description": "RNG seed; null = random"},148        },149    }150151152@app.post("/v1/images/generations")153def generate(req: ImageRequest):154    if pipe is None:155        raise HTTPException(503, "Model not loaded")156    seed = req.seed if req.seed is not None else secrets.randbelow(2**31)157    t0 = time.time()158    with lock:159        images = run(req.prompt, seed)160    elapsed = time.time() - t0161    print(f"Generated {len(images)} image(s) in {elapsed:.1f}s (seed {seed})", flush=True)162    data = []163    for img in images:164        buf = io.BytesIO()165        img.save(buf, format="PNG")166        data.append({"b64_json": base64.b64encode(buf.getvalue()).decode()})167    return JSONResponse({168        "created": int(time.time()),169        "model": SERVED_MODEL_ID,170        "seed": seed,171        "elapsed_seconds": round(elapsed, 2),172        "data": data,173        "params_used": {**baked(), "seed": seed},174    })175176177def main() -> None:178    try:179        load_model()180    except Exception:181        traceback.print_exc()182        status("error", detail="model load failed; see container logs")183        raise184    status("ready")185    uvicorn.run(app, host="0.0.0.0", port=PORT)  # the gate reaches us over the capsule network186187188if __name__ == "__main__":189    main()