""" ShukGEN Advanced — Premium Gradio Application Identity-Preserving VAE + HD Super Resolution + 8 Style Variants """ import os import io import math import random import base64 import numpy as np from scipy.ndimage import gaussian_filter import torch import torch.nn as nn import torch.nn.functional as F from torchvision import transforms, models import torchvision.transforms.functional as TF from PIL import Image, ImageFilter, ImageEnhance import timm import gradio as gr # ───────────────────────────────────────────────────────────────────────────── # DEVICE # ───────────────────────────────────────────────────────────────────────────── DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # ───────────────────────────────────────────────────────────────────────────── # STYLE METADATA # ───────────────────────────────────────────────────────────────────────────── STYLE_NAMES = [ "Youthful", "Aged / Mature", "Dramatic Light", "Soft Glow", "Intense / Bold", "Warm Golden Hour", "Cool / Moody", "Sketch / Artistic", ] STYLE_EMOJIS = ["✨", "🎞️", "🎭", "🌸", "🔥", "🌅", "🌊", "🎨"] STYLE_COLORS = [ "#27ae60", "#c0392b", "#8e44ad", "#f39c12", "#e74c3c", "#d35400", "#2980b9", "#16a085" ] # ───────────────────────────────────────────────────────────────────────────── # MODEL ARCHITECTURE (exact copy from notebook) # ───────────────────────────────────────────────────────────────────────────── def norm(ch): for g in [8, 4, 2, 1]: if ch % g == 0: return nn.GroupNorm(g, ch) class ResBlock(nn.Module): def __init__(self, ch): super().__init__() self.net = nn.Sequential( norm(ch), nn.SiLU(), nn.Conv2d(ch, ch, 3, padding=1, bias=False), norm(ch), nn.SiLU(), nn.Conv2d(ch, ch, 3, padding=1, bias=False), ) def forward(self, x): return x + self.net(x) class SE(nn.Module): def __init__(self, ch, ratio=8): super().__init__() mid = max(ch // ratio, 4) self.net = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(ch, mid), nn.SiLU(), nn.Linear(mid, ch), nn.Sigmoid(), ) def forward(self, x): return x * self.net(x).view(x.size(0), x.size(1), 1, 1) class AttnBlock(nn.Module): def __init__(self, ch, num_heads=8): super().__init__() while ch % num_heads != 0 and num_heads > 1: num_heads //= 2 self.norm = norm(ch) self.attn = nn.MultiheadAttention(ch, num_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) h = h.reshape(B, C, H * W).permute(0, 2, 1) h, _ = self.attn(h, h, h) h = h.permute(0, 2, 1).reshape(B, C, H, W) return x + self.proj(h) class DownBlock(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv = nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False), norm(out_ch), nn.SiLU()) self.res = ResBlock(out_ch) self.se = SE(out_ch) self.down = nn.Conv2d(out_ch, out_ch, 4, stride=2, padding=1, bias=False) def forward(self, x): x = self.se(self.res(self.conv(x))) skip = x return self.down(x), skip class UpBlock(nn.Module): def __init__(self, in_ch, skip_ch, out_ch): super().__init__() self.up_conv = nn.Conv2d(in_ch, in_ch * 4, 1, bias=False) self.shuffle = nn.PixelShuffle(2) self.conv = nn.Conv2d(in_ch + skip_ch, out_ch, 3, padding=1, bias=False) self.norm_l = norm(out_ch) self.act = nn.SiLU() self.res = ResBlock(out_ch) def forward(self, x, skip): x = self.shuffle(self.up_conv(x)) if x.shape[2:] != skip.shape[2:]: x = F.interpolate(x, size=skip.shape[2:], mode='bilinear', align_corners=False) x = torch.cat([x, skip], dim=1) return self.res(self.act(self.norm_l(self.conv(x)))) class FaceVAE(nn.Module): def __init__(self, latent_dim=512, base_filters=48): super().__init__() f = base_filters self.f = f self.latent = latent_dim self.pool_sz = 4 flat = f * 16 * self.pool_sz * self.pool_sz # ── Encoder ────────────────────────────────────────────────────────── self.inc = nn.Sequential(nn.Conv2d(3, f, 3, padding=1, bias=False), norm(f), nn.SiLU()) self.down1 = DownBlock(f, f*2) self.down2 = DownBlock(f*2, f*4) self.down3 = DownBlock(f*4, f*8) self.down4 = DownBlock(f*8, f*16) self.btn = nn.Sequential(ResBlock(f*16), AttnBlock(f*16), ResBlock(f*16), SE(f*16)) self.pool = nn.AdaptiveAvgPool2d(self.pool_sz) self.fc_mu = nn.Linear(flat, latent_dim) self.fc_logvar = nn.Linear(flat, latent_dim) # ── Decoder ────────────────────────────────────────────────────────── self.fc_dec = nn.Linear(latent_dim, flat) # Up-blocks: up1 = first (right after bottleneck), up4 = last (near output) self.up1 = UpBlock(f*16, f*16, f*8) # uses s4 → f*8 output self.up2 = UpBlock(f*8, f*8, f*4) # uses s3 → f*4 output # Mid-decoder attention at f*4 resolution (192 channels when base_filters=48) self.attn_d = AttnBlock(f*4) self.up3 = UpBlock(f*4, f*4, f*2) # uses s2 → f*2 output self.up4 = UpBlock(f*2, f*2, f) # uses s1 → f output # Output head self.outc = nn.Sequential( norm(f), nn.SiLU(), nn.Conv2d(f, 3, 3, padding=1), nn.Tanh() ) def encode(self, x): x1 = self.inc(x) x2, s1 = self.down1(x1) x3, s2 = self.down2(x2) x4, s3 = self.down3(x3) x5, s4 = self.down4(x4) b = self.btn(x5) p = self.pool(b).flatten(1) return self.fc_mu(p), self.fc_logvar(p), (s1, s2, s3, s4) def reparameterize(self, mu, logvar): if self.training: std = (0.5 * logvar).exp() return mu + std * torch.randn_like(std) return mu def decode(self, z, skips): s1, s2, s3, s4 = skips f = self.f h = self.fc_dec(z).view(z.size(0), f*16, self.pool_sz, self.pool_sz) h = F.interpolate(h, size=s4.shape[2:], mode='bilinear', align_corners=False) h = self.up1(h, s4) # f*16 → f*8 h = self.up2(h, s3) # f*8 → f*4 h = self.attn_d(h) # mid-decoder attention at f*4 (192 ch) h = self.up3(h, s2) # f*4 → f*2 h = self.up4(h, s1) # f*2 → f return self.outc(h) def forward(self, x): mu, logvar, skips = self.encode(x) z = self.reparameterize(mu, logvar) return self.decode(z, skips), mu, logvar def reconstruct(self, x): mu, _, skips = self.encode(x) return self.decode(mu, skips) class SRResBlock(nn.Module): def __init__(self, ch=64): super().__init__() self.net = nn.Sequential( nn.Conv2d(ch, ch, 3, padding=1, bias=False), nn.BatchNorm2d(ch), nn.PReLU(num_parameters=ch), nn.Conv2d(ch, ch, 3, padding=1, bias=False), nn.BatchNorm2d(ch), ) def forward(self, x): return x + self.net(x) class SRHead(nn.Module): def __init__(self, n_blocks=10, ch=64): super().__init__() self.head = nn.Sequential( nn.Conv2d(3, ch, 9, padding=4), nn.PReLU(num_parameters=ch) ) self.body = nn.Sequential(*[SRResBlock(ch) for _ in range(n_blocks)]) self.body_end = nn.Sequential( nn.Conv2d(ch, ch, 3, padding=1, bias=False), nn.BatchNorm2d(ch) ) self.upsample = nn.Sequential( nn.Conv2d(ch, ch * 4, 3, padding=1), nn.PixelShuffle(2), nn.PReLU(num_parameters=ch) ) self.tail = nn.Sequential( nn.Conv2d(ch, 3, 9, padding=4), nn.Tanh() ) def forward(self, x): head = self.head(x) body = self.body(head) body = self.body_end(body) + head up = self.upsample(body) return self.tail(up) # ───────────────────────────────────────────────────────────────────────────── # STYLE TRANSFORM BANK (exact copy from notebook) # ───────────────────────────────────────────────────────────────────────────── class StyleTransformBank: @staticmethod def _np(img): return np.array(img, dtype=np.float32) / 255.0 @staticmethod def _pil(arr): return Image.fromarray((arr.clip(0, 1) * 255).astype(np.uint8)) @staticmethod def _gamma(arr, g): return np.power(arr.clip(1e-6, 1.0), g) @staticmethod def _rgb_to_hsv(rgb): r, g, b = rgb[...,0], rgb[...,1], rgb[...,2] maxc = np.maximum(np.maximum(r, g), b) minc = np.minimum(np.minimum(r, g), b) v = maxc s = np.where(maxc != 0, (maxc - minc) / (maxc + 1e-10), 0.0) d = maxc - minc + 1e-10 h = np.zeros_like(r) m = maxc != minc mr = m & (maxc == r) mg = m & (maxc == g) mb = m & (maxc == b) h[mr] = (60 * ((g[mr] - b[mr]) / d[mr])) % 360 h[mg] = 60 * ((b[mg] - r[mg]) / d[mg]) + 120 h[mb] = 60 * ((r[mb] - g[mb]) / d[mb]) + 240 return np.stack([h/360, s, v], axis=-1) @staticmethod def _hsv_to_rgb(hsv): h, s, v = hsv[...,0]*360, hsv[...,1], hsv[...,2] hi = (h / 60).astype(int) % 6 f = (h / 60) - np.floor(h / 60) p = v * (1 - s) q = v * (1 - f * s) t = v * (1 - (1 - f) * s) rgb = np.zeros(hsv.shape, dtype=np.float32) for i, (rv, gv, bv) in enumerate([(v,t,p),(q,v,p),(p,v,t),(p,q,v),(t,p,v),(v,p,q)]): m = hi == i rgb[...,0][m] = rv[m]; rgb[...,1][m] = gv[m]; rgb[...,2][m] = bv[m] return rgb @classmethod def _saturate(cls, arr, factor): hsv = cls._rgb_to_hsv(arr) hsv[...,1] = (hsv[...,1] * factor).clip(0, 1) return cls._hsv_to_rgb(hsv) @staticmethod def _vignette(arr, strength=0.5): H, W = arr.shape[:2] y = np.linspace(-1, 1, H)[:, None] x = np.linspace(-1, 1, W)[None, :] mask = 1.0 - strength * (x**2 + y**2) return arr * mask[..., None].clip(0, 1) @staticmethod def _grain(arr, amount=0.04, seed=42): rng = np.random.default_rng(seed=seed) return (arr + rng.normal(0, amount, arr.shape).astype(np.float32)).clip(0, 1) @staticmethod def _blur(arr, sigma): return np.stack( [gaussian_filter(arr[...,c], sigma) for c in range(3)], axis=-1 ).astype(np.float32) @classmethod def style_0_youthful(cls, img): arr = cls._np(img) arr = (arr * 1.22).clip(0, 1) arr = cls._saturate(arr, 1.55) blr = cls._blur(arr, sigma=3.0) lum = arr.mean(axis=-1, keepdims=True) mask = np.exp(-((lum - 0.6)**2) / (2*0.2**2)) arr = arr * (1 - 0.45*mask) + blr * (0.45*mask) arr[...,0] = (arr[...,0] * 1.06).clip(0,1) pil = cls._pil(arr) return ImageEnhance.Sharpness(pil).enhance(0.5) @classmethod def style_1_aged_mature(cls, img): arr = cls._np(img) arr = cls._saturate(arr, 0.25) grey = arr.mean(axis=-1, keepdims=True) sepia = np.concatenate([grey * 1.12, grey * 0.88, grey * 0.65], axis=-1) arr = arr * 0.35 + sepia * 0.65 arr = cls._gamma(arr, 1.55) arr = arr * 0.78 + 0.10 rng = np.random.default_rng(seed=7) tex = rng.normal(0, 0.03, arr.shape[:2]).astype(np.float32) tex = gaussian_filter(tex, 0.7)[..., None] arr = (arr + tex * 0.55).clip(0, 1) return cls._pil(arr) @classmethod def style_2_dramatic_light(cls, img): arr = cls._np(img) H, W = arr.shape[:2] x = np.linspace(0, 1, W)[None, :] shadow = 1.0 / (1.0 + np.exp(-14 * (x - 0.52))) shadow3 = shadow[..., None] light = cls._gamma(arr, 0.60) * 1.35 dark = cls._gamma(arr, 2.50) * 0.18 arr = light * (1 - shadow3) + dark * shadow3 arr = (arr - 0.5) * 3.0 + 0.5 arr = cls._saturate(arr, 0.80) arr = cls._vignette(arr, strength=0.75) return cls._pil(arr) @classmethod def style_3_soft_glow(cls, img): arr = cls._np(img) hi = np.clip(arr - 0.45, 0, 1) bloom = cls._blur(hi, sigma=14) arr = (arr + bloom * 2.2).clip(0, 1) soft = cls._blur(arr, sigma=1.8) arr = arr * 0.68 + soft * 0.32 arr[...,0] = (arr[...,0] * 1.18).clip(0,1) arr[...,1] = (arr[...,1] * 1.06).clip(0,1) arr[...,2] = (arr[...,2] * 0.82).clip(0,1) arr = arr * 0.82 + 0.15 arr = cls._saturate(arr, 1.35) return cls._pil(arr) @classmethod def style_4_intense_bold(cls, img): arr = cls._np(img) arr = cls._gamma(arr, 0.68) arr = cls._saturate(arr, 2.6) arr = (arr - 0.5) * 1.9 + 0.5 lum = arr.mean(axis=-1, keepdims=True) mid = np.exp(-((lum - 0.5)**2) / (2*0.25**2)) arr = (arr + mid * 0.07).clip(0, 1) pil = cls._pil(arr) pil = pil.filter(ImageFilter.UnsharpMask(radius=2, percent=200, threshold=3)) return pil.filter(ImageFilter.UnsharpMask(radius=1, percent=90, threshold=1)) @classmethod def style_5_warm_golden(cls, img): arr = cls._np(img) r = (arr[...,0] * 1.40 + 0.10).clip(0, 1) g = (arr[...,1] * 1.12 + 0.05).clip(0, 1) b = (arr[...,2] * 0.48 - 0.03).clip(0, 1) arr = np.stack([r, g, b], axis=-1) dark = np.exp(-arr.mean(axis=-1, keepdims=True) / 0.22) arr = (arr + dark * 0.22 * np.array([0.95, 0.62, 0.08])).clip(0,1) hi = np.clip(arr - 0.60, 0, 1) arr = (arr + cls._blur(hi, 10) * 1.4).clip(0, 1) arr = cls._saturate(arr, 1.25) return cls._pil(arr) @classmethod def style_6_cool_moody(cls, img): arr = cls._np(img) arr[...,0] = (arr[...,0] * 0.68 - 0.05).clip(0, 1) arr[...,1] = (arr[...,1] * 0.87 + 0.02).clip(0, 1) arr[...,2] = (arr[...,2] * 1.35 + 0.08).clip(0, 1) arr = cls._gamma(arr, 1.30) lum = arr.mean(axis=-1, keepdims=True) mid = np.exp(-((lum - 0.5)**2) / (2*0.18**2)) hsv = cls._rgb_to_hsv(arr) hsv[...,1] = (hsv[...,1] * (1 - mid[...,0] * 0.55)).clip(0, 1) arr = cls._hsv_to_rgb(hsv) arr = (arr - 0.5) * 1.45 + 0.5 arr = cls._vignette(arr, strength=0.60) arr = cls._grain(arr, amount=0.038, seed=42) return cls._pil(arr) @classmethod def style_7_sketch_artistic(cls, img): arr = cls._np(img) steps = 5 arr_post = np.floor(arr * steps) / steps lum = arr.mean(axis=-1) edges = np.abs(lum - gaussian_filter(lum, sigma=1.2)) * 10.0 edges = edges.clip(0, 1)[..., None] arr_e = arr_post * (1 - edges * 0.88) arr_e = cls._saturate(arr_e, 1.9) canvas = np.array([0.95, 0.92, 0.85], dtype=np.float32) arr_e = arr_e * 0.86 + canvas * 0.14 rng = np.random.default_rng(seed=99) grain = rng.uniform(-0.03, 0.03, arr.shape).astype(np.float32) arr_e = (arr_e + grain).clip(0, 1) pil = cls._pil(arr_e) return pil.filter(ImageFilter.UnsharpMask(radius=1, percent=130, threshold=2)) TRANSFORMS = [ style_0_youthful.__func__, style_1_aged_mature.__func__, style_2_dramatic_light.__func__, style_3_soft_glow.__func__, style_4_intense_bold.__func__, style_5_warm_golden.__func__, style_6_cool_moody.__func__, style_7_sketch_artistic.__func__, ] @classmethod def apply(cls, img, style_idx): return cls.TRANSFORMS[style_idx](cls, img) @classmethod def apply_all(cls, img): return [cls.apply(img, i) for i in range(len(cls.TRANSFORMS))] # ───────────────────────────────────────────────────────────────────────────── # MODEL LOADING # ───────────────────────────────────────────────────────────────────────────── _vae = None _sr = None _cfg = None def load_model(path: str): global _vae, _sr, _cfg ckpt = torch.load(path, map_location=DEVICE) c = ckpt['config'] vae = FaceVAE(c['latent_dim'], c['base_filters']).to(DEVICE) sr = SRHead(n_blocks=10, ch=64).to(DEVICE) vae.load_state_dict(ckpt['vae_state']) sr.load_state_dict(ckpt['sr_state']) vae.eval(); sr.eval() _vae = vae; _sr = sr; _cfg = c return ( f"✅ Model loaded from: {os.path.basename(path)}\n" f" FaceVAE : {sum(p.numel() for p in vae.parameters())/1e6:.2f}M params\n" f" SRHead : {sum(p.numel() for p in sr.parameters())/1e6:.2f}M params\n" f" Pipeline : {c['image_size']}×{c['image_size']} → {c['sr_size']}×{c['sr_size']} HD\n" f" Device : {DEVICE}" ) # ───────────────────────────────────────────────────────────────────────────── # INFERENCE HELPERS # ───────────────────────────────────────────────────────────────────────────── def preprocess(pil_img: Image.Image, size: int) -> torch.Tensor: t = TF.to_tensor(TF.resize(pil_img.convert("RGB"), (size, size))) return TF.normalize(t, [0.5]*3, [0.5]*3).unsqueeze(0).to(DEVICE) def tensor_to_pil(t: torch.Tensor) -> Image.Image: t = t.squeeze(0).cpu().clamp(-1, 1) return transforms.ToPILImage()((t + 1.0) / 2.0) def apply_with_strength(hd_pil, style_idx, alpha): styled = StyleTransformBank.apply(hd_pil, style_idx) rec_arr = np.array(hd_pil, dtype=np.float32) sty_arr = np.array(styled, dtype=np.float32) blended = rec_arr * (1 - alpha) + sty_arr * alpha return Image.fromarray(blended.clip(0, 255).astype(np.uint8)) # ───────────────────────────────────────────────────────────────────────────── # MAIN INFERENCE FUNCTION # ───────────────────────────────────────────────────────────────────────────── def run_inference(image, style_strength): if _vae is None or _sr is None: raise gr.Error("⚠️ Please load a model first using the Model Path field above.") if image is None: raise gr.Error("⚠️ Please upload an image.") lr_size = _cfg.get("image_size", 256) hr_size = _cfg.get("sr_size", 512) pil_input = Image.fromarray(image).convert("RGB") x_lr = preprocess(pil_input, lr_size) with torch.no_grad(): recon_256_tensor = _vae.reconstruct(x_lr) recon_512_tensor = _sr(recon_256_tensor) recon_256_pil = tensor_to_pil(recon_256_tensor) recon_512_pil = tensor_to_pil(recon_512_tensor) # Apply strength-blended styles alpha = style_strength / 100.0 styled_pils = [ apply_with_strength(recon_512_pil, i, alpha) for i in range(len(STYLE_NAMES)) ] orig_hr = pil_input.resize((hr_size, hr_size), Image.LANCZOS) return ( np.array(orig_hr), np.array(recon_256_pil), np.array(recon_512_pil), *[np.array(s) for s in styled_pils], ) def latent_walk(image, style_a_idx, style_b_idx, steps): if _vae is None or _sr is None: raise gr.Error("⚠️ Please load a model first.") if image is None: raise gr.Error("⚠️ Please upload an image.") lr_size = _cfg.get("image_size", 256) pil_input = Image.fromarray(image).convert("RGB") x_lr = preprocess(pil_input, lr_size) steps = int(steps) walk_frames = [] with torch.no_grad(): mu, logvar, skips = _vae.encode(x_lr) direction = F.normalize(torch.randn_like(mu), dim=-1) for i, t_val in enumerate(torch.linspace(-2.5, 2.5, steps)): z_walk = mu + t_val.item() * direction rec_256 = _vae.decode(z_walk, skips) rec_512 = _sr(rec_256) rec_pil = tensor_to_pil(rec_512) alpha = i / (steps - 1) sa = np.array(StyleTransformBank.apply(rec_pil, style_a_idx), dtype=np.float32) sb = np.array(StyleTransformBank.apply(rec_pil, style_b_idx), dtype=np.float32) blended = Image.fromarray((sa*(1-alpha) + sb*alpha).clip(0,255).astype(np.uint8)) walk_frames.append(np.array(blended)) return walk_frames # ───────────────────────────────────────────────────────────────────────────── # CUSTOM CSS — Warm Ivory · Premium Editorial Light Theme # ───────────────────────────────────────────────────────────────────────────── CUSTOM_CSS = """ @import url('https://fonts.googleapis.com/css2?family=DM+Sans:ital,opsz,wght@0,9..40,300;0,9..40,400;0,9..40,500;0,9..40,600;0,9..40,700;1,9..40,300&family=DM+Serif+Display:ital@0;1&family=DM+Mono:wght@400;500&display=swap'); /* ── Design tokens ───────────────────────────────────────────────── */ :root { --cream: #faf7f2; --parchment: #f3ede3; --warm-white: #fefcf8; --sand: #e8dece; --sand-dark: #d4c8b8; --ink: #2c2520; --ink-mid: #5a4f47; --ink-light: #8c7f75; --ink-faint: #b8ada5; --terracotta: #c4714a; --terracotta-l:#e8956d; --sage: #6b8f71; --sage-l: #9ab89f; --gold: #b8943a; --gold-l: #d4b05a; --dusty-rose: #c47a7a; --slate-blue: #6a7d96; --radius: 10px; --radius-lg: 16px; --shadow-soft: 0 2px 16px rgba(44,37,32,0.08); --shadow-card: 0 4px 24px rgba(44,37,32,0.10); --border: #ddd5c8; --border-light:#e8e0d6; } /* ── Base ────────────────────────────────────────────────────────── */ body, .gradio-container { background: var(--cream) !important; font-family: 'DM Sans', sans-serif !important; color: var(--ink) !important; } /* Subtle linen texture overlay */ .gradio-container::before { content: ''; position: fixed; inset: 0; background-image: url("data:image/svg+xml,%3Csvg xmlns='http://www.w3.org/2000/svg' width='300' height='300'%3E%3Cfilter id='n'%3E%3CfeTurbulence type='fractalNoise' baseFrequency='0.75' numOctaves='4' stitchTiles='stitch'/%3E%3C/filter%3E%3Crect width='300' height='300' filter='url(%23n)' opacity='0.025'/%3E%3C/svg%3E"); pointer-events: none; z-index: 0; opacity: 0.6; } /* ── Hero ────────────────────────────────────────────────────────── */ .shuk-hero { text-align: center; padding: 56px 24px 40px; background: linear-gradient(160deg, var(--warm-white) 0%, var(--parchment) 100%); border-bottom: 1px solid var(--border); position: relative; overflow: hidden; } /* Warm glow blob behind title */ .shuk-hero::before { content: ''; position: absolute; top: -40px; left: 50%; transform: translateX(-50%); width: 500px; height: 260px; background: radial-gradient(ellipse, rgba(196,113,74,0.10) 0%, rgba(184,148,58,0.06) 50%, transparent 80%); pointer-events: none; animation: warmPulse 6s ease-in-out infinite; } @keyframes warmPulse { 0%, 100% { opacity: 0.7; transform: translateX(-50%) scale(1); } 50% { opacity: 1; transform: translateX(-50%) scale(1.08); } } .shuk-hero h1 { font-family: 'DM Serif Display', serif !important; font-size: clamp(2.6rem, 5.5vw, 4rem) !important; font-weight: 400 !important; font-style: italic; color: var(--ink) !important; letter-spacing: -0.5px; line-height: 1.1; margin: 0 0 10px !important; } .shuk-hero h1 span { color: var(--terracotta) !important; } .shuk-hero .tagline { font-size: 0.95rem; color: var(--ink-light); letter-spacing: 0.02em; font-weight: 400; margin-top: 4px; } .shuk-hero .badge-row { display: flex; gap: 8px; justify-content: center; flex-wrap: wrap; margin-top: 20px; } .badge { background: var(--warm-white); border: 1px solid var(--sand-dark); border-radius: 100px; padding: 4px 13px; font-size: 0.72rem; font-family: 'DM Mono', monospace; color: var(--ink-mid); letter-spacing: 0.04em; box-shadow: 0 1px 4px rgba(44,37,32,0.06); } /* ── Pipeline bar ────────────────────────────────────────────────── */ .pipeline-bar { display: flex; align-items: center; justify-content: center; flex-wrap: wrap; gap: 4px; padding: 18px 0 0; } .pipe-step { background: var(--warm-white); border: 1px solid var(--sand-dark); border-radius: 7px; padding: 5px 12px; font-size: 0.70rem; font-family: 'DM Mono', monospace; color: var(--ink-mid); white-space: nowrap; box-shadow: 0 1px 3px rgba(44,37,32,0.06); } .pipe-arrow { color: var(--ink-faint); font-size: 1rem; margin: 0 1px; } /* ── Section labels ──────────────────────────────────────────────── */ .section-label { font-family: 'DM Sans', sans-serif; font-size: 0.68rem; font-weight: 600; letter-spacing: 0.13em; text-transform: uppercase; color: var(--ink-light); margin-bottom: 10px; display: flex; align-items: center; gap: 10px; } .section-label::after { content: ''; flex: 1; height: 1px; background: var(--border-light); } /* ── Buttons ─────────────────────────────────────────────────────── */ button.primary, .gr-button-primary { background: var(--terracotta) !important; border: none !important; border-radius: var(--radius) !important; color: #fff !important; font-family: 'DM Sans', sans-serif !important; font-weight: 600 !important; font-size: 0.92rem !important; padding: 11px 26px !important; cursor: pointer !important; transition: background 0.2s, transform 0.15s, box-shadow 0.2s !important; box-shadow: 0 3px 14px rgba(196,113,74,0.35) !important; letter-spacing: 0.01em; } button.primary:hover { background: var(--terracotta-l) !important; transform: translateY(-1px) !important; box-shadow: 0 6px 22px rgba(196,113,74,0.45) !important; } button.primary:active { transform: translateY(0) !important; } button.secondary, .gr-button-secondary { background: var(--warm-white) !important; border: 1px solid var(--border) !important; border-radius: var(--radius) !important; color: var(--ink-mid) !important; font-family: 'DM Sans', sans-serif !important; font-weight: 500 !important; transition: border-color 0.2s, color 0.2s, box-shadow 0.2s !important; } button.secondary:hover { border-color: var(--terracotta) !important; color: var(--terracotta) !important; box-shadow: 0 2px 10px rgba(196,113,74,0.15) !important; } /* ── Inputs ──────────────────────────────────────────────────────── */ input[type="text"], textarea, .gr-textbox input { background: var(--warm-white) !important; border: 1px solid var(--border) !important; border-radius: var(--radius) !important; color: var(--ink) !important; font-family: 'DM Mono', monospace !important; font-size: 0.83rem !important; padding: 10px 14px !important; transition: border-color 0.2s, box-shadow 0.2s !important; } input[type="text"]:focus, textarea:focus { border-color: var(--terracotta) !important; box-shadow: 0 0 0 3px rgba(196,113,74,0.12) !important; outline: none !important; } /* ── Textbox status ──────────────────────────────────────────────── */ .gr-textbox textarea { background: var(--warm-white) !important; color: var(--sage) !important; border: 1px solid var(--border) !important; font-family: 'DM Mono', monospace !important; font-size: 0.8rem !important; line-height: 1.7 !important; border-radius: var(--radius) !important; } /* ── Sliders ─────────────────────────────────────────────────────── */ input[type="range"] { accent-color: var(--terracotta) !important; } /* ── Block / card overrides ──────────────────────────────────────── */ .block, .gr-block { background: var(--warm-white) !important; border: 1px solid var(--border-light) !important; border-radius: var(--radius-lg) !important; box-shadow: var(--shadow-soft) !important; } /* ── Image component ─────────────────────────────────────────────── */ .gr-image, [data-testid="image"] { border-radius: var(--radius) !important; border: 1px solid var(--border) !important; background: var(--parchment) !important; overflow: hidden; } /* ── Tab nav ─────────────────────────────────────────────────────── */ .tab-nav { border-bottom: 1px solid var(--border) !important; background: var(--parchment) !important; } .tab-nav button { font-family: 'DM Sans', sans-serif !important; font-weight: 500 !important; font-size: 0.88rem !important; color: var(--ink-light) !important; border: none !important; background: transparent !important; padding: 11px 20px !important; transition: color 0.2s !important; border-bottom: 2px solid transparent !important; } .tab-nav button.selected, .tab-nav button[aria-selected="true"] { color: var(--terracotta) !important; border-bottom: 2px solid var(--terracotta) !important; background: rgba(196,113,74,0.04) !important; font-weight: 600 !important; } .tab-nav button:hover { color: var(--ink-mid) !important; } /* ── Labels ──────────────────────────────────────────────────────── */ label, .gr-label, span.svelte-1gfkn6j { font-family: 'DM Sans', sans-serif !important; font-size: 0.82rem !important; font-weight: 500 !important; color: var(--ink-mid) !important; letter-spacing: 0.01em; } /* ── Dropdown ────────────────────────────────────────────────────── */ .gr-dropdown select, select { background: var(--warm-white) !important; border: 1px solid var(--border) !important; border-radius: var(--radius) !important; color: var(--ink) !important; font-family: 'DM Sans', sans-serif !important; } /* ── Scrollbar ───────────────────────────────────────────────────── */ ::-webkit-scrollbar { width: 6px; height: 6px; } ::-webkit-scrollbar-track { background: var(--parchment); } ::-webkit-scrollbar-thumb { background: var(--sand-dark); border-radius: 4px; } ::-webkit-scrollbar-thumb:hover { background: var(--terracotta); } /* ── Divider ─────────────────────────────────────────────────────── */ .divider { border: none; border-top: 1px solid var(--border-light); margin: 8px 0; } /* ── Footer ──────────────────────────────────────────────────────── */ .footer { text-align: center; padding: 28px; color: #3ba729 !important; font-size: 0.88rem; border-top: 1px solid var(--border-light); margin-top: 40px; font-family: 'DM Mono', monospace; letter-spacing: 0.03em; background: var(--parchment); } .footer span { color: var(--terracotta); font-style: italic; } /* ── Gallery ─────────────────────────────────────────────────────── */ .gr-gallery { background: var(--parchment) !important; border-radius: var(--radius) !important; border: 1px solid var(--border) !important; } /* ── Accordion ───────────────────────────────────────────────────── */ details, .gr-accordion { background: var(--warm-white) !important; border: 1px solid var(--border-light) !important; border-radius: var(--radius) !important; } .load-model-btn{ margin-top: 24px !important; } """ # ───────────────────────────────────────────────────────────────────────────── # GRADIO UI # ───────────────────────────────────────────────────────────────────────────── def build_ui(): with gr.Blocks( css=CUSTOM_CSS, theme=gr.themes.Base( primary_hue="orange", secondary_hue="yellow", neutral_hue="stone", font=["DM Sans", "sans-serif"], ).set( body_background_fill="#faf7f2", body_text_color="#2c2520", block_background_fill="#fefcf8", block_border_color="#ddd5c8", block_label_text_color="#8c7f75", input_background_fill="#fefcf8", input_border_color="#ddd5c8", button_primary_background_fill="#c4714a", button_primary_text_color="#ffffff", button_secondary_background_fill="#fefcf8", button_secondary_border_color="#ddd5c8", button_secondary_text_color="#5a4f47", ), title="ShukGEN Advanced", ) as demo: # ── HERO ───────────────────────────────────────────────────────────── gr.HTML("""

