Update app.py

app.py

@@ -8,18 +8,14 @@ from transformers import AutoImageProcessor, AutoModelForImageClassification
 from diffusers import AutoPipelineForImage2Image
 
 # -------------------------
-# 1) 年龄估计模型（HF 上可用）
-#    说明：本示例使用 Hugging Face 的 ViT 年龄估计模型。
-#    我们把分类标签转换成年龄（若是"0-2"取区间中点；若是"23"就取整）。
+# 1) 年龄估计模型
 # -------------------------
 AGE_MODEL_ID = "nateraw/vit-age-classifier"
-
 age_processor = AutoImageProcessor.from_pretrained(AGE_MODEL_ID)
 age_model = AutoModelForImageClassification.from_pretrained(AGE_MODEL_ID)
 age_model.eval()
 
 def _label_to_age(label: str) -> float:
-    # 尝试解析类似 "(0-2)"、"0-2"、"3-9" 的标签
     label = label.strip().replace("(", "").replace(")", "")
     if "-" in label:
         a, b = label.split("-")
@@ -27,11 +23,9 @@ def _label_to_age(label: str) -> float:
             return (float(a) + float(b)) / 2.0
         except:
             pass
-    # 若是单值，如 "23"
     try:
         return float(label)
     except:
-        # 兜底：无法解析就返回 NaN
         return float("nan")
 
 @torch.inference_mode()
@@ -39,19 +33,16 @@ def estimate_age(image: Image.Image) -> dict:
     inputs = age_processor(images=image, return_tensors="pt")
     logits = age_model(**inputs).logits
     probs = torch.softmax(logits, dim=-1)[0]
-
-    # 取 top-5 以便展示
     id2label = age_model.config.id2label
+
     topk = torch.topk(probs, k=min(5, probs.shape[0]))
-    items = []
-    ages = []
+    items, ages = [], []
     for score, idx in zip(topk.values.tolist(), topk.indices.tolist()):
         label = id2label[idx]
         age = _label_to_age(label)
         ages.append((age, score))
         items.append(f"{label}: {score*100:.1f}%")
 
-    # 期望年龄（加权平均）
     ages_valid = [(a, p) for a, p in ages if not math.isnan(a)]
     if ages_valid:
         num = sum(a * p for a, p in ages_valid)
@@ -67,19 +58,15 @@ def estimate_age(image: Image.Image) -> dict:
 
 # -------------------------
 # 2) 漫画风格生成（img2img）
-#    说明：使用 "stabilityai/sd-turbo" 的图生图，速度较快，提示词主打漫画/卡通风。
-#    CPU 也能跑，但较慢；有 GPU（T4/A10）体验最佳。
 # -------------------------
 IMG2IMG_MODEL_ID = "stabilityai/sd-turbo"
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if device == "cuda" else torch.float32
 
 pipe = AutoPipelineForImage2Image.from_pretrained(
     IMG2IMG_MODEL_ID,
     torch_dtype=dtype
-)
-pipe = pipe.to(device)
+).to(device)
 
 DEFAULT_PROMPT = (
     "comic style, manga, cel-shaded, bold ink outlines, clean lineart, high contrast, "
@@ -96,32 +83,34 @@ def stylize_to_comic(
     steps: int = 4,
     seed: int | None = 42
 ) -> Image.Image:
-    if seed is None or seed < 0:
-        generator = None
-    else:
-        generator = torch.Generator(device=device).manual_seed(seed)
-
+    generator = None if (seed is None or seed < 0) else torch.Generator(device=device).manual_seed(int(seed))
     image = image.convert("RGB")
     out = pipe(
         prompt=prompt,
         negative_prompt=NEG_PROMPT,
         image=image,
-        strength=strength,           # 0.2~0.7：数值越大改动越明显
-        num_inference_steps=steps,   # sd-turbo 推荐极少步数（2~6）
-        guidance_scale=guidance_scale,  # sd-turbo 常用 0~1
-        generator=generator
+        strength=float(strength),
+        num_inference_steps=int(steps),
+        guidance_scale=float(guidance_scale),
+        generator=generator,
     )
     return out.images[0]
 
 # -------------------------
-# 3) Gradio 界面
+# 3) Gradio 界面（两个按钮都在最上面）
 # -------------------------
-def process(image, prompt, strength, guidance, steps, seed):
+def ui_estimate_age(image):
     if image is None:
-        return "请先上传图片", None
+        return "请先上传图片"
+    res = estimate_age(image)
+    if res["expected_age"] is None:
+        return "年龄估计：解析失败（可能检测不到年龄标签）"
+    return f"年龄估计：≈ {res['expected_age']} 岁\nTop-5: " + " | ".join(res["top5"])
 
-    age_result = estimate_age(image)
-    styled = stylize_to_comic(
+def ui_stylize(image, prompt, strength, guidance, steps, seed):
+    if image is None:
+        return None
+    return stylize_to_comic(
         image=image,
         prompt=prompt,
         strength=strength,
@@ -130,19 +119,16 @@ def process(image, prompt, strength, guidance, steps, seed):
         seed=int(seed) if seed is not None else 42
     )
 
-    # 结果文字
-    if age_result["expected_age"] is None:
-        age_text = "年龄估计：解析失败（可能检测不到年龄标签）"
-    else:
-        age_text = f"年龄估计：≈ {age_result['expected_age']} 岁\nTop-5: " + " | ".join(age_result["top5"])
+with gr.Blocks(title="Age & Comicify Agent") as demo:
+    gr.Markdown("# 🧠 Age & Comicify Agent\n上传图片 → ① 估计年龄 ② 生成���画风格图片")
 
-    return age_text, styled
+    # 顶部两个按钮
+    with gr.Row():
+        btn_est = gr.Button("🧮 估计年龄", variant="primary")
+        btn_gen = gr.Button("🎨 生成漫画图片", variant="secondary")
 
-with gr.Blocks(title="Age & Comicify Agent") as demo:
-    gr.Markdown("# 🧠 Age & Comicify Agent\n上传图片 → 年龄估计 → 漫画风格生成")
     with gr.Row():
         with gr.Column(scale=1):
-            run_btn = gr.Button("🚀 运行")
             in_img = gr.Image(label="上传图片", type="pil")
             prompt = gr.Textbox(label="风格提示词", value=DEFAULT_PROMPT)
             strength = gr.Slider(0.1, 0.9, value=0.6, step=0.05, label="风格强度（strength）")
@@ -153,11 +139,10 @@ with gr.Blocks(title="Age & Comicify Agent") as demo:
             age_txt = gr.Textbox(label="年龄估计结果")
             out_img = gr.Image(label="漫画风格输出")
 
-    run_btn.click(
-        fn=process,
-        inputs=[in_img, prompt, strength, guidance, steps, seed],
-        outputs=[age_txt, out_img]
-    )
+    # 绑定：按钮各自只触发一个功能
+    btn_est.click(fn=ui_estimate_age, inputs=[in_img], outputs=[age_txt])
+    btn_gen.click(fn=ui_stylize, inputs=[in_img, prompt, strength, guidance, steps, seed], outputs=[out_img])
 
 if __name__ == "__main__":
-    demo.launch()
+    # 可选：并发/队列
+    demo.queue().launch()