Update app.py

app.py

@@ -1,46 +1,201 @@
 """
 FLUX.2-klein-4B GGUF — CPU-only Inference Space
-硬件约束：无 GPU · CPU · 18 GB RAM
+硬件：无 GPU · CPU · 18 GB RAM
 """
 
 import os
-import time
+import sys
+import json
 import torch
-import gradio as gr
-from diffusers import FluxPipeline
+import numpy as np
+from huggingface_hub import login, hf_hub_download, list_repo_files
 
 # ── 环境配置 ──────────────────────────────────────────────────
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
-os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 torch.set_num_threads(os.cpu_count() or 4)
 
-MODEL_ID = "unsloth/FLUX-2-klein-4B-GGUF"
+GGUF_REPO  = "unsloth/FLUX-2-klein-4B-GGUF"
+TOKEN      = os.environ.get("HF_TOKEN", "")
 
-# ── 加载模型（启动时执行一次） ──────────────────────────────────
+# ── 认证 ──────────────────────────────────────────────────────
 print("╔═══════════════════════════════════════════════════╗")
 print("║   FLUX.2-klein-4B GGUF — CPU Inference          ║")
 print("╚═══════════════════════════════════════════════════╝")
-print(f"  Model   : {MODEL_ID}")
-print(f"  Threads : {torch.get_num_threads()}")
-print(f"  PyTorch : {torch.__version__}")
-print("  Loading...")
 
-pipe = FluxPipeline.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.float32,   # CPU 上必须用 float32
-)
-pipe = pipe.to("cpu")
+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")
 
