Update app.py

app.py

@@ -7,24 +7,6 @@ import gradio as gr
 from PIL import Image
 from huggingface_hub import hf_hub_download
 
-# ── Model — exactly as in the Pony V7 Captioner notebook ───────────────────────
-class AestheticScorer(nn.Module):
-    def __init__(self, input_size: int = 768):
-        super().__init__()
-        self.model = nn.Sequential(
-            nn.Linear(input_size, 1024),
-            nn.ReLU(),
-            nn.Dropout(0.5),
-            nn.Linear(1024, 512),
-            nn.ReLU(),
-            nn.Dropout(0.3),
-            nn.Linear(512, 1),
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.model(x)
-
-
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"[info] device: {DEVICE}")
 
@@ -39,35 +21,39 @@ ckpt_path = hf_hub_download(
 )
 checkpoint_data = torch.load(ckpt_path, map_location=DEVICE)
 state_dict = checkpoint_data["state_dict"]
-# Strip the "model." prefix from keys (same as notebook)
+# Strip "model." prefix — keys become "0.weight", "3.weight", "6.weight"
 state_dict = {k.replace("model.", ""): v for k, v in state_dict.items()}
 
-aesthetic_model = AestheticScorer(input_size=768).to(DEVICE)
+# Build Sequential directly so keys match ("0.weight", "3.weight", "6.weight")
+aesthetic_model = nn.Sequential(
+    nn.Linear(768, 1024),
+    nn.ReLU(),
+    nn.Dropout(0.5),
+    nn.Linear(1024, 512),
+    nn.ReLU(),
+    nn.Dropout(0.3),
+    nn.Linear(512, 1),
+).to(DEVICE)
 aesthetic_model.load_state_dict(state_dict)
 aesthetic_model.eval()
 print("[info] Model ready.")
 
 
-# ── Scoring — identical to notebook ────────────────────────────────────────────
 @torch.no_grad()
 def get_score(image: Image.Image) -> float:
-    """Returns raw float score (typically 0-1 range)."""
     image_tensor = preprocess(image.convert("RGB")).unsqueeze(0).to(DEVICE)
     features = clip_model.encode_image(image_tensor).cpu().numpy()
     norm = np.linalg.norm(features, axis=1, keepdims=True)
     norm[norm == 0] = 1
     features = features / norm
     features_t = torch.tensor(features, dtype=torch.float32, device=DEVICE)
-    raw = aesthetic_model(features_t).item()
-    return raw
+    return aesthetic_model(features_t).item()
 
 
 def raw_to_pony(raw: float) -> int:
-    """Convert raw score to pony score_0...score_9 (same formula as notebook)."""
     return int(max(0.0, min(0.99, raw)) * 10)
 
 
-# ── Colour palette ─────────────────────────────────────────────────────────────
 SCORE_COLOURS = [
     "#c0392b", "#e74c3c", "#e67e22", "#f39c12", "#d4ac0d",
     "#27ae60", "#1e8449", "#148f77", "#0e6655", "#0a4f42",
@@ -75,7 +61,7 @@ SCORE_COLOURS = [
 
 
 def build_html(raw: float) -> str:
-    pony = raw_to_pony(raw)
+    pony   = raw_to_pony(raw)
     colour = SCORE_COLOURS[pony]
 
     tiles_html = ""
@@ -86,12 +72,13 @@ def build_html(raw: float) -> str:
         weight = "700" if active else "400"
         scale  = "scale(1.12)" if active else "scale(1)"
         opac   = "1" if active else "0.45"
-        tiles_html += f"""<div style="background:{bg};border:{border};border-radius:8px;
-            padding:10px 0;text-align:center;font-size:.82rem;font-weight:{weight};color:#fff;
-            transform:{scale};opacity:{opac};transition:all .2s;user-select:none;">score_{i}</div>"""
+        tiles_html += (
+            f'<div style="background:{bg};border:{border};border-radius:8px;'
+            f'padding:10px 0;text-align:center;font-size:.82rem;font-weight:{weight};color:#fff;'
+            f'transform:{scale};opacity:{opac};transition:all .2s;user-select:none;">score_{i}</div>'
+        )
 
     bar_w = min(raw, 1.0) * 100
-
     return f"""
 <div style="font-family:'Inter',sans-serif;padding:8px 0;">
   <div style="text-align:center;margin-bottom:20px;">
@@ -99,7 +86,7 @@ def build_html(raw: float) -> str:
         padding:14px 36px;font-size:2rem;font-weight:800;letter-spacing:.04em;
         box-shadow:0 4px 20px {colour}66;">score_{pony}</div>
     <div style="color:#aaa;font-size:.85rem;margin-top:8px;">
-      raw score: <code style="color:#ddd">{raw:.4f}</code>
+      raw: <code style="color:#ddd">{raw:.4f}</code>
     </div>
   </div>
   <div style="display:grid;grid-template-columns:repeat(10,1fr);gap:6px;margin-bottom:16px;">
@@ -108,7 +95,7 @@ def build_html(raw: float) -> str:
   <div style="background:rgba(255,255,255,.1);border-radius:6px;height:8px;overflow:hidden;">
     <div style="width:{bar_w:.1f}%;height:100%;
         background:linear-gradient(90deg,#c0392b,#f39c12,#27ae60);
-        border-radius:6px;transition:width .4s;"></div>
+        border-radius:6px;"></div>
   </div>
   <div style="display:flex;justify-content:space-between;font-size:.72rem;color:#777;margin-top:4px;">
     <span>score_0</span><span>score_9</span>
@@ -119,26 +106,24 @@ def build_html(raw: float) -> str:
 def classify(image):
     if image is None:
         return "<p style='color:#888;text-align:center'>Upload an image to score it.</p>"
-    raw = get_score(image)
-    return build_html(raw)
+    return build_html(get_score(image))
 
 
-# ── Gradio UI ───────────────────────────────────────────────────────────────────
 with gr.Blocks(
     title="Aesthetic Classifier - PurpleSmartAI",
     theme=gr.themes.Soft(primary_hue="purple"),
     css=".gradio-container{max-width:860px!important;margin:auto} #title{text-align:center} #sub{text-align:center;color:#888;font-size:.9rem;margin-bottom:1.5rem}",
 ) as demo:
-    gr.Markdown("# Aesthetic Classifier", elem_id="title")
+    gr.Markdown("# 🎨 Aesthetic Classifier", elem_id="title")
     gr.Markdown(
         "CLIP ViT-L/14 regression model by **PurpleSmartAI** for Pony V7 captioning. "
-        "Outputs a **score_0...score_9** tag used directly in training captions.",
+        "Outputs a **score_0…score_9** tag used directly in training captions.",
         elem_id="sub",
     )
     with gr.Row():
         with gr.Column(scale=1):
             img_input = gr.Image(type="pil", label="Input Image", height=340)
-            run_btn   = gr.Button("Score image", variant="primary", size="lg")
+            run_btn   = gr.Button("✨ Score image", variant="primary", size="lg")
         with gr.Column(scale=1):
             out_html  = gr.HTML(
                 value="<p style='color:#888;text-align:center;padding:40px 0'>Upload an image to see its score.</p>",