ShukGEN Advanced

Identity-Preserving Face Reconstruction  ·  HD Super Resolution  ·  8 Style Variants

FaceVAE SRHead 512px EMA Weights Identity Loss FFT Sharpness 8 Style Transforms
Input 256×256 → FaceVAE Encode → z · latent 512d → FaceVAE Decode → SRHead 512×512 → StyleTransformBank → 8 HD Variants ✦
""") # ── MODEL LOADER ────────────────────────────────────────────────────── gr.HTML('
Model Configuration
') with gr.Row(): with gr.Column(scale=4): model_path = gr.Textbox( label="Model Path", placeholder="e.g. shukgen_v4_final.pth or checkpoints_shukgen_v4/best.pth", value="shukgen_v4_final.pth", ) with gr.Column(scale=1): load_btn = gr.Button("Load Model", variant="primary", elem_classes="load-model-btn") model_status = gr.Textbox( label="Status", interactive=False, placeholder="Load a model to get started...", lines=4, ) load_btn.click(fn=load_model, inputs=model_path, outputs=model_status) gr.HTML('
') # ── TABS ────────────────────────────────────────────────────────────── with gr.Tabs(): # ── TAB 1: Main Inference ───────────────────────────────────────── with gr.TabItem(" Reconstruct & Style "): gr.HTML('
Input & Controls
') with gr.Row(): with gr.Column(scale=1): input_img = gr.Image( label="Upload Face Image", type="numpy", height=310, image_mode="RGB", ) style_strength = gr.Slider( minimum=0, maximum=100, value=100, step=1, label="Style Strength (%)", info="0 = pure HD reconstruction · 100 = full style applied", ) run_btn = gr.Button("Generate All Outputs ✦", variant="primary", size="lg") with gr.Column(scale=2): gr.HTML('
Core Outputs
') with gr.Row(): out_orig = gr.Image(label="Original (HR Reference)", height=200) out_256 = gr.Image(label="VAE Recon · 256×256", height=200) out_512 = gr.Image(label="HD Recon ✦ · 512×512", height=200) gr.HTML('
8 Style Variants · HD 512×512
') style_label_colors = [ ("✦ Youthful", "#6b8f71"), ("⌛ Aged / Mature", "#8c6a4a"), ("◑ Dramatic Light", "#6a5a8c"), ("◌ Soft Glow", "#b8943a"), ("◈ Intense / Bold", "#c4714a"), ("☀ Warm Golden Hour", "#b87c3a"), ("◐ Cool / Moody", "#4a7a8c"), ("✏ Sketch / Artistic", "#4a7a5a"), ] style_outputs = [] for row in range(2): with gr.Row(): for col in range(4): idx = row * 4 + col label, color = style_label_colors[idx] img = gr.Image( label=label, height=175, elem_id=f"style_{idx}", ) style_outputs.append(img) run_btn.click( fn=run_inference, inputs=[input_img, style_strength], outputs=[out_orig, out_256, out_512] + style_outputs, ) # ── TAB 2: Latent Walk ──────────────────────────────────────────── with gr.TabItem(" Latent Space Walk "): gr.HTML("""
Explore the latent space of your image. The model encodes the face to a 512-d vector, then walks along a random direction while gradually blending between two chosen styles.
""") with gr.Row(): walk_input = gr.Image(label="Input Image", type="numpy", height=270) with gr.Column(): style_a = gr.Dropdown( choices=[(f"{STYLE_EMOJIS[i]} {n}", i) for i, n in enumerate(STYLE_NAMES)], value=2, label="Style A (start)", ) style_b = gr.Dropdown( choices=[(f"{STYLE_EMOJIS[i]} {n}", i) for i, n in enumerate(STYLE_NAMES)], value=6, label="Style B (end)", ) walk_steps = gr.Slider(4, 12, value=8, step=1, label="Walk Steps") walk_btn = gr.Button("Generate Latent Walk", variant="primary") gr.HTML('
Walk Frames · Style A → Style B
') walk_gallery = gr.Gallery( label="", columns=8, height=220, ) walk_btn.click( fn=latent_walk, inputs=[walk_input, style_a, style_b, walk_steps], outputs=walk_gallery, ) # ── TAB 3: Style Strength Explorer ─────────────────────────────── with gr.TabItem(" Style Strength "): gr.HTML("""
Choose a single style and preview how it blends from pure HD reconstruction (0%) to full effect (100%) across five evenly-spaced steps.
""") with gr.Row(): exp_input = gr.Image(label="Input Image", type="numpy", height=270) with gr.Column(): exp_style = gr.Dropdown( choices=[(f"{STYLE_EMOJIS[i]} {n}", i) for i, n in enumerate(STYLE_NAMES)], value=0, label="Choose Style", ) exp_btn = gr.Button("Explore Strength", variant="primary") gr.HTML('
Blend Steps · 0% → 25% → 50% → 75% → 100%
') exp_outs = [] with gr.Row(): for a in [0, 25, 50, 75, 100]: img = gr.Image(label=f"α = {a}%", height=195) exp_outs.append(img) def explore_strength(image, style_idx): if _vae is None or _sr is None: raise gr.Error("Load model first.") if image is None: raise gr.Error("Upload an image first.") lr_size = _cfg.get("image_size", 256) pil_input = Image.fromarray(image).convert("RGB") x_lr = preprocess(pil_input, lr_size) with torch.no_grad(): recon_256 = _vae.reconstruct(x_lr) recon_512 = _sr(recon_256) hd_pil = tensor_to_pil(recon_512) results = [] for a in [0.0, 0.25, 0.50, 0.75, 1.0]: results.append(np.array(apply_with_strength(hd_pil, int(style_idx), a))) return results exp_btn.click(fn=explore_strength, inputs=[exp_input, exp_style], outputs=exp_outs) # ── TAB 4: Architecture ─────────────────────────────────────────── with gr.TabItem(" Architecture "): gr.HTML("""

