Blame for app.py

"""
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("""
    <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>""")

    gr.HTML("""
    <div class="warn-bar">
        &#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...", 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)