""" FLUX.2-klein-4B GGUF — CPU-only Inference Space 硬件：无 GPU · CPU · 18 GB RAM """ import os import sys import json import torch import numpy as np from huggingface_hub import login, hf_hub_download, list_repo_files # ── 环境配置 ────────────────────────────────────────────────── os.environ["TOKENIZERS_PARALLELISM"] = "false" torch.set_num_threads(os.cpu_count() or 4) GGUF_REPO = "unsloth/FLUX-2-klein-4B-GGUF" TOKEN = os.environ.get("HF_TOKEN", "") # ── 认证 ────────────────────────────────────────────────────── print("╔═══════════════════════════════════════════════════╗") print("║ FLUX.2-klein-4B GGUF — CPU Inference ║") print("╚═══════════════════════════════════════════════════╝") if TOKEN: login(token=TOKEN, add_to_git_credential=False) print(" ✓ HF_TOKEN authenticated") else: print(" ✗ HF_TOKEN not set!") print(" → Go to Space Settings → Repository Secrets → Add HF_TOKEN") print(" → Get token: https://huggingface.co/settings/tokens") print(" → Also accept model license on the model page first!\n") print(f" Model : {GGUF_REPO}") print(f" Threads : {torch.get_num_threads()}") print(f" PyTorch : {torch.__version__}") # ── 加载模型（双重策略） ───────────────────────────────────── from diffusers import FluxPipeline, FluxTransformer2DModel pipe = None # ═══════════════════════════════════════════════════════════ # Strategy A：直接加载（diffusers 原生 GGUF 支持） # ═══════════════════════════════════════════════════════════ if pipe is None: try: print("\n[A] Trying direct diffusers loading...") pipe = FluxPipeline.from_pretrained( GGUF_REPO, torch_dtype=torch.float32, token=TOKEN or None, ) pipe.to("cpu") print("[A] ✓ Direct loading succeeded!\n") except Exception as e: print(f"[A] ✗ Failed: {e}\n") # ═══════════════════════════════════════════════════════════ # Strategy B：手动解析 GGUF → 注入 transformer # ═══════════════════════════════════════════════════════════ if pipe is None: try: print("[B] Manual GGUF loading...") from gguf import GGUFReader # 1) 列出仓库文件，找到 .gguf print(" Listing repo files...") repo_files = list_repo_files(GGUF_REPO, token=TOKEN or None) gguf_candidates = [f for f in repo_files if f.endswith(".gguf")] if not gguf_candidates: raise FileNotFoundError("No .gguf file found in repo") # 优先选 Q4_K_M（精度与速度平衡），没有则选第一个 preferred = [f for f in gguf_candidates if "Q4_K_M" in f.upper()] gguf_fname = preferred[0] if preferred else gguf_candidates[0] print(f" Selected: {gguf_fname}") # 2) 下载 GGUF 文件 print(" Downloading GGUF file...") gguf_path = hf_hub_download( GGUF_REPO, gguf_fname, token=TOKEN or None ) print(f" Saved to: {gguf_path}") # 3) 读取 GGUF 张量 print(" Reading GGUF tensors...") reader = GGUFReader(gguf_path) gguf_sd = {} for tensor in reader.tensors: name = tensor.name if isinstance(name, bytes): name = name.decode("utf-8") data = np.array(tensor.data, dtype=np.float32) data = data.reshape(tensor.shape).copy() gguf_sd[name] = torch.from_numpy(data) print(f" Loaded {len(gguf_sd)} tensors") sample_keys = list(gguf_sd.keys())[:5] print(f" Sample keys: {sample_keys}") # 4) 下载 transformer config print(" Loading transformer config...") config_path = hf_hub_download( GGUF_REPO, "transformer/config.json", token=TOKEN or None ) with open(config_path, "r") as f: tf_config = json.load(f) # 5) 创建空 transformer 并加载 GGUF 权重 print(" Building transformer and loading weights...") transformer = FluxTransformer2DModel.from_config(tf_config) # 尝试多种 key 映射 loaded_ok = False # B1: 直接映射 try: info = transformer.load_state_dict(gguf_sd, strict=False) missing = [k for k in info.missing_keys if "bias" not in k.lower()] if len(missing) < 5: loaded_ok = True print(f" B1 ✓ Direct mapping (missing={len(missing)})") except Exception: pass # B2: 加 "transformer." 前缀 if not loaded_ok: try: mapped = {"transformer." + k: v for k, v in gguf_sd.items()} info = transformer.load_state_dict(mapped, strict=False) missing = [k for k in info.missing_keys if "bias" not in k.lower()] if len(missing) < 5: loaded_ok = True print(f" B2 ✓ Prefix mapping (missing={len(missing)})") except Exception: pass # B3: 去掉已有前缀再映射 if not loaded_ok: try: mapped = {} for k, v in gguf_sd.items(): new_key = k for prefix in ["transformer.", "model."]: if k.startswith(prefix): new_key = k[len(prefix):] break mapped[new_key] = v info = transformer.load_state_dict(mapped, strict=False) missing = [k for k in info.missing_keys if "bias" not in k.lower()] if len(missing) < 5: loaded_ok = True print(f" B3 ✓ Strip-prefix mapping (missing={len(missing)})") except Exception: pass if not loaded_ok: # 打印诊断信息帮助调试 model_keys = set(transformer.state_dict().keys()) gguf_keys = set(gguf_sd.keys()) print(f"\n ── Key mismatch diagnostics ──") print(f" Model expects {len(model_keys)} keys") print(f" GGUF has {len(gguf_keys)} keys") print(f" Model key sample : {list(model_keys)[:3]}") print(f" GGUF key sample : {list(gguf_keys)[:3]}") raise RuntimeError("Could not map GGUF keys to model keys") # 6) 构建完整 pipeline（传入 transformer 跳过磁盘加载） print(" Building full pipeline...") pipe = FluxPipeline.from_pretrained( GGUF_REPO, transformer=transformer, torch_dtype=torch.float32, token=TOKEN or None, ) pipe.to("cpu") print("[B] ✓ Manual GGUF loading succeeded!\n") except Exception as e: print(f"[B] ✗ Failed: {e}") import traceback traceback.print_exc() print() # ── 最终检查 ────────────────────────────────────────────────── if pipe is None: print("=" * 58) print("FATAL: Could not load model.") print("") print("Please check:") print(" 1. HF_TOKEN is set as a Space secret") print(" 2. You have accepted the model license at:") print(f" https://huggingface.co/{GGUF_REPO}") print("=" * 58) sys.exit(1) print("✓ Model ready for inference!\n") # ── 推理函数 ────────────────────────────────────────────────── @torch.no_grad() def generate( prompt: str, negative_prompt: str, width: int, height: int, steps: int, guidance: float, seed: int, progress=gr.Progress(track_tqdm=True), ): if not prompt.strip(): raise gr.Error("请输入提示词") seed = seed if seed >= 0 else int(time.time() * 1000) % (2**32) generator = torch.Generator("cpu").manual_seed(seed) t0 = time.time() result = pipe( prompt=prompt, negative_prompt=negative_prompt if negative_prompt.strip() else None, width=width, height=height, num_inference_steps=steps, guidance_scale=guidance, generator=generator, ) dt = time.time() - t0 return result.images[0], seed, f"完成 · 耗时 {dt:.1f}s · seed {seed}" # ── 界面样式 ────────────────────────────────────────────────── CSS = """ :root{--bg:#0a0a0f;--card:#13131a;--border:#23233a;--fg:#e4e4ec;--muted:#6b6b80;--accent:#7c5cfc;--glow:rgba(124,92,252,.25)} .gradio-container{background:var(--bg)!important;color:var(--fg)!important;max-width:1080px!important;padding-top:1rem!important} .gr-block,.gr-panel{background:var(--card)!important;border:1px solid var(--border)!important;border-radius:14px!important} .gr-input,textarea.gr-textbox,input[type=number].gr-number{background:#0e0e16!important;border:1px solid var(--border)!important;color:var(--fg)!important;border-radius:10px!important;font-size:.92rem!important} textarea.gr-textbox{resize:vertical!important} .gr-label{color:var(--muted)!important;font-size:.78rem!important;text-transform:uppercase!important;letter-spacing:.06em!important;font-weight:600!important} .gr-button-primary{background:linear-gradient(135deg,#7c5cfc,#a78bfa)!important;border:none!important;font-weight:700!important;font-size:1rem!important;border-radius:10px!important;padding:.7rem 1.5rem!important;box-shadow:0 4px 24px var(--glow)!important;transition:all .2s ease!important} .gr-button-primary:hover{transform:translateY(-1px)!important;box-shadow:0 6px 32px var(--glow)!important;filter:brightness(1.1)!important} .gr-slider input[type=range]{accent-color:var(--accent)!important} .gr-image{border-radius:12px!important;overflow:hidden!important} footer{display:none!important} ::-webkit-scrollbar{width:6px}::-webkit-scrollbar-track{background:transparent}::-webkit-scrollbar-thumb{background:#333;border-radius:3px} .warn-bar{background:linear-gradient(90deg,#1c1a2e,#1a1428);border:1px solid #2e2848;border-radius:10px;padding:.7rem 1rem;margin-bottom:1.2rem;font-size:.82rem;color:#a0a0b8;line-height:1.5} .warn-bar strong{color:#c4b5fd} .info-bar{background:#0e0e16;border:1px solid var(--border);border-radius:8px;padding:.5rem .8rem;font-size:.82rem;color:var(--muted);font-family:'SF Mono','Fira Code',monospace} """ import time import gradio as gr with gr.Blocks(title="FLUX.2-klein-4B GGUF", css=CSS, theme=gr.themes.Soft(primary_hue="violet")) as demo: gr.HTML("""