ShukGEN Advanced Architecture

Identity-Preserving VAE · Super Resolution · Deterministic Style Transforms

Component What It Does
FaceVAE U-Net VAE with multi-head attention bottleneck. Encodes face → 512-d latent → reconstructs at 256×256.
Identity Loss InceptionResNetV2 (timm) measures cosine similarity of face embeddings — preserves who the person is.
FFT Loss Penalises missing high-frequency detail — hair strands, skin pores, sharp edges.
SSIM Loss Structural Similarity Index — evaluates quality the way human eyes do.
SRHead SRRESNET-style: 10 residual blocks + PixelShuffle 2× — 256×256 → 512×512 HD, no checkerboard artifacts.
EMA Weights Exponential Moving Average (decay 0.999) — inference uses averaged weights for smoother, sharper outputs.
StyleTransformBank 8 deterministic pixel-space transforms on HD output: Youthful, Aged, Dramatic Light, Soft Glow, Intense, Golden Hour, Cool/Moody, Sketch.

Training Losses

VAE Losses · 256×256
L1 Pixel  ·  Perceptual VGG (4 layers)
SSIM  ·  Identity (FaceNet)  ·  KL Divergence
SR Losses · 512×512
L1 Pixel  ·  FFT Sharpness  ·  Perceptual VGG
""") # ── FOOTER ──────────────────────────────────────────────────────────── gr.HTML("""
ShukGEN Advanced  ·  Identity-Preserving HD Face Reconstruction  ·  VAE + SRHead + StyleTransformBank  ·  PyTorch & Gradio
""") return demo # ───────────────────────────────────────────────────────────────────────────── # ENTRY POINT # ───────────────────────────────────────────────────────────────────────────── if __name__ == "__main__": demo = build_ui() demo.launch( share=True )