Create app.py

app.py

@@ -0,0 +1,163 @@
+import os
+import math
+import torch
+import gradio as gr
+from PIL import Image
+
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+from diffusers import AutoPipelineForImage2Image
+
+# -------------------------
+# 1) 年龄估计模型（HF 上可用）
+#    说明：本示例使用 Hugging Face 的 ViT 年龄估计模型。
+#    我们把分类标签转换成年龄（若是"0-2"取区间中点；若是"23"就取整）。
+# -------------------------
+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("-")
+        try:
+            return (float(a) + float(b)) / 2.0
+        except:
+            pass
+    # 若是单值，如 "23"
+    try:
+        return float(label)
+    except:
+        # 兜底：无法解析就返回 NaN
+        return float("nan")
+
+@torch.inference_mode()
+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 = []
+    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)
+        den = sum(p for _, p in ages_valid)
+        expected_age = num / den
+    else:
+        expected_age = float("nan")
+
+    return {
+        "expected_age": None if math.isnan(expected_age) else round(expected_age, 1),
+        "top5": items
+    }
+
+# -------------------------
+# 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)
+
+DEFAULT_PROMPT = (
+    "comic style, manga, cel-shaded, bold ink outlines, clean lineart, high contrast, "
+    "professional illustration, vibrant"
+)
+NEG_PROMPT = "realistic, photorealistic, blurry, noisy, artifacts, watermark, text"
+
+@torch.inference_mode()
+def stylize_to_comic(
+    image: Image.Image,
+    prompt: str = DEFAULT_PROMPT,
+    strength: float = 0.6,
+    guidance_scale: float = 0.0,
+    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)
+
+    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
+    )
+    return out.images[0]
+
+# -------------------------
+# 3) Gradio 界面
+# -------------------------
+def process(image, prompt, strength, guidance, steps, seed):
+    if image is None:
+        return "请先上传图片", None
+
+    age_result = estimate_age(image)
+    styled = stylize_to_comic(
+        image=image,
+        prompt=prompt,
+        strength=strength,
+        guidance_scale=guidance,
+        steps=int(steps),
+        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"])
+
+    return age_text, styled
+
+with gr.Blocks(title="Age & Comicify Agent") as demo:
+    gr.Markdown("# 🧠 Age & Comicify Agent\n上传图片 → 年龄估计 → 漫画风格生成")
+    with gr.Row():
+        with gr.Column(scale=1):
+            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）")
+            guidance = gr.Slider(0.0, 3.0, value=0.0, step=0.1, label="引导系数（guidance_scale）")
+            steps = gr.Slider(2, 12, value=4, step=1, label="步数（num_inference_steps）")
+            seed = gr.Number(value=42, precision=0, label="随机种子（固定可复现）")
+            run_btn = gr.Button("🚀 运行")
+        with gr.Column(scale=1):
+            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]
+    )
+
+if __name__ == "__main__":
+    demo.launch()