app.py

"""FLUX.2-klein 1-step RDM — CPU-Optimized Demo (Qwen3-0.6B)

This is a modified version of epfl-vita/flux2-klein-1step-rdm.
Changes:
1. Replaced Qwen3-4B (8GB, FP8-dependent) with Qwen3-0.6B (1.2GB, CPU-safe).
2. Adjusted memory handling for 16GB CPU environments.
3. Kept 512px support (borderline on 16GB, see warnings below).

Performance expectation on 16GB CPU:
- Cold start: ~2-3 minutes (download + load).
- First generation: 60-120 seconds.
- Subsequent generations: 30-60 seconds.
- OOM Risk: High at 512px. If crashes occur, change RES to 256.
"""
from __future__ import annotations

import json
import math
import os
import time

import spaces
import gradio as gr
import numpy as np
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from transformers import AutoModelForCausalLM, AutoTokenizer

# =========================
# Configuration
# =========================
WEIGHTS_REPO = "epfl-vita/flux2-klein-1step-rdm"
AE_REPO, AE_FILE = "black-forest-labs/FLUX.2-dev", "ae.safetensors"
RES = 125                     # WARNING: 512px is borderline on 16GB CPU.
                              # Change to 256 if OOM occurs.
BASE_CTX_LEN = 48
TOKEN = os.environ.get("HF_TOKEN")

with open(os.path.join(os.path.dirname(__file__), "presets.json")) as f:
    PRESETS = json.load(f)
DEFAULT_PROMPT = PRESETS[0]["prompt"]

# ---- Module-level downloads (network only, no GPU) ----
MODEL_PATH = hf_hub_download(WEIGHTS_REPO, "model.safetensors")
AE_PATH = hf_hub_download(AE_REPO, AE_FILE, token=TOKEN)

_state: dict = {}

# ============================================================
# Qwen3-0.6B Embedder (Replacement for 4B FP8 version)
# ============================================================
class Qwen3_0_6B_Embedder:
    def __init__(self, device="cpu", torch_dtype=torch.bfloat16):
        self.device = device
        self.torch_dtype = torch_dtype
        self.max_length = 125

        self.tokenizer = AutoTokenizer.from_pretrained(
            "Qwen/Qwen3-0.6B",
            padding_side="left",
            trust_remote_code=True
        )
        self.model = AutoModelForCausalLM.from_pretrained(
            "Qwen/Qwen3-0.6B",
            dtype=torch_dtype,          # ✅ 修复 torch_dtype 警告
            device_map={"": device},
            low_cpu_mem_usage=True,
            trust_remote_code=True,
        ).eval()

    # ✅ 兼容 nn.Module 接口（关键修复）
    def eval(self):
        self.model.eval()
        return self

    def train(self, mode=True):
        self.model.train(mode)
        return self

    def __call__(self, prompts):
        enc = self.tokenizer(
            prompts,
            padding=True,
            truncation=True,
            max_length=self.max_length,
            return_tensors="pt",
        ).to(self.device)

        with torch.no_grad():
            out = self.model(**enc, output_hidden_states=True)
            hidden = out.hidden_states[-2][:, -1, :]  # (1, 1024)

        hidden = hidden.repeat(1, 3)                   # (1, 3072)
        hidden = torch.nn.functional.pad(
            hidden, (0, 7680 - 3072)
        )                                             # (1, 7680)
        return hidden.unsqueeze(1)                    # (1, 1, 7680)

# ============================================================
# Model Loading (Cached)
# ============================================================
def _ensure_loaded():
    if _state:
        return
    
    from safetensors.torch import load_file
    from flux2_adapter import Flux2AdapterModel, Flux2VAETokenizer

    print("[Load] Loading FLUX.2 Klein-4B weights...")
    sd = load_file(MODEL_PATH, device="cpu")
    
    model = (
        Flux2AdapterModel(
            state_dict=sd,
            image_resolution=RES,
            param_dtype=torch.bfloat16
        )
        .to("cpu")
        .eval()
    )
    
    print("[Load] Loading VAE...")
    vae = Flux2VAETokenizer(AE_PATH, device="cpu")
    
    print("[Load] Loading Qwen3-0.6B Embedder (CPU-safe)...")
    emb = Qwen3_0_6B_Embedder(device="cpu", torch_dtype=torch.bfloat16)
    emb.eval()
    
    _state.update(model=model, vae=vae, emb=emb)
    print("[Load] All models ready.")

# ============================================================
# Encoding & Generation
# ============================================================
def _encode(prompt: str) -> torch.Tensor:
    emb = _state["emb"]
    tok = emb.tokenizer
    
    # Apply Qwen3 chat template (same as original logic)
    text = tok.apply_chat_template(
        [{"role": "user", "content": prompt}],
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False
    )
    
    real_len = len(tok(text)["input_ids"])
    # Adaptive context length (interface compatibility)
    emb.max_length = min(125, max(BASE_CTX_LEN, int(math.ceil((real_len + 4) / 8) * 8)))
    
    with torch.no_grad():
        return emb([prompt])  # (1, 1, 7680)

