Spaces:
Sleeping
Sleeping
| """ | |
| 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 ← 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 | |
| import copy | |
| import random | |
| import numpy as np | |
| import gradio as gr | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from PIL import Image | |
| # ─── Konstanten (müssen exakt mit train_diffusion.py übereinstimmen) ────────── | |
| IMG_SIZE = 64 | |
| CHANNELS = 4 | |
| EMBED_DIM = 256 | |
| T_STEPS = 500 | |
| BETA_START = 1e-4 | |
| BETA_END = 0.02 | |
| # ─── Tags (identisch mit Training) ──────────────────────────────────────────── | |
| BEREICHE = ["head","body","arm_l","arm_r","leg_l","leg_r"] | |
| FARBEN = ["orange","red","blue","green","cyan","yellow","pink","purple","black","white","gray","brown","beige"] | |
| HELL = ["bright","medium","dark"] | |
| KLEIDUNG = ["hoodie","shirt","tshirt","jacket","coat","armor","robe","suit","dress","cape","vest","sweater","uniform","casual","formal","jeans","pants","shorts","skirt"] | |
| STIL = ["player_skin","mob_skin","zombie","enderman","skeleton_like","custom","unknown","fantasy","modern","medieval","sci_fi","ninja","pirate","wizard","knight","archer","mage"] | |
| HAUTTONE = ["skin_light","skin_medium","skin_dark","skin_pale","skin_tan"] | |
| ACCESSOIRES = ["hat","helmet","crown","glasses","beard","hair_long","hair_short","wings","tail","horns","mask"] | |
| ALL_TAGS = [] | |
| for b in BEREICHE: | |
| for f in FARBEN: ALL_TAGS.append(f"{b}_{f}") | |
| for b in BEREICHE: | |
| for h in HELL: ALL_TAGS.append(f"{b}_{h}") | |
| ALL_TAGS += STIL + KLEIDUNG + HAUTTONE + ACCESSOIRES | |
| TAG2IDX = {t: i for i, t in enumerate(ALL_TAGS)} | |
| NUM_TAGS = len(ALL_TAGS) | |
| # ─── Prompt-Parser ──────────────────────────────────────────────────────────── | |
| PROMPT_KEYWORDS = { | |
| "rot":"red","rote":"red","roter":"red","blau":"blue","blaue":"blue", | |
| "gruen":"green","grune":"green","gelb":"yellow","orange":"orange", | |
| "cyan":"cyan","trkis":"cyan","pink":"pink","rosa":"pink", | |
| "lila":"purple","violett":"purple","schwarz":"black","schwarze":"black", | |
| "weiss":"white","grau":"gray","braun":"brown","beige":"beige", | |
| "red":"red","blue":"blue","green":"green","yellow":"yellow","cyan":"cyan", | |
| "pink":"pink","purple":"purple","black":"black","white":"white", | |
| "gray":"gray","grey":"gray","brown":"brown", | |
| "hell":"bright","bright":"bright","dunkel":"dark","dark":"dark","mittel":"medium","medium":"medium", | |
| "zombie":"zombie","enderman":"enderman","skelett":"skeleton_like","skeleton":"skeleton_like", | |
| "armor":"armor","player":"player_skin","custom":"custom", | |
| "hoodie":"hoodie","hemd":"shirt","shirt":"shirt", | |
| "tshirt":"tshirt","jacke":"jacket","jacket":"jacket", | |
| "mantel":"coat","coat":"coat","robe":"robe","anzug":"suit","suit":"suit", | |
| "kleid":"dress","dress":"dress","umhang":"cape","cape":"cape", | |
| "weste":"vest","vest":"vest","pullover":"sweater","sweater":"sweater", | |
| "uniform":"uniform","casual":"casual","formal":"formal", | |
| "jeans":"jeans","hose":"pants","pants":"pants","shorts":"shorts","skirt":"skirt", | |
| "fantasy":"fantasy","modern":"modern","medieval":"medieval", | |
| "scifi":"sci_fi","ninja":"ninja","pirate":"pirate", | |
| "wizard":"wizard","knight":"knight","archer":"archer","mage":"mage", | |
| "pale":"skin_pale","tan":"skin_tan", | |
| "hat":"hat","helmet":"helmet","crown":"crown", | |
| "glasses":"glasses","beard":"beard", | |
| "wings":"wings","horns":"horns","mask":"mask", | |
| } | |
| _COLOR_BODY_PARTS = { | |
| "hoodie":["body","arm_l","arm_r"],"shirt":["body"],"tshirt":["body"], | |
| "jacket":["body","arm_l","arm_r"],"coat":["body","arm_l","arm_r"], | |
| "jeans":["leg_l","leg_r"],"pants":["leg_l","leg_r"],"shorts":["leg_l","leg_r"],"skirt":["leg_l","leg_r"], | |
| "default":["head","body","arm_l","arm_r","leg_l","leg_r"], | |
| } | |
| def parse_prompt(prompt: str) -> list: | |
| words = prompt.lower().replace(","," ").replace("-"," ").split() | |
| tags, pending_color, pending_garment = set(), None, None | |
| for w in words: | |
| resolved = PROMPT_KEYWORDS.get(w) | |
| if resolved in FARBEN: | |
| pending_color = resolved | |
| if pending_garment is None: | |
| for b in _COLOR_BODY_PARTS["default"]: tags.add(f"{b}_{resolved}") | |
| elif resolved in KLEIDUNG: | |
| pending_garment = resolved | |
| tags.add(resolved) | |
| if pending_color: | |
| for b in _COLOR_BODY_PARTS.get(resolved, _COLOR_BODY_PARTS["default"]): | |
| tags.add(f"{b}_{pending_color}") | |
| elif resolved in HELL: | |
| for b in ["head","body"]: tags.add(f"{b}_{resolved}") | |
| elif resolved and resolved in TAG2IDX: | |
| tags.add(resolved) | |
| elif w in TAG2IDX: | |
| tags.add(w) | |
| if not any(t in tags for t in STIL): | |
| tags.add("player_skin") | |
| return list(tags) | |
| def tags_to_vector(tags: list) -> torch.Tensor: | |
| vec = torch.zeros(NUM_TAGS) | |
| for t in tags: | |
| if t in TAG2IDX: vec[TAG2IDX[t]] = 1.0 | |
| return vec | |
| # ─── UV-Masken (identisch mit train_diffusion.py) ───────────────────────────── | |
| SKIN_REGIONS = { | |
| "head": (0, 0, 32, 16), | |
| "body": (16, 16, 40, 32), | |
| "arm_r": (40, 16, 56, 32), | |
| "leg_r": (0, 16, 16, 32), | |
| "arm_l": (32, 48, 48, 64), | |
| "leg_l": (16, 48, 32, 64), | |
| } | |
| OVERLAY_REGIONS = { | |
| "head_overlay": (32, 0, 64, 16), | |
| "body_overlay": (16, 32, 40, 48), | |
| "arm_r_overlay": (40, 32, 56, 48), | |
| "leg_r_overlay": (0, 32, 16, 48), | |
| "arm_l_overlay": (48, 48, 64, 64), | |
| "leg_l_overlay": (0, 48, 16, 64), | |
| } | |
| # [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(): | |
| 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_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_new], dim=1) | |
| # ─── UNet (identisch mit train_diffusion.py) ────────────────────────────────── | |
| class SinusoidalPE(nn.Module): | |
| def __init__(self, dim): | |
| super().__init__() | |
| self.dim = dim | |
| def forward(self, t): | |
| 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) | |
| class CondEmbed(nn.Module): | |
| def __init__(self, num_tags, embed_dim): | |
| super().__init__() | |
| self.net = nn.Sequential( | |
| nn.Linear(num_tags, embed_dim * 2), nn.SiLU(), | |
| nn.Linear(embed_dim * 2, embed_dim), nn.SiLU(), | |
| ) | |
| def forward(self, c): return self.net(c) | |
| class ResBlock(nn.Module): | |
| def __init__(self, in_ch, out_ch, time_dim, dropout=0.1): | |
| super().__init__() | |
| self.norm1 = nn.GroupNorm(min(8, in_ch), in_ch) | |
| self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1) | |
| self.norm2 = nn.GroupNorm(min(8, out_ch), out_ch) | |
| self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1) | |
| self.time_mlp = nn.Sequential(nn.SiLU(), nn.Linear(time_dim, out_ch * 2)) | |
| self.skip = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity() | |
| self.dropout = nn.Dropout(dropout) | |
| self.act = nn.SiLU() | |
| def forward(self, x, t_emb): | |
| h = self.conv1(self.act(self.norm1(x))) | |
| t = self.time_mlp(t_emb)[:, :, None, None] | |
| scale, shift = t.chunk(2, dim=1) | |
| h = self.norm2(h) * (1 + scale) + shift | |
| h = self.conv2(self.dropout(self.act(h))) | |
| return h + self.skip(x) | |
| class AttentionBlock(nn.Module): | |
| def __init__(self, ch, heads=4): | |
| super().__init__() | |
| self.norm = nn.GroupNorm(min(8, ch), ch) | |
| self.attn = nn.MultiheadAttention(ch, heads, batch_first=True) | |
| self.proj = nn.Conv2d(ch, ch, 1) | |
| def forward(self, x): | |
| B, C, H, W = x.shape | |
| h = self.norm(x).view(B, C, H * W).permute(0, 2, 1) | |
| h, _ = self.attn(h, h, h) | |
| h = h.permute(0, 2, 1).view(B, C, H, W) | |
| return x + self.proj(h) | |
| class UNet(nn.Module): | |
| def __init__(self, channels=CHANNELS, base_ch=96, embed_dim=EMBED_DIM): | |
| super().__init__() | |
| time_dim = embed_dim * 2 | |
| cond_dim = embed_dim | |
| self.time_pe = SinusoidalPE(embed_dim) | |
| self.time_mlp = nn.Sequential( | |
| nn.Linear(embed_dim, time_dim), nn.SiLU(), nn.Linear(time_dim, time_dim) | |
| ) | |
| self.cond_emb = CondEmbed(NUM_TAGS, cond_dim) | |
| self.cond_mlp = nn.Linear(cond_dim, time_dim) | |
| ch = base_ch | |
| self.enc_in = nn.Conv2d(channels, ch, 3, padding=1) | |
| self.enc1 = ResBlock(ch, ch, time_dim) | |
| self.enc1b = ResBlock(ch, ch, time_dim) | |
| self.down1 = nn.Conv2d(ch, ch, 4, stride=2, padding=1) | |
| self.enc2 = ResBlock(ch, ch*2, time_dim) | |
| self.enc2b = ResBlock(ch*2, ch*2, time_dim) | |
| self.down2 = nn.Conv2d(ch*2, ch*2, 4, stride=2, padding=1) | |
| self.enc3 = ResBlock(ch*2, ch*4, time_dim) | |
| self.enc3b = ResBlock(ch*4, ch*4, time_dim) | |
| self.attn3 = AttentionBlock(ch*4) | |
| self.down3 = nn.Conv2d(ch*4, ch*4, 4, stride=2, padding=1) | |
| self.mid1 = ResBlock(ch*4, ch*4, time_dim) | |
| self.mid_att = AttentionBlock(ch*4) | |
| self.mid2 = ResBlock(ch*4, ch*4, time_dim) | |
| self.up3 = nn.ConvTranspose2d(ch*4, ch*4, 4, stride=2, padding=1) | |
| self.dec3 = ResBlock(ch*8, ch*4, time_dim) | |
| self.dec3b = ResBlock(ch*4, ch*4, time_dim) | |
| self.attn_d3 = AttentionBlock(ch*4) | |
| self.up2 = nn.ConvTranspose2d(ch*4, ch*2, 4, stride=2, padding=1) | |
| self.dec2 = ResBlock(ch*4, ch*2, time_dim) | |
| self.dec2b = ResBlock(ch*2, ch*2, time_dim) | |
| self.up1 = nn.ConvTranspose2d(ch*2, ch, 4, stride=2, padding=1) | |
| self.dec1 = ResBlock(ch*2, ch, time_dim) | |
| self.dec1b = ResBlock(ch, ch, time_dim) | |
| self.out = nn.Sequential( | |
| nn.GroupNorm(min(8, ch), ch), nn.SiLU(), nn.Conv2d(ch, channels, 3, padding=1) | |
| ) | |
| def forward(self, x, t, cond): | |
| t_emb = self.time_mlp(self.time_pe(t)) | |
| c_emb = self.cond_mlp(self.cond_emb(cond)) | |
| emb = t_emb + c_emb | |
| h0 = self.enc_in(x) | |
| h1 = self.enc1b(self.enc1(h0, emb), emb) | |
| h2 = self.enc2b(self.enc2(self.down1(h1), emb), emb) | |
| h3 = self.attn3(self.enc3b(self.enc3(self.down2(h2), emb), emb)) | |
| h = self.mid2(self.mid_att(self.mid1(self.down3(h3), emb)), emb) | |
| h = self.attn_d3(self.dec3b(self.dec3(torch.cat([self.up3(h), h3], 1), emb), emb)) | |
| h = self.dec2b(self.dec2(torch.cat([self.up2(h), h2], 1), emb), emb) | |
| h = self.dec1b(self.dec1(torch.cat([self.up1(h), h1], 1), emb), emb) | |
| return self.out(h) | |
| # ─── Diffusion Schedule ─────────────────────────────────────────────────────── | |
| class DiffusionSchedule: | |
| def __init__(self, T=T_STEPS, device="cpu"): | |
| self.T = T | |
| self.device = device | |
| steps = T + 1 | |
| x = torch.linspace(0, T, steps) | |
| alphas = torch.cos(((x / T) + 0.008) / 1.008 * math.pi / 2) ** 2 | |
| alphas = alphas / alphas[0] | |
| betas = (1 - alphas[1:] / alphas[:-1]).clamp(0, 0.999) | |
| self.betas = betas.to(device) | |
| self.alphas = 1.0 - self.betas | |
| self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) | |
| self.alphas_cumprod_prev = F.pad(self.alphas_cumprod[:-1], (1, 0), value=1.0) | |
| self.posterior_variance = ( | |
| self.betas * (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) | |
| ) | |
| def p_sample(self, model, x, t_idx, cond, guidance_scale): | |
| t_tensor = torch.full((x.shape[0],), t_idx, device=self.device, dtype=torch.long) | |
| null_cond = torch.zeros_like(cond) | |
| x2 = torch.cat([x, x]) | |
| t2 = torch.cat([t_tensor, t_tensor]) | |
| c2 = torch.cat([cond, null_cond]) | |
| out = model(x2, t2, c2) | |
| 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] | |
| beta = self.betas[t_idx] | |
| coef = beta / (1 - alpha_bar).sqrt() | |
| mean = (1 / alpha.sqrt()) * (x - coef * noise_pred) | |
| if t_idx > 0: | |
| return mean + self.posterior_variance[t_idx].sqrt() * torch.randn_like(x) | |
| return mean | |
| def sample(self, model, cond, n=1, steps=50, guidance_scale=3.0): | |
| model.eval() | |
| x = torch.randn(n, CHANNELS, IMG_SIZE, IMG_SIZE, device=self.device) | |
| for t_idx in torch.linspace(self.T - 1, 0, steps, dtype=torch.long, device=self.device): | |
| x = self.p_sample(model, x, t_idx.item(), cond, guidance_scale) | |
| return force_alpha_mask(x).clamp(-1, 1) | |
| # ─── Modell laden ───────────────────────────────────────────────────────────── | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Device: {device}") | |
| 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"): | |
| if key in sd: | |
| base_ch = sd[key].shape[0] | |
| print(f"base_ch aus state_dict ermittelt: {base_ch}") | |
| break | |
| 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) | |
| model.load_state_dict(sd, strict=False) | |
| model.eval() | |
| try: | |
| torch._dynamo.disable(model) | |
| except Exception: | |
| pass | |
| schedule = DiffusionSchedule(device=device) | |
| 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 ──────────────────────────────────────────────────── | |
| def generate(prompt, num_skins, steps, guidance_scale, seed, randomize_seed): | |
| if randomize_seed: | |
| seed = random.randint(0, 2**31) | |
| torch.manual_seed(seed) | |
| tags = parse_prompt(prompt) | |
| tag_str = ", ".join(tags) if tags else "–" | |
| cond = tags_to_vector(tags).to(device).unsqueeze(0).expand(num_skins, -1) | |
| with torch.inference_mode(): | |
| imgs = schedule.sample(model, cond, n=num_skins, steps=steps, guidance_scale=guidance_scale) | |
| results = [] | |
| for img_t in imgs: | |
| arr = ((img_t.cpu().permute(1, 2, 0).numpy() + 1) * 127.5).clip(0, 255).astype(np.uint8) | |
| pil = Image.fromarray(arr, "RGBA").resize((512, 512), Image.NEAREST) | |
| results.append(pil) | |
| return results, f"Tags erkannt: {tag_str}", seed | |
| # ─── Gradio UI ──────────────────────────────────────────────────────────────── | |
| EXAMPLES = [ | |
| ["roter hoodie blaue jeans", 4, 50, 3.0], | |
| ["zombie", 4, 50, 4.0], | |
| ["wizard fantasy purple", 4, 50, 3.5], | |
| ["knight medieval armor", 4, 50, 3.0], | |
| ["ninja black dark", 4, 50, 4.0], | |
| ["enderman", 2, 50, 3.0], | |
| ] | |
| css = "#gallery { min-height: 300px; }" | |
| with gr.Blocks(css=css, title="Minecraft Skin Generator") as demo: | |
| gr.Markdown(""" | |
| # 🎮 Minecraft Skin Generator | |
| Generiert 64×64 Minecraft Skins aus einem Text-Prompt. Trainiert mit DDPM auf ~41k Skins. | |
| **Beispiel-Prompts:** `roter hoodie blaue jeans` · `zombie` · `knight medieval armor` · `wizard fantasy purple` | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| prompt = gr.Text( | |
| label="Prompt", | |
| placeholder="z.B. roter hoodie blaue jeans", | |
| lines=1, | |
| ) | |
| run_btn = gr.Button("Generieren", variant="primary", size="lg") | |
| with gr.Accordion("Einstellungen", open=False): | |
| num_skins = gr.Slider(label="Anzahl Skins", minimum=1, maximum=8, step=1, value=4) | |
| steps = gr.Slider(label="Diffusion-Schritte", minimum=10, maximum=100, step=5, value=50) | |
| guidance = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=15.0,step=0.5, value=3.0) | |
| 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) | |
| seed_out = gr.Number(label="Verwendeter Seed", interactive=False) | |
| with gr.Column(scale=3): | |
| gallery = gr.Gallery( | |
| label="Generierte Skins (512×512 hochskaliert)", | |
| elem_id="gallery", | |
| columns=4, | |
| rows=2, | |
| object_fit="contain", | |
| height=400, | |
| ) | |
| gr.Examples( | |
| examples=EXAMPLES, | |
| inputs=[prompt, num_skins, steps, guidance], | |
| label="Beispiele", | |
| ) | |
| gr.on( | |
| triggers=[run_btn.click, prompt.submit], | |
| fn=generate, | |
| inputs=[prompt, num_skins, steps, guidance, seed, rand_seed], | |
| outputs=[gallery, tag_info, seed_out], | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |