"""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 = """

⚡ FLUX.2-klein · 1-step · RDM

One diffusion step, distilled with Representation Distribution Matching at EPFL VITA — Running on CPU (Qwen3-0.6B).

GenEval 0.826 PickScore-COCO 22.76 ⚠️ 125px Borderline 🤗 model card

""" 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('
FLUX.2-klein-4B distilled 4→1 steps · Qwen3-0.6B for CPU · ' 'First generation loads models (~2-3 min) · EPFL VITA lab
') 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()