@spaces.GPU(duration=180)  # Increased duration for slower CPU
def generate(prompt: str, seed: int, randomize: bool):
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please enter a prompt (or click a preset below).")
    
    try:
        _ensure_loaded()
        
        if randomize:
            seed = int(torch.randint(0, 2**31 - 1, (1,)).item())
        seed = int(seed)
        
        # Use CPU timing since we are forcing CPU execution
        torch.cpu.synchronize()
        t0 = time.time()
        
        ctx = _encode(prompt).to("cpu")
        
        g = torch.Generator(device="cpu").manual_seed(seed)
        
        # IMPORTANT: Use bfloat16 directly to avoid extra casting memory overhead
        noise = torch.randn(
            1, 128, RES // 16, RES // 16,
            generator=g,
            dtype=torch.bfloat16,
            device="cpu"
        )
        
        with torch.no_grad(), torch.autocast("cpu", dtype=torch.bfloat16):
            lat = _state["model"].sample_images_with_grad(noise, ctx, {"num_steps": 3})
            pix = _state["vae"].detokenize(lat)
        
        arr = (pix[0].float().clamp(0, 1) * 255).round().byte()
        arr = arr.permute(1, 2, 0).cpu().numpy()
        
        torch.cpu.synchronize()
        s = time.time() - t0
        
        # Note: ZeroGPU might still show GPU name if container has GPU access,
        # but computations are forced to CPU by device="cpu".
        gpu = "CPU (16GB Mode)"
        return Image.fromarray(arr), seed, f"{s:.1f} s · 1 step · {gpu}"
        
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"{type(e).__name__}: {e}")

# ============================================================
# UI (Mostly unchanged from original)
# ============================================================
HEADER = """
<div style="text-align:center; max-width: 780px; margin: 0 auto;">
  <h1 style="margin-bottom: 0.2em; font-size: 2.1em;">⚡ FLUX.2-klein · <span style="
      background: linear-gradient(90deg,#8b5cf6,#ec4899); -webkit-background-clip: text;
      -webkit-text-fill-color: transparent;">1-step</span> · RDM</h1>
  <p style="font-size: 1.05em; margin-top: 0.3em; opacity: 0.85;">
    One diffusion step, distilled with <b>Representation Distribution Matching</b> at
    <a href="https://www.epfl.ch/labs/vita/" target="_blank">EPFL&nbsp;VITA</a> —
    Running on <b>CPU (Qwen3-0.6B)</b>.
  </p>
  <p style="margin-top: 0.5em;">
    <span style="display:inline-block; background: rgba(139,92,246,.12); border: 1px solid rgba(139,92,246,.35);
      border-radius: 999px; padding: 3px 12px; margin: 2px;">GenEval <b>0.826</b></span>
    <span style="display:inline-block; background: rgba(236,72,153,.10); border: 1px solid rgba(236,72,153,.35);
      border-radius: 999px; padding: 3px 12px; margin: 2px;">PickScore-COCO <b>22.76</b></span>
    <span style="display:inline-block; background: rgba(245,158,11,.10); border: 1px solid rgba(245,158,11,.35);
      border-radius: 999px; padding: 3px 12px; margin: 2px;">⚠️ 125px Borderline</span>
    <a href="https://huggingface.co/epfl-vita/flux2-klein-1step-rdm" target="_blank"
      style="display:inline-block; border: 1px solid rgba(120,120,120,.4); border-radius: 999px;
      padding: 3px 12px; margin: 2px; text-decoration: none;">🤗 model card</a>
  </p>
</div>
"""

CSS = """
.gradio-container { max-width: 1120px !important; margin: 0 auto !important; }
#gen-btn { font-size: 1.1em; font-weight: 700; }
#out-img img { border-radius: 14px; }
#latency textarea { text-align: center; font-weight: 600; }
footer { visibility: hidden; }
.pfooter { text-align: center; opacity: 0.65; font-size: 0.9em; margin-top: 6px; }
"""

theme = gr.themes.Soft(primary_hue="violet", secondary_hue="pink", neutral_hue="slate")

with gr.Blocks(title="FLUX.2-klein 1-step RDM (CPU)", theme=theme, css=CSS) as demo:
    gr.HTML(HEADER)
    with gr.Row(equal_height=False):
        with gr.Column(scale=5):
            prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT, lines=3,
                                placeholder="describe the image…")
            btn = gr.Button("✨ Generate (1 step)", variant="primary", size="lg", elem_id="gen-btn")
            with gr.Accordion("Seed", open=False):
                with gr.Row():
                    seed = gr.Number(label="Seed", value=20240623, precision=0)
                    randomize = gr.Checkbox(label="Random seed", value=True)
            gr.Markdown("**🎨 Aesthetic presets** — click to fill the prompt:")
            gr.Examples(examples=[[p["prompt"]] for p in PRESETS], inputs=[prompt],
                        label=None, examples_per_page=8)
        with gr.Column(scale=4):
            out = gr.Image(label=f"{RES} × {RES} · one step", type="pil", elem_id="out-img", height=520)
            with gr.Row():
                latency = gr.Textbox(label="Latency", interactive=False, elem_id="latency", scale=2)
                used_seed = gr.Number(label="Seed used", precision=0, interactive=False, scale=1)
    gr.HTML('<div class="pfooter">FLUX.2-klein-4B distilled 4→1 steps · Qwen3-0.6B for CPU · '
            'First generation loads models (~2-3 min) · EPFL VITA lab</div>')
    
    btn.click(generate, inputs=[prompt, seed, randomize], outputs=[out, used_seed, latency])
    prompt.submit(generate, inputs=[prompt, seed, randomize], outputs=[out, used_seed, latency])

demo.queue(max_size=40).launch()