Update app.py

app.py

@@ -1,94 +1,324 @@
-import os, sys, subprocess, time, json, shutil
-from pathlib import Path
-from datetime import datetime
+"""
+FLUX.2-klein-4B GGUF — CPU-only Inference Space
+硬件约束：无 GPU · CPU · 18 GB RAM
+"""
 
-BASE = Path.cwd()
-COMFYUI = BASE / "ComfyUI"
-OUTPUTS = BASE / "outputs"
-MODEL_FILE = "flux-2-klein-4b-Q4_K_M.gguf"
+import os
+import time
+import torch
+import gradio as gr
+from diffusers import FluxPipeline
+
+# ── 环境配置 ──────────────────────────────────────────────────
+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"
+
+# ── 加载模型（启动时执行一次） ──────────────────────────────────
+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")
+
+# 如果 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
 
-def log(msg):
-    print(f"[{datetime.now().strftime('%H:%M:%S')}] {msg}")
+print("  ✓ Model loaded successfully!\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),
+):
+    """文本生成图像"""
 
-def generate(prompt, neg="", w=1024, h=1024, steps=20, cfg=7.0, seed=-1):
-    import requests
-    
     if not prompt.strip():
-        return None, "Enter a prompt!"
-    
-    if seed < 0:
-        seed = int(time.time() * 1000) % (2**32)
-    
-    workflow = {
-        "1": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["4", 0]}},
-        "2": {"class_type": "CLIPTextEncode", "inputs": {"text": neg or "blurry", "clip": ["4", 0]}},
-        "3": {"class_type": "VAEDecode", "inputs": {"samples": ["6", 0], "vae": ["5", 0]}},
-        "4": {"class_type": "CLIPLoader", "inputs": {"clip_name": "t5xxl_fp16.safetensors", "type": "l"}},
-        "5": {"class_type": "VAELoader", "inputs": {"vae_name": "ae.safetensors"}},
-        "6": {"class_type": "KSampler", "inputs": {
-            "model": ["7", 0], "positive": ["1", 0], "negative": ["2", 0],
-            "latent_image": ["8", 0], "seed": seed, "steps": steps,
-            "cfg": cfg, "sampler_name": "euler", "scheduler": "normal", "denoise": 1.0
-        }},
-        "7": {"class_type": "GGUFLoader", "inputs": {"unet_name": MODEL_FILE, "weight_dtype": "fp16_e4m3fn"}},
-        "8": {"class_type": "EmptyLatentImage", "inputs": {"width": w, "height": h, "batch_size": 1}},
-        "9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "output", "images": ["3", 0]}}
-    }
-    
-    # Start ComfyUI if not running
-    global proc
-    if 'proc' not in globals():
-        log("Starting ComfyUI...")
-        proc = subprocess.Popen(
-            [sys.executable, str(COMFYUI / "main.py"), "--listen", "127.0.0.1", "--port", "8188"],
-            cwd=str(COMFYUI), stdout=subprocess.PIPE, stderr=subprocess.PIPE
-        )
-        for _ in range(30):
-            try:
-                requests.get("http://127.0.0.1:8188/system_stats", timeout=2)
-                break
-            except:
-                time.sleep(1)
-    
-    # Queue and wait
-    resp = requests.post("http://127.0.0.1:8188/prompt", json={"prompt": workflow}, timeout=10)
-    pid = resp.json()["prompt_id"]
-    
-    for _ in range(150):
-        time.sleep(2)
-        hist = requests.get(f"http://127.0.0.1:8188/history/{pid}", timeout=5).json()
-        if pid in hist:
-            imgs = hist[pid].get("outputs", {}).get("9", {}).get("images", [])
-            if imgs:
-                src = COMFYUI / "output" / imgs[0]["filename"]
-                OUTPUT.mkdir(exist_ok=True)
-                dst = OUTPUT / f"{datetime.now().strftime('%H%M%S')}.png"
-                shutil.copy2(src, dst)
-                return str(dst), f"✅ Done! ({dst.stat().st_size/(1024*1024):.1f}MB)"
-            break
-    
-    return None, "Timeout or error"
+        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
+    img = result.images[0]
+
+    return img, 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;
+}
+"""
+
+# ── 构建界面 ──────────────────────────────────────────────────
+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, 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,
+            )
+
+            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],
+    )
+
+    # 回车也能触发
+    prompt.submit(
+        fn=generate,
+        inputs=[prompt, negative_prompt, w, h, steps, guidance, seed],
+        outputs=[output_img, seed, info_text],
+    )
 
-import gradio as gr
 
-with gr.Blocks(title="FLUX Generator") as demo:
-    gr.Markdown("# 🎨 FLUX.2-klein-4B (Q4_K_M)\n### ComfyUI Backend")
-    
-    with gr.Row():
-        with gr.Column(scale=2):
-            prompt = gr.Textbox(label="Prompt", lines=3, value="a red fox in snow, detailed, 8k")
-            neg = gr.Textbox(label="Negative", lines=2)
-        with gr.Column():
-            w = gr.Slider(512, 1024, 1024, 64, label="Width")
-            h = gr.Slider(512, 1024, 1024, 64, label="Height")
-            steps = gr.Slider(1, 50, 20, 1, label="Steps")
-            cfg = gr.Slider(1, 20, 7.0, 0.5, label="CFG")
-            seed = gr.Number(-1, label="Seed")
-    
-    gen = gr.Button("Generate", variant="primary")
-    img = gr.Image(label="Output")
-    status = gr.Textbox(label="Status")
-    
-    gen.click(generate, [prompt, neg, w, h, steps, cfg, seed], [img, status])
-
-demo.launch(server_name="0.0.0.0", port=7860)
+# ── 启动 ───────────────────────────────────────────────────��──
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)