-# 如果 diffusers 支持 GGUF 原生加载失败，取消下方注释使用备用方案
-# 备用方案：从原始模型加载 pipeline 骨架，再替换 transformer 为 GGUF 权重
-#
-# from diffusers import FluxPipeline
-# BASE_ID = "black-forest-labs/FLUX.2-klein-4B"  # 原始非量化模型
-# pipe = FluxPipeline.from_pretrained(BASE_ID, torch_dtype=torch.float32)
-# # 然后用 gguf 包手动加载量化权重替换 pipe.transformer
-
-print("  ✓ Model loaded successfully!\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()
@@ -54,17 +209,13 @@ def generate(
     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,
@@ -74,251 +225,80 @@ def generate(
         guidance_scale=guidance,
         generator=generator,
     )
-
     dt = time.time() - t0
-    img = result.images[0]
 
-    return img, seed, f"完成 · 耗时 {dt:.1f}s · seed {seed}"
+    return result.images[0], seed, f"完成 · 耗时 {dt:.1f}s · seed {seed}"
 
 
-# ── 样式 ──────────────────────────────────────────────────────
+# ── 界面样式 ──────────────────────────────────────────────────
 CSS = """
-:root {
-    --bg: #0a0a0f;
-    --card: #13131a;
-    --card-hover: #1a1a24;
-    --border: #23233a;
-    --fg: #e4e4ec;
-    --muted: #6b6b80;
-    --accent: #7c5cfc;
-    --accent-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, .gr-text-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(--accent-glow) !important;
-    transition: all .2s ease !important;
-}
-.gr-button-primary:hover {
-    transform: translateY(-1px) !important;
-    box-shadow: 0 6px 32px var(--accent-glow) !important;
-    filter: brightness(1.1) !important;
-}
-.gr-button-primary:active {
-    transform: translateY(0) !important;
-}
-
-/* 滑块 */
-.gr-slider input[type=range] { accent-color: var(--accent) !important; }
-
-/* 图片输出 */
-.gr-image { border-radius: 12px !important; overflow: hidden !important; }
-.gr-image img { border-radius: 12px !important; }
-
-/* Examples */
-.gr-examples { border-radius: 12px !important; }
-
-/* 隐藏 footer */
-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;
-}
+: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}
 """
 
-# ── 构建界面 ──────────────────────────────────────────────────
-with gr.Blocks(
-    title="FLUX.2-klein-4B GGUF",
-    css=CSS,
-    theme=gr.themes.Soft(primary_hue="violet"),
-) as demo:
+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("""
-    <div style="text-align:center; padding:2rem 0 .6rem">
-        <h1 style="font-size:2.4rem; font-weight:800; letter-spacing:-.03em;
+    <div style="text-align:center;padding:2rem 0 .6rem">
+        <h1 style="font-size:2.4rem;font-weight:800;letter-spacing:-.03em;
                     background:linear-gradient(135deg,#7c5cfc 0%,#c084fc 50%,#f0abfc 100%);
-                    -webkit-background-clip:text; -webkit-text-fill-color:transparent;
-                    margin:0; line-height:1.2">
-            FLUX.2-klein-4B
-        </h1>
-        <p style="color:#555; font-size:.88rem; margin-top:.5rem; font-weight:500;
-                  letter-spacing:.08em; text-transform:uppercase">
-            GGUF Quantized · CPU Inference · 4B Parameters
-        </p>
-    </div>
-    """)
-
-    # ── 警告 ──
+                    -webkit-background-clip:text;-webkit-text-fill-color:transparent;
+                    margin:0;line-height:1.2">FLUX.2-klein-4B</h1>
+        <p style="color:#555;font-size:.88rem;margin-top:.5rem;font-weight:500;
+                  letter-spacing:.08em;text-transform:uppercase">
+            GGUF Quantized · CPU Inference · 4B Parameters</p>
+    </div>""")
+
     gr.HTML("""
     <div class="warn-bar">
-        &#9888;&nbsp; 纯 CPU 推理，速度较慢。建议 <strong>512×512</strong> 以下分辨率 +
-        <strong>8~12 步</strong>，首次生成会因 JIT 编译更慢，后续会加速。
-        如果超时，请进一步降低分辨率或步数。
-    </div>
-    """)
+        &#9888;&nbsp; 纯 CPU 推理，速度较慢。建议 <strong>512×512</strong> 以下 +
+        <strong>8~12 步</strong>。首次生成因 JIT 编译更慢，后续加速。
+        超时请降低分辨率或步数。
+    </div>""")
 
     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, cinematic lighting...",
-                lines=3,
-                show_label=True,
-            )
-
-            negative_prompt = gr.Textbox(
-                label="Negative Prompt",
-                placeholder="blurry, low quality, distorted, watermark...",
-                lines=1,
-                show_label=True,
-            )
-
+            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 falling over a stone bridge, koi pond below, golden hour lighting, 4K",
-                        ""
-                    ],
-                    [
-                        "Futuristic space station interior, holographic star maps floating in zero gravity, bioluminescent plants, cinematic wide shot",
-                        ""
-                    ],
-                    [
-                        "Oil painting of a cozy wooden cabin in snowy mountains, warm orange light from windows, smoke from chimney, impressionist style",
-                        ""
-                    ],
-                    [
-                        "Macro photograph of a dewdrop on a spider web at sunrise, prismatic light refraction, extreme shallow depth of field",
-                        ""
-                    ],
-                    [
-                        "A whimsical treehouse library at twilight, floating books with glowing pages, warm lanterns, fairy tale atmosphere",
-                        "blurry, low quality"
-                    ],
-                    [
-                        "Underwater coral reef city with bioluminescent architecture, schools of tropical fish, volumetric god rays, wide angle",
-                        ""
-                    ],
-                ],
-                inputs=[prompt, negative_prompt],
-                label="Quick Examples",
-            )
-
-        # ── 右栏：输出 ──
-        with gr.Column(scale=4):
-            output_img = gr.Image(
-                label="Output",
-                show_label=False,
-                height=540,
-                elem_classes="output-image",
-            )
-            info_text = gr.Textbox(
-                label="Status",
-                interactive=False,
-                show_label=True,
-                max_lines=1,
-                elem_classes="info-bar",
-                value="Ready — enter a prompt and click Generate",
-            )
-
-    # ── 绑定事件 ──
-    btn.click(
-        fn=generate,
-        inputs=[prompt, negative_prompt, w, h, steps, guidance, seed],
-        outputs=[output_img, seed, info_text],
-    )
+            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")
 
-    # 回车也能触发
-    prompt.submit(
-        fn=generate,
-        inputs=[prompt, negative_prompt, w, h, steps, guidance, seed],
-        outputs=[output_img, seed, info_text],
-    )
+        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])
 
-# ── 启动 ──────────────────────────────────────────────────────
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch(server_name="0.0.0.0", server_port=7860)