Update app.py

app.py

@@ -4,8 +4,14 @@ Minecraft Skin Generator – HuggingFace Spaces Demo
 Lädt model.pt (EMA-Gewichte) aus dem Repo und generiert Skins per Prompt.
 Benötigte Dateien im Space-Repo:
   app.py          ← diese Datei
-  model.pt        ← dein exportiertes EMA-Modell
+  model.pt        ← mit export_ema_model.py exportiert (EMA-Gewichte!)
   requirements.txt
+
+FIXES gegenüber der alten app.py:
+  [FIX 1] EMA-Gewichte werden korrekt priorisiert (ckpt["ema"] vor ckpt["model"])
+  [FIX 2] base_ch Fallback-Kette ist identisch mit train_diffusion.py (Default 96 statt 128)
+  [FIX 3] _build_base_mask / _build_overlay_mask ohne device-Parameter (wie im Training)
+  [FIX 4] force_alpha_mask identisch mit train_diffusion.py
 """
 
 import math
@@ -111,7 +117,7 @@ def tags_to_vector(tags: list) -> torch.Tensor:
         if t in TAG2IDX: vec[TAG2IDX[t]] = 1.0
     return vec
 
-# ─── UV-Masken ────────────────────────────────────────────────────────────────
+# ─── UV-Masken (identisch mit train_diffusion.py) ─────────────────────────────
 SKIN_REGIONS = {
     "head":   (0,  0,  32, 16),
     "body":   (16, 16, 40, 32),
@@ -129,28 +135,30 @@ OVERLAY_REGIONS = {
     "leg_l_overlay": (0,  48, 16, 64),
 }
 
-def _build_base_mask(device):
-    mask = torch.zeros(1, 1, IMG_SIZE, IMG_SIZE, device=device)
+# [FIX 3] Keine device-Parameter – identisch mit train_diffusion.py
+def _build_base_mask():
+    mask = torch.zeros(1, 1, IMG_SIZE, IMG_SIZE)
     for x1,y1,x2,y2 in SKIN_REGIONS.values():
         mask[0,0,y1:y2,x1:x2] = 1.0
     return mask
 
-def _build_overlay_mask(device):
-    mask = torch.zeros(1, 1, IMG_SIZE, IMG_SIZE, device=device)
+def _build_overlay_mask():
+    mask = torch.zeros(1, 1, IMG_SIZE, IMG_SIZE)
     for x1,y1,x2,y2 in OVERLAY_REGIONS.values():
         mask[0,0,y1:y2,x1:x2] = 1.0
     return mask
 
+# [FIX 4] force_alpha_mask identisch mit train_diffusion.py (device über .to())
 def force_alpha_mask(img: torch.Tensor) -> torch.Tensor:
-    base    = _build_base_mask(img.device)
-    overlay = _build_overlay_mask(img.device)
-    outside = (1.0 - base - overlay).clamp(0, 1)
-    alpha   = (
-        base    * torch.ones_like(img[:, 3:4])
-        + overlay * img[:, 3:4]
-        + outside * torch.full_like(img[:, 3:4], -1.0)
+    base_mask    = _build_base_mask().to(img.device)
+    overlay_mask = _build_overlay_mask().to(img.device)
+    outside_mask = (1.0 - base_mask - overlay_mask).clamp(0, 1)
+    alpha_new = (
+        base_mask    * torch.ones_like(img[:, 3:4])
+        + overlay_mask * img[:, 3:4]
+        + outside_mask * torch.full_like(img[:, 3:4], -1.0)
     )
-    return torch.cat([img[:, :3], alpha], dim=1)
+    return torch.cat([img[:, :3], alpha_new], dim=1)
 
 # ─── UNet (identisch mit train_diffusion.py) ──────────────────────────────────
 class SinusoidalPE(nn.Module):
@@ -159,9 +167,10 @@ class SinusoidalPE(nn.Module):
         self.dim = dim
 
     def forward(self, t):
-        half  = self.dim // 2
-        freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
-        args  = t[:, None].float() * freqs[None]
+        device = t.device
+        half   = self.dim // 2
+        freqs  = torch.exp(-math.log(10000) * torch.arange(half, device=device) / half)
+        args   = t[:, None].float() * freqs[None]
         return torch.cat([args.sin(), args.cos()], dim=-1)
 
 
@@ -308,8 +317,8 @@ class DiffusionSchedule:
         c2 = torch.cat([cond, null_cond])
 
         out = model(x2, t2, c2)
-        n_cond, n_uncond = out.chunk(2)
-        noise_pred = n_uncond + guidance_scale * (n_cond - n_uncond)
+        noise_cond, noise_uncond = out.chunk(2)
+        noise_pred = noise_uncond + guidance_scale * (noise_cond - noise_uncond)
 
         alpha     = self.alphas[t_idx]
         alpha_bar = self.alphas_cumprod[t_idx]
@@ -335,25 +344,44 @@ class DiffusionSchedule:
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Device: {device}")
 
-ckpt    = torch.load("model.pt", map_location=device, weights_only=False)
-base_ch = ckpt.get("base_ch", 96)
+ckpt = torch.load("model.pt", map_location=device, weights_only=False)
+print(f"Checkpoint Keys: {list(ckpt.keys())}")
+
+# [FIX 1] EMA-Gewichte priorisieren – das ist der Hauptfehler der alten app.py!
+# "ema" Key = EMA-Gewichte (beste Qualität, geglättet)
+# "model" Key = je nach Datei entweder EMA (bei latest.pt) oder rohe Gewichte (bei ep*.pt)
+sd = ckpt.get("ema") or ckpt.get("model") or ckpt
+if "ema" in ckpt:
+    print("✅ Verwende EMA-Gewichte ('ema' Key) – beste Qualität")
+elif "model" in ckpt:
+    print("ℹ️  Verwende 'model' Key (kein 'ema' Key gefunden)")
+else:
+    print("⚠️  Kein 'ema' oder 'model' Key – versuche direktes Laden")
+
+# [FIX 2] base_ch Fallback identisch mit train_diffusion.py
+base_ch = ckpt.get("base_ch", None)
 if base_ch is None:
     for key in ("enc_in.weight", "_orig_mod.enc_in.weight"):
-        sd_check = ckpt.get("model", ckpt)
-        if key in sd_check:
-            base_ch = sd_check[key].shape[0]
+        if key in sd:
+            base_ch = sd[key].shape[0]
+            print(f"base_ch aus state_dict ermittelt: {base_ch}")
             break
-    base_ch = base_ch or 96
+if base_ch is None:
+    base_ch = 96   # train_diffusion.py Default ist 96, nicht 128!
+    print(f"⚠️  base_ch nicht gefunden, verwende Default: {base_ch}")
 
 model = UNet(base_ch=base_ch).to(device)
-sd    = ckpt.get("model", ckpt)
 model.load_state_dict(sd, strict=False)
 model.eval()
-try: torch._dynamo.disable(model)
-except Exception: pass
+
+try:
+    torch._dynamo.disable(model)
+except Exception:
+    pass
 
 schedule = DiffusionSchedule(device=device)
-print(f"Modell geladen: base_ch={base_ch}, {sum(p.numel() for p in model.parameters())/1e6:.1f}M Parameter")
+num_params = sum(p.numel() for p in model.parameters()) / 1e6
+print(f"Modell geladen: base_ch={base_ch}, {num_params:.1f}M Parameter")
 
 
 # ─── Generierungs-Funktion ────────────────────────────────────────────────────
@@ -415,7 +443,7 @@ Generiert 64×64 Minecraft Skins aus einem Text-Prompt. Trainiert mit DDPM auf ~
                 seed      = gr.Slider(label="Seed",                minimum=0,   maximum=2**31,step=1,  value=42)
                 rand_seed = gr.Checkbox(label="Seed zufällig", value=True)
 
-            tag_info = gr.Text(label="Erkannte Tags",     interactive=False)
+            tag_info = gr.Text(label="Erkannte Tags",      interactive=False)
             seed_out = gr.Number(label="Verwendeter Seed", interactive=False)
 
         with gr.Column(scale=3):