FLUX.2-klein-4B

GGUF Quantized · CPU Inference · 4B Parameters

""") gr.HTML("""
⚠  纯 CPU 推理，速度较慢。建议 512×512 以下 + 8~12 步。首次生成因 JIT 编译更慢，后续加速。 超时请降低分辨率或步数。
""") with gr.Row(equal_height=False): with gr.Column(scale=3): prompt = gr.Textbox(label="Prompt", placeholder="A cyberpunk city at night with neon signs reflecting in rain-soaked streets...", lines=3) negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="blurry, low quality, distorted...", lines=1) with gr.Row(): w = gr.Slider(256, 768, value=512, step=64, label="Width") h = gr.Slider(256, 768, value=512, step=64, label="Height") with gr.Row(): steps = gr.Slider(4, 30, value=10, step=1, label="Steps") guidance = gr.Slider(1.0, 10.0, value=3.5, step=0.5, label="Guidance Scale") seed = gr.Number(label="Seed（-1 随机）", value=-1, precision=0) btn = gr.Button("✦ Generate", variant="primary", size="lg") gr.Examples(examples=[ ["A serene Japanese garden with cherry blossoms over a stone bridge, koi pond, golden hour, 4K", ""], ["Futuristic space station interior, holographic star maps, bioluminescent plants, cinematic wide shot", ""], ["Oil painting of a cozy wooden cabin in snowy mountains, warm light from windows, impressionist style", ""], ["Macro photograph of a dewdrop on a spider web at sunrise, prismatic light, shallow depth of field", ""], ["Whimsical treehouse library at twilight, floating books with glowing pages, fairy tale atmosphere", "blurry"], ["Underwater coral reef city, bioluminescent architecture, tropical fish, volumetric god rays", ""], ], inputs=[prompt, negative_prompt], label="Quick Examples") with gr.Column(scale=4): output_img = gr.Image(label="Output", show_label=False, height=540) info_text = gr.Textbox(label="Status", interactive=False, max_lines=1, elem_classes="info-bar", value="Ready") btn.click(fn=generate, inputs=[prompt, negative_prompt, w, h, steps, guidance, seed], outputs=[output_img, seed, info_text]) prompt.submit(fn=generate, inputs=[prompt, negative_prompt, w, h, steps, guidance, seed], outputs=[output_img, seed, info_text]) demo.launch(server_name="0.0.0.0", server_port=7860)