Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py — Age-first + FAST group cartoons (SD-Turbo), single page
+# app.py — Age-first + FAST group cartoons (SD-Turbo), single page (HF Spaces safe)
 
 import os
 os.environ["TRANSFORMERS_NO_TF"] = "1"
@@ -6,6 +6,8 @@ os.environ["TRANSFORMERS_NO_FLAX"] = "1"
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
 
 import math
+from typing import Optional
+
 import gradio as gr
 from PIL import Image, ImageDraw
 import numpy as np
@@ -21,7 +23,7 @@ AGE_RANGE_TO_MID = {
 }
 
 class PretrainedAgeEstimator:
-    def __init__(self, model_id: str = HF_MODEL_ID, device: str | None = None):
+    def __init__(self, model_id: str = HF_MODEL_ID, device: Optional[str] = None):
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         self.processor = AutoImageProcessor.from_pretrained(model_id, use_fast=True)
         self.model = AutoModelForImageClassification.from_pretrained(model_id)
@@ -52,7 +54,7 @@ class FaceCropper:
     - detect_all_wide: returns (list[crops], annotated, list[boxes])
     Boxes are (x1,y1,x2,y2) floats.
     """
-    def __init__(self, device: str | None = None, margin_scale: float = 1.8):
+    def __init__(self, device: Optional[str] = None, margin_scale: float = 1.8):
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         self.mtcnn = MTCNN(keep_all=True, device=self.device)
         self.margin_scale = margin_scale
@@ -62,13 +64,13 @@ class FaceCropper:
             return img.convert("RGB")
         return Image.fromarray(img).convert("RGB")
 
-    def _expand_box(self, box, W, H, aspect=0.8):  # 4:5 portrait (w/h=0.8)
+    def _expand_box(self, box, W, H, aspect=0.8):  # ~4:5 portrait (w/h=0.8)
         x1, y1, x2, y2 = box
         cx, cy = (x1 + x2)/2, (y1 + y2)/2
         w, h = (x2 - x1), (y2 - y1)
         side = max(w, h) * self.margin_scale
         tw = side
-        th = side / aspect  # make it taller than wide
+        th = side / aspect  # taller than wide
         nx1 = int(max(0, cx - tw/2)); nx2 = int(min(W, cx + tw/2))
         ny1 = int(max(0, cy - th/2)); ny2 = int(min(H, cy + th/2))
         return nx1, ny1, nx2, ny2
@@ -108,6 +110,7 @@ class FaceCropper:
         if boxes is None or len(boxes) == 0:
             return crops, annotated, []
 
+        # sort roughly left->right for table order
         for b, p in sorted(zip(boxes, probs), key=lambda x: (x[0][0]+x[0][2])/2):
             bx1, by1, bx2, by2 = map(float, b)
             draw.rectangle([bx1, by1, bx2, by2], outline=(0, 200, 255), width=3)
@@ -121,15 +124,24 @@ class FaceCropper:
 
 # ------------------ FAST Cartoonizer (SD-Turbo) ------------------
 from diffusers import AutoPipelineForImage2Image
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from transformers import AutoFeatureExtractor
+
 TURBO_ID = "stabilityai/sd-turbo"
 
 def load_turbo_pipe(device):
     dtype = torch.float16 if (device == "cuda") else torch.float32
     pipe = AutoPipelineForImage2Image.from_pretrained(
         TURBO_ID,
-        torch_dtype=dtype,
-        safety_checker=None,
+        dtype=dtype,           # ✅ no deprecation warning
     ).to(device)
+    # safety checker ON for public Spaces
+    pipe.safety_checker = StableDiffusionSafetyChecker.from_pretrained(
+        "CompVis/stable-diffusion-safety-checker"
+    )
+    pipe.feature_extractor = AutoFeatureExtractor.from_pretrained(
+        "CompVis/stable-diffusion-safety-checker"
+    )
     try:
         pipe.enable_attention_slicing()
     except Exception:
@@ -186,7 +198,7 @@ def predict_age(img, group_mode=False, auto_crop=True):
             top1, p1 = top[0]
             rows.append(f"| {i} | {age:.1f} | {top1} | {p1:.2f} |")
         md = "\n".join(rows)
-        # also return a simple dict from the largest (first) face just to feed Label
+        # also return a simple dict from the first face just to feed Label
         age0, top0 = age_est.predict(crops[0], topk=5)
         probs0 = {lbl: float(p) for lbl, p in top0}
         return probs0, md, annotated
@@ -222,7 +234,69 @@ def cartoonize(img, prompt="", group_mode=False, auto_crop=True, strength=0.5, s
         if not crops:
             crops = [pil]  # fallback
 
-        # resize each to 384 for speed/variety
         proc = []
         for c in crops:
-            c = _resiz_
+            c = _resize_512(c)
+            out = sd_pipe(
+                prompt=pos, negative_prompt=neg, image=c,
+                strength=float(strength), guidance_scale=0.0,
+                num_inference_steps=int(steps), generator=generator
+            )
+            proc.append(out.images[0])
+
+        # tile into a grid
+        n = len(proc)
+        cols = int(math.ceil(math.sqrt(n)))
+        rows = int(math.ceil(n / cols))
+        cell_w = max(im.width for im in proc)
+        cell_h = max(im.height for im in proc)
+        grid = Image.new("RGB", (cols * cell_w, rows * cell_h), (240, 240, 240))
+        for i, im in enumerate(proc):
+            r, c = divmod(i, cols)
+            grid.paste(im, (c * cell_w, r * cell_h))
+        return grid
+
+    # single person
+    face_wide = None
+    if auto_crop:
+        face_wide, _ = cropper.detect_one_wide(pil)
+    base = face_wide if face_wide is not None else pil
+    base = _resize_512(base)
+    out = sd_pipe(
+        prompt=pos, negative_prompt=neg, image=base,
+        strength=float(strength), guidance_scale=0.0,
+        num_inference_steps=int(steps), generator=generator
+    )
+    return out.images[0]
+
+# ------------------ UI ------------------
+with gr.Blocks(title="Group Age + Cartoons (Fast)") as demo:
+    gr.Markdown("# Predict ages and make fast cartoons — single or group photos")
+    with gr.Row():
+        with gr.Column(scale=1):
+            img_in = gr.Image(sources=["upload", "webcam"], type="pil", label="Upload / Webcam")
+            group_mode = gr.Checkbox(False, label="Group photo (detect everyone)")
+            auto = gr.Checkbox(True, label="Auto face crop (wide)")
+            prompt = gr.Textbox(label="(Optional) Extra cartoon style",
+                                placeholder="e.g., studio ghibli watercolor, soft bokeh, pastel palette")
+            with gr.Row():
+                strength = gr.Slider(0.3, 0.8, value=0.5, step=0.05, label="Cartoon strength")
+                steps = gr.Slider(1, 4, value=2, step=1, label="Turbo steps (1–4)")
+                seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
+            btn_age = gr.Button("Predict Age(s) (fast)", variant="primary")
+            btn_cartoon = gr.Button("Make Cartoon(s) (fast)", variant="secondary")
+
+        with gr.Column(scale=1):
+            probs_out = gr.Label(num_top_classes=5, label="Age Prediction (probabilities, first face)")
+            age_md = gr.Markdown(label="Age Table / Summary")
+            preview = gr.Image(label="Detection Preview (boxes)")
+            cartoon_out = gr.Image(label="Cartoon Result (grid for groups)")
+
+    btn_age.click(fn=predict_age, inputs=[img_in, group_mode, auto], outputs=[probs_out, age_md, preview])
+    btn_cartoon.click(fn=cartoonize, inputs=[img_in, prompt, group_mode, auto, strength, steps, seed], outputs=cartoon_out)
+
+# Expose for Hugging Face Spaces
+app = demo
+
+if __name__ == "__main__":
+    app.queue().launch()