Upload anime_t2i_tpu.py with huggingface_hub

anime_t2i_tpu.py

@@ -0,0 +1,739 @@
+#!/usr/bin/env python
+"""
+Anime yüzü text-to-image (pixel-space diffusion) - Google TPU v5e-1 (PyTorch/XLA) sürümü.
+
+Mimari : diffusers.UNet2DConditionModel (~17-18M parametre, sıfırdan eğitiliyor)
+Metin  : CLIP ViT-B/32 text encoder (DONDURULMUŞ; embedding'ler önceden hesaplanır)
+Veri   : huggan/anime-faces -> CLIP zero-shot ile saç rengi / saç uzunluğu / göz rengi / gülümseme
+         etiketleri -> "anime face, blue long hair, red eyes, smiling"
+
+GPU sürümüne göre TPU için yapılan değişiklikler
+  * Accelerate / DDP / GradScaler / fp16 yok. TPU v5e-1 tek çip -> tek süreç.
+    Hesap bf16 (torch.autocast("xla")), master ağırlıklar + EMA + optimizer durumu FP32.
+    bf16'nın aralığı fp32 ile aynı olduğu için "FP32 attention" hack'ine ve loss-scaling'e gerek kalmadı.
+  * Tüm şekiller statik (batch, çözünürlük, token sayısı sabit) -> XLA grafiği 1 kez derlenir.
+  * Adım içinde HİÇ host<->device senkronizasyonu yok (.item(), python if, isfinite kontrolü vs. yok).
+    lr, warmup, cosine, Adam adım sayacı, EMA decay hepsi cihaz üzerinde TENSÖR olarak hesaplanır
+    (python float lr kullanmak XLA'da her adımda yeniden derlemeye yol açar).
+  * Özel AdamW + grad-clip + NaN/inf koruması (bozuk adım cihaz üzerinde atlanır, host'a hiç dönülmez).
+  * Bellek-verimli (query-chunked, flash-benzeri, özel backward) self-attention:
+    XLA'da SDPA tam (B,H,S,S) skor matrisini HBM'de tutar; 32x32 seviyesinde (S=1024) batch 128+ OOM olurdu.
+  * Veri seti (uint8, NCHW) tamamen HBM'de; batch örnekleme/flip/caption seçimi cihazda -> DataLoader yok.
+  * DDIM örnekleyici elle yazıldı: her adımın grafiği aynı (t ve alpha'lar girdi tensörü) -> 50 adım TEK derleme.
+  * Kalıcı XLA derleme önbelleği (xla_cache/): ikinci çalıştırmada / resume'da derleme süresi ~0.
+  * prepare: CLIP görsel encoder'ı da TPU'da (sabit batch=256, son batch padlenir).
+  * Checkpoint'ler CPU'ya alınıp yazılır; torchvision bağımlılığı kaldırıldı (grid'i PIL ile yazıyoruz).
+
+KULLANIM (Colab: Runtime > TPU v5e-1)
+    !pip -q install diffusers transformers                # torch_xla hazır değilse aşağıya bak
+    !python anime_t2i_tpu.py prepare
+    !python anime_t2i_tpu.py train --steps 40000 --batch 128 --lr 1e-4
+    !python anime_t2i_tpu.py sample --prompt "anime face, blue long hair, red eyes, smiling"
+
+    torch_xla yoksa:  !pip -q install torch==2.8.0 'torch_xla[tpu]==2.8.0'
+    Kalıcı klasör  :  WORK_DIR=/content/drive/MyDrive/anime_t2i  (checkpoint'ler Drive'a yazılır)
+    Devam etme     :  ... train --resume <WORK>/ckpt/train_state_4000.pt --lr 5e-5
+                      (eski GPU checkpoint'leri de yüklenir; optimizer durumu sıfırdan başlar)
+    Ayar önerileri :  --batch 256 (HBM izin veriyorsa; lr'yi 1.5e-4 - 2e-4 yap), --base_ch 128 (MXU dostu, ~4x param)
+    Test (TPU'suz) :  FORCE_CPU=1 python anime_t2i_tpu.py train --steps 5 --batch 8 --base_ch 32 ...
+"""
+import os
+import time
+import math
+import argparse
+import zipfile
+import contextlib
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+# ---------------------------------------------------------------- XLA (opsiyonel)
+HAS_XLA = False
+if os.environ.get("FORCE_CPU", "0") != "1":
+    try:
+        import torch_xla
+        import torch_xla.core.xla_model as xm
+        HAS_XLA = True
+    except Exception:
+        HAS_XLA = False
+
+# ------------------------------------------------------------------ AYARLAR
+RES = 64                 # çözünürlük (4'e bölünebilmeli)
+T_STEPS = 1000           # difüzyon adım sayısı
+TOK_LEN = 24             # caption token uzunluğu
+CLIP_ID = "openai/clip-vit-base-patch32"
+P_UNCOND = 0.1           # classifier-free guidance için caption'ı düşürme olasılığı
+MIN_SNR_GAMMA = 5.0      # Min-SNR loss ağırlığı
+SEED = 42
+WARMUP = 1000
+LOG_EVERY = 100
+ATTN_BUDGET = 2 ** 27    # chunked attention: bir chunk'taki skor elemanı sayısı üst sınırı (fp32 ~512MB)
+
+if os.environ.get("WORK_DIR"):
+    WORK = Path(os.environ["WORK_DIR"])
+elif Path("/kaggle/working").exists():
+    WORK = Path("/kaggle/working")
+elif Path("/content").exists():
+    WORK = Path("/content/work")
+else:
+    WORK = Path("./work")
+DATA = WORK / "data"
+CKPT = WORK / "ckpt"
+OUT = WORK / "samples"
+
+HAIR_COLORS = ["blonde", "brown", "black", "blue", "pink", "purple", "green", "red", "silver", "white", "orange"]
+EYE_COLORS = ["blue", "red", "green", "brown", "purple", "yellow", "orange", "pink", "black"]
+HAIR_LEN = ["long", "short"]
+
+
+# ------------------------------------------------------------------ CİHAZ YARDIMCILARI
+def init_xla_cache():
+    """Kalıcı derleme önbelleği: cihaz kullanılmadan ÖNCE çağrılmalı."""
+    if not HAS_XLA:
+        return
+    try:
+        import torch_xla.runtime as xr
+        (WORK / "xla_cache").mkdir(parents=True, exist_ok=True)
+        xr.initialize_cache(str(WORK / "xla_cache"), readonly=False)
+    except Exception as e:
+        print("XLA derleme önbelleği açılamadı (sorun değil):", e)
+
+
+def get_device():
+    if HAS_XLA:
+        os.environ.setdefault("PJRT_DEVICE", "TPU")
+        try:
+            return torch_xla.device()
+        except AttributeError:
+            return xm.xla_device()
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def sync():
+    """Lazy grafiği cihaza gönder (adım sınırı). XLA dışında no-op."""
+    if HAS_XLA:
+        if hasattr(torch_xla, "sync"):
+            torch_xla.sync()
+        else:
+            xm.mark_step()
+
+
+def amp(dev):
+    """bf16 autocast (TPU'da doğal). CPU'da kapalı (test için fp32)."""
+    if dev.type == "xla":
+        try:
+            return torch.autocast("xla", dtype=torch.bfloat16)
+        except Exception:
+            from torch_xla.amp import autocast as xla_autocast
+            return xla_autocast(dev, dtype=torch.bfloat16)
+    if dev.type == "cuda" and torch.cuda.is_bf16_supported():
+        return torch.autocast("cuda", dtype=torch.bfloat16)
+    return contextlib.nullcontext()
+
+
+def seed_all(seed):
+    torch.manual_seed(seed)
+    np.random.seed(seed % (2 ** 32))
+    if HAS_XLA:
+        try:
+            xm.set_rng_state(seed)
+        except Exception:
+            pass
+
+
+# ------------------------------------------------------------------ BELLEK-VERİMLİ ATTENTION
+class _ChunkedAttnFn(torch.autograd.Function):
+    """
+    softmax(QK^T * scale) V  — query'leri parçalara bölerek hesaplar, (B,H,S,S) matrisi asla tamamen tutulmaz.
+    Backward'da skorlar yeniden hesaplanır (FlashAttention mantığı; yalnızca q,k,v,o,logsumexp saklanır).
+    softmax/exp fp32, matmul'lar girdi dtype'ında (bf16) yapılır.
+    """
+
+    @staticmethod
+    def forward(ctx, q, k, v, scale, chunk):
+        outs, lses = [], []
+        kt = k.transpose(-1, -2)
+        for i in range(0, q.shape[2], chunk):
+            s = torch.matmul(q[:, :, i:i + chunk], kt).float() * scale
+            lse = torch.logsumexp(s, dim=-1, keepdim=True)
+            p = torch.exp(s - lse).to(v.dtype)
+            outs.append(torch.matmul(p, v))
+            lses.append(lse)
+        o = torch.cat(outs, dim=2)
+        lse = torch.cat(lses, dim=2)
+        ctx.save_for_backward(q, k, v, o, lse)
+        ctx.scale, ctx.chunk = scale, chunk
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        q, k, v, o, lse = ctx.saved_tensors
+        scale, chunk = ctx.scale, ctx.chunk
+        do = do.to(q.dtype)
+        delta = (do.float() * o.float()).sum(-1, keepdim=True)          # (B,H,S,1)
+        kt, vt = k.transpose(-1, -2), v.transpose(-1, -2)
+        dk = torch.zeros(k.shape, dtype=torch.float32, device=k.device)
+        dv = torch.zeros(v.shape, dtype=torch.float32, device=v.device)
+        dqs = []
+        for i in range(0, q.shape[2], chunk):
+            qc, doc = q[:, :, i:i + chunk], do[:, :, i:i + chunk]
+            s = torch.matmul(qc, kt).float() * scale
+            p = torch.exp(s - lse[:, :, i:i + chunk])
+            dp = torch.matmul(doc, vt).float()
+            ds = (p * (dp - delta[:, :, i:i + chunk]) * scale).to(q.dtype)
+            dqs.append(torch.matmul(ds, k))
+            dk = dk + torch.matmul(ds.transpose(-1, -2), qc).float()
+            dv = dv + torch.matmul(p.to(q.dtype).transpose(-1, -2), doc).float()
+        return torch.cat(dqs, dim=2), dk.to(k.dtype), dv.to(v.dtype), None, None
+
+
+def chunked_attention(q, k, v, scale, budget=ATTN_BUDGET):
+    """q,k,v: (B,H,S,D)."""
+    B, H, S, _ = q.shape
+    c = max(16, min(S, budget // (B * H * k.shape[2])))
+    c = 1 << (int(c).bit_length() - 1)            # 2'nin kuvveti (S de 2'nin kuvveti -> tam bölünür)
+    return _ChunkedAttnFn.apply(q, k, v, scale, min(c, S))
+
+
+class ChunkedAttnProcessor:
+    """diffusers Attention için processor: self-attn -> chunked, cross-attn (K uzunluğu 24) -> SDPA."""
+
+    def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, temb=None,
+                 *args, **kwargs):
+        if attention_mask is not None or getattr(attn, "group_norm", None) is not None \
+                or getattr(attn, "spatial_norm", None) is not None:
+            raise NotImplementedError("ChunkedAttnProcessor bu attention yapılandırmasını desteklemiyor; --attn sdpa kullan.")
+        is_cross = encoder_hidden_states is not None
+        ctx = encoder_hidden_states if is_cross else hidden_states
+        if is_cross and getattr(attn, "norm_cross", None):
+            ctx = attn.norm_encoder_hidden_states(ctx)
+        B = hidden_states.shape[0]
+        q, k, v = attn.to_q(hidden_states), attn.to_k(ctx), attn.to_v(ctx)
+        H = attn.heads
+        D = q.shape[-1] // H
+        q = q.reshape(B, -1, H, D).transpose(1, 2)
+        k = k.reshape(B, -1, H, D).transpose(1, 2)
+        v = v.reshape(B, -1, H, D).transpose(1, 2)
+        if is_cross:
+            o = F.scaled_dot_product_attention(q, k, v)
+        else:
+            o = chunked_attention(q, k, v, attn.scale)
+        o = o.transpose(1, 2).reshape(B, -1, H * D).to(q.dtype)
+        o = attn.to_out[0](o)
+        return attn.to_out[1](o)
+
+
+def set_attn(unet, mode):
+    if mode == "chunked":
+        unet.set_attn_processor(ChunkedAttnProcessor())
+    return unet
+
+
+# ------------------------------------------------------------------ MODEL
+def build_unet(base_ch=64, attn="chunked"):
+    """Hazır diffusers UNet'i, küçük config ile (base_ch=64 -> ~17-18M param)."""
+    from diffusers import UNet2DConditionModel
+    unet = UNet2DConditionModel(
+        sample_size=RES,
+        in_channels=3,
+        out_channels=3,
+        down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D"),
+        up_block_types=("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D"),
+        block_out_channels=(96, 192, 384),   # base_ch=128 -> MXU (128x128) dostu
+        layers_per_block=1,
+        cross_attention_dim=512,             # CLIP ViT-B/32 text hidden size
+        attention_head_dim=8,
+        norm_num_groups=32,
+    )
+    torch.nn.init.zeros_(unet.conv_out.weight)
+    torch.nn.init.zeros_(unet.conv_out.bias)
+    return set_attn(unet, attn)
+
+
+class EMA:
+    """Cihaz üzerinde EMA. decay warmup'ı tensör olarak hesaplanır -> derleme tekrarı yok."""
+
+    def __init__(self, params, dev, decay=0.9995):
+        self.shadow = [p.detach().clone() for p in params]
+        self.decay = decay
+        self.n = torch.tensor(0.0, device=dev)
+
+    @torch.no_grad()
+    def update(self, params):
+        self.n += 1.0
+        d = torch.clamp((1.0 + self.n) / (10.0 + self.n), max=self.decay)
+        w = 1.0 - d
+        for s, p in zip(self.shadow, params):
+            s.add_((p.detach() - s) * w)
+
+    @torch.no_grad()
+    def copy_to(self, model):
+        for p, s in zip(model.parameters(), self.shadow):
+            p.copy_(s)
+
+
+class DeviceAdamW:
+    """
+    AdamW (decoupled wd). lr ve adım sayacı cihaz tensörüdür; `ok` False ise (grad NaN/inf)
+    m, v, p ve sayaç güncellenmez (adım tamamen atlanır) — host'a dönmeden.
+    """
+
+    def __init__(self, params, dev, betas=(0.9, 0.99), eps=1e-6, weight_decay=0.01):
+        self.params = params
+        self.b1, self.b2 = betas
+        self.eps, self.wd = eps, weight_decay
+        self.m = [torch.zeros_like(p) for p in params]
+        self.v = [torch.zeros_like(p) for p in params]
+        self.t = torch.tensor(0.0, device=dev)
+
+    @torch.no_grad()
+    def step(self, lr, ok):
+        self.t += ok.float()
+        t = torch.clamp(self.t, min=1.0)
+        bc1 = 1.0 - self.b1 ** t
+        bc2 = 1.0 - self.b2 ** t
+        step_size = lr / bc1
+        inv_sqrt_bc2 = torch.rsqrt(bc2)
+        decay = 1.0 - lr * self.wd
+        for p, m, v in zip(self.params, self.m, self.v):
+            if p.grad is None:
+                continue
+            g = torch.where(ok, p.grad, torch.zeros_like(p.grad))
+            m_new = m + (g - m) * (1.0 - self.b1)
+            v_new = v + (g * g - v) * (1.0 - self.b2)
+            p_new = p * decay - step_size * m_new / (v_new.sqrt() * inv_sqrt_bc2 + self.eps)
+            p.copy_(torch.where(ok, p_new, p))
+            m.copy_(torch.where(ok, m_new, m))
+            v.copy_(torch.where(ok, v_new, v))
+
+    def zero_grad(self):
+        for p in self.params:
+            p.grad = None
+
+    def state_dict(self):
+        return {"m": [x.detach().cpu() for x in self.m], "v": [x.detach().cpu() for x in self.v],
+                "t": float(self.t.item())}
+
+    @torch.no_grad()
+    def load_state_dict(self, sd):
+        for a, b in zip(self.m, sd["m"]):
+            a.copy_(b.to(a.device))
+        for a, b in zip(self.v, sd["v"]):
+            a.copy_(b.to(a.device))
+        self.t.fill_(float(sd["t"]))
+
+
+def clip_grads(params, max_norm):
+    """Global grad-norm clip, host senkronizasyonu yok. Clip ÖNCESİ normu döner."""
+    grads = [p.grad for p in params if p.grad is not None]
+    total = torch.sqrt(sum(g.float().pow(2).sum() for g in grads))
+    coef = torch.clamp(max_norm / (total + 1e-6), max=1.0)
+    for g in grads:
+        g.mul_(coef)
+    return total
+
+
+def lr_tensor(step_t, base_lr, total_steps):
+    """Warmup + cosine, tamamen cihazda (step_t: float tensör)."""
+    warm = base_lr * (step_t + 1.0) / WARMUP
+    prog = torch.clamp((step_t - WARMUP) / max(1, total_steps - WARMUP), min=0.0, max=1.0)
+    cos = base_lr * (0.05 + 0.95 * 0.5 * (1.0 + torch.cos(math.pi * prog)))
+    return torch.where(step_t < WARMUP, warm, cos)
+
+
+# ------------------------------------------------------------------ CLIP YARDIMCILARI (CPU)
+def load_clip_text():
+    from transformers import CLIPTokenizer, CLIPTextModelWithProjection
+    tok = CLIPTokenizer.from_pretrained(CLIP_ID)
+    txt = CLIPTextModelWithProjection.from_pretrained(CLIP_ID).eval()
+    return tok, txt
+
+
+@torch.no_grad()
+def encode_prompts(tok, txt, prompts):
+    t = tok(prompts, padding="max_length", max_length=TOK_LEN, truncation=True, return_tensors="pt")
+    return txt(**t).last_hidden_state  # (B, TOK_LEN, 512)
+
+
+# ------------------------------------------------------------------ ÖRNEK ÜRETİMİ (DDIM + CFG)
+@torch.no_grad()
+def generate(unet, cond, null, steps=50, cfg=4.0, seed=0):
+    """
+    cond: (B,L,D), null: (1,L,D) -> [0,1] aralığında (B,3,H,W).
+    DDIM (eta=0, clip_sample) elle yazıldı: t / alpha değerleri her adımda GİRDİ tensörü,
+    böylece adım grafiği tek kez derlenir ve 50 adımda yeniden kullanılır.
+    """
+    from diffusers import DDIMScheduler
+    dev = cond.device
+    B = cond.shape[0]
+    sched = DDIMScheduler(num_train_timesteps=T_STEPS, beta_schedule="squaredcos_cap_v2", clip_sample=True,
+                          prediction_type="epsilon", timestep_spacing="trailing")
+    sched.set_timesteps(steps)                                   # CPU
+    ts = sched.timesteps.long()
+    prev = ts - T_STEPS // steps
+    ac = sched.alphas_cumprod.float()
+    a_t = ac[ts]
+    a_p = torch.where(prev >= 0, ac[prev.clamp(min=0)], sched.final_alpha_cumprod.float())
+    table = torch.stack([ts.float(), a_t, a_p], dim=1)           # (steps,3) CPU
+
+    g = torch.Generator().manual_seed(seed)
+    x = torch.randn(B, 3, RES, RES, generator=g).to(dev)
+    ctx = torch.cat([null.to(dev).float().expand(B, -1, -1), cond.float()])
+    for i in range(steps):
+        row = table[i].to(dev)
+        t = row[0].expand(2 * B)
+        with amp(dev):
+            eps = unet(torch.cat([x, x]), t, encoder_hidden_states=ctx).sample
+        e_u, e_c = eps.float().chunk(2)
+        eps = e_u + cfg * (e_c - e_u)
+        at, ap = row[1], row[2]
+        x0 = ((x - (1.0 - at).sqrt() * eps) / at.sqrt()).clamp(-1, 1)
+        x = ap.sqrt() * x0 + (1.0 - ap).sqrt() * eps
+        sync()
+    return (x.clamp(-1, 1) + 1) / 2
+
+
+def save_grid(x, path, nrow, pad=2):
+    x = (x.detach().float().cpu().clamp(0, 1) * 255 + 0.5).byte().permute(0, 2, 3, 1).numpy()
+    n, h, w, _ = x.shape
+    rows = (n + nrow - 1) // nrow
+    canvas = np.zeros((rows * (h + pad) + pad, nrow * (w + pad) + pad, 3), dtype=np.uint8)
+    for i in range(n):
+        r, c = divmod(i, nrow)
+        y0, x0 = pad + r * (h + pad), pad + c * (w + pad)
+        canvas[y0:y0 + h, x0:x0 + w] = x[i]
+    Path(path).parent.mkdir(parents=True, exist_ok=True)
+    Image.fromarray(canvas).save(str(path))
+
+
+# ------------------------------------------------------------------ 1) VERİ HAZIRLAMA
+def collect_files(image_dir):
+    exts = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
+    if image_dir is None:
+        from huggingface_hub import hf_hub_download
+        z = hf_hub_download(repo_id="huggan/anime-faces", filename="data.zip", repo_type="dataset")
+        out = WORK / "raw"
+        if not out.exists():
+            print("data.zip açılıyor...")
+            with zipfile.ZipFile(z) as zf:
+                zf.extractall(out)
+        image_dir = out
+    files = [p for p in Path(image_dir).rglob("*") if p.suffix.lower() in exts and not p.name.startswith("._")]
+    return sorted(files)
+
+
+def _load_square(p):
+    try:
+        im = Image.open(p).convert("RGB")
+    except Exception:
+        return None
+    w, h = im.size
+    s = min(w, h)
+    l, t = (w - s) // 2, (h - s) // 2
+    im = im.crop((l, t, l + s, t + s)).resize((RES, RES), Image.LANCZOS)
+    return np.asarray(im, dtype=np.uint8)
+
+
+def cmd_prepare(args):
+    from tqdm import tqdm
+    from collections import Counter
+    from concurrent.futures import ThreadPoolExecutor
+    from transformers import CLIPVisionModelWithProjection
+
+    DATA.mkdir(parents=True, exist_ok=True)
+    dev = get_device()
+
+    files = collect_files(args.image_dir)
+    print(f"{len(files)} dosya bulundu, {RES}x{RES} boyutuna getiriliyor...")
+    with ThreadPoolExecutor(max_workers=os.cpu_count() or 4) as ex:
+        arr = [a for a in tqdm(ex.map(_load_square, files), total=len(files)) if a is not None]
+    if not arr:
+        raise SystemExit("Hiç görsel yüklenemedi. --image_dir ile klasör ver.")
+    images = np.stack(arr)  # (N,H,W,3) uint8
+    N = len(images)
+    print("Toplam görsel:", N)
+
+    # --- CLIP zero-shot ile pseudo-caption (görsel encoder TPU'da, sabit batch)
+    tok, txt = load_clip_text()                                     # metin tarafı CPU'da (çok küçük iş)
+    vis = CLIPVisionModelWithProjection.from_pretrained(CLIP_ID).to(dev).eval()
+    mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1)
+    std = torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1)
+    BS = 256
+    feats = []
+    with torch.no_grad():
+        for i in tqdm(range(0, N, BS), desc="CLIP"):
+            x = torch.from_numpy(images[i:i + BS]).permute(0, 3, 1, 2).float() / 255.0
+            n = x.shape[0]
+            x = F.interpolate(x, size=224, mode="bicubic", align_corners=False).clamp(0, 1)
+            x = (x - mean) / std
+            if n < BS:                                              # statik şekil -> tek derleme
+                x = torch.cat([x, x.new_zeros(BS - n, *x.shape[1:])])
+            f = vis(pixel_values=x.to(dev)).image_embeds
+            feats.append(F.normalize(f.float(), dim=-1)[:n].cpu())
+    feats = torch.cat(feats)                                        # (N,512) CPU
+    del vis
+
+    @torch.no_grad()
+    def classify(prompts):
+        t = tok(prompts, padding=True, return_tensors="pt")
+        te = F.normalize(txt(**t).text_embeds, dim=-1)
+        return (feats @ te.T).argmax(1).numpy()
+
+    hair_c = classify([f"an anime character with {c} hair" for c in HAIR_COLORS])
+    eye_c = classify([f"an anime character with {c} eyes" for c in EYE_COLORS])
+    hair_l = classify([f"an anime character with {c} hair" for c in HAIR_LEN])
+    smile = classify(["an anime face with a smile", "an anime face with a neutral expression"])
+
+    for name, a, names in [("saç rengi", hair_c, HAIR_COLORS), ("göz rengi", eye_c, EYE_COLORS),
+                           ("saç uzunluğu", hair_l, HAIR_LEN), ("gülümseme", smile, ["gülüyor", "nötr"])]:
+        cnt = Counter(a.tolist())
+        print(name, {names[k]: v for k, v in sorted(cnt.items())})
+
+    caps = []
+    for i in range(N):
+        c = f"anime face, {HAIR_COLORS[hair_c[i]]} {HAIR_LEN[hair_l[i]]} hair, {EYE_COLORS[eye_c[i]]} eyes"
+        if smile[i] == 0:
+            c += ", smiling"
+        caps.append(c)
+
+    uniq = sorted(set(caps))
+    captions = [""] + uniq                      # index 0 = boş caption (CFG için null)
+    cmap = {c: i for i, c in enumerate(captions)}
+    labels = np.array([cmap[c] for c in caps], dtype=np.int64)
+
+    emb = torch.cat([encode_prompts(tok, txt, captions[i:i + 128]).half()
+                     for i in range(0, len(captions), 128)])
+
+    np.save(DATA / "images.npy", images)
+    np.save(DATA / "labels.npy", labels)
+    torch.save({"emb": emb, "captions": captions}, DATA / "text_emb.pt")
+
+    prev = torch.from_numpy(images[:32]).permute(0, 3, 1, 2).float() / 255.0
+    save_grid(prev, WORK / "preview_labels.png", nrow=8)
+    print(f"\n{len(captions) - 1} benzersiz caption. İlk 32 görselin captionları (8'erli satırlar, preview_labels.png ile karşılaştır):")
+    for i in range(32):
+        print(i, "|", caps[i])
+    print("\nHazır ->", DATA)
+
+
+# ------------------------------------------------------------------ 2) EĞİTİM
+def save_train_state(path, unet, ema, opt, step):
+    state = {"unet": {k: v.detach().cpu() for k, v in unet.state_dict().items()},
+             "ema": [s.detach().cpu() for s in ema.shadow], "ema_n": float(ema.n.item()),
+             "opt": opt.state_dict(), "step": step}
+    tmp = CKPT / "train_state.tmp"
+    torch.save(state, tmp)
+    os.replace(tmp, path)
+
+
+def export_ema(ema, base_ch, path):
+    """EMA ağırlıklarını CPU'da diffusers formatında (save_pretrained) yaz."""
+    m = build_unet(base_ch, attn="sdpa")
+    with torch.no_grad():
+        for p, s in zip(m.parameters(), ema.shadow):
+            p.copy_(s.detach().cpu())
+    m.save_pretrained(str(path))
+
+
+def cmd_train(args):
+    from diffusers import DDPMScheduler
+
+    dev = get_device()
+    seed_all(SEED)
+
+    # tüm veri HBM'de (uint8 NCHW, ~260MB)
+    images = torch.from_numpy(np.load(DATA / "images.npy")).permute(0, 3, 1, 2).contiguous().to(dev)
+    labels = torch.from_numpy(np.load(DATA / "labels.npy")).long().to(dev)
+    te = torch.load(DATA / "text_emb.pt", weights_only=False)
+    text_emb = te["emb"].float().to(dev)                  # (K,L,512) fp32
+    captions = te["captions"]
+    null = text_emb[0]                                    # (L,512) boş caption
+    N = images.shape[0]
+
+    unet = build_unet(args.base_ch, args.attn).to(dev)
+    params = [p for p in unet.parameters() if p.requires_grad]
+    nparam = sum(p.numel() for p in params)
+    print(f"cihaz: {dev} | UNet: {nparam / 1e6:.2f}M param | görsel: {N} | batch: {args.batch} "
+          f"| lr: {args.lr} | attn: {args.attn} | bf16 autocast")
+
+    ck, start_step = None, 0
+    if args.resume:
+        ck = torch.load(args.resume, map_location="cpu", weights_only=False)
+        unet.load_state_dict(ck["unet"])
+        start_step = ck["step"]
+        print("Devam ediliyor, adım:", start_step)
+
+    opt = DeviceAdamW(params, dev, betas=(0.9, 0.99), eps=1e-6, weight_decay=0.01)
+    if ck is not None:
+        if isinstance(ck.get("opt"), dict) and "m" in ck["opt"]:
+            opt.load_state_dict(ck["opt"])
+        else:
+            print("Uyarı: checkpoint eski (GPU) formatında, optimizer durumu sıfırdan başlıyor.")
+
+    ema = EMA(params, dev, decay=0.9995)
+    if ck is not None:
+        for s, p in zip(ema.shadow, ck["ema"]):
+            s.copy_(p.to(dev))
+        ema.n.fill_(float(ck.get("ema_n", 0)))
+    ema_model = build_unet(args.base_ch, args.attn).to(dev).eval()
+
+    ac = DDPMScheduler(num_train_timesteps=T_STEPS, beta_schedule="squaredcos_cap_v2").alphas_cumprod.float().to(dev)
+    snr = ac / (1.0 - ac)
+    minsnr_w = torch.clamp(snr, max=MIN_SNR_GAMMA) / snr
+
+    # periyodik örnekler için en sık geçen 6 caption
+    freq = np.bincount(labels.cpu().numpy(), minlength=len(captions))
+    top = [int(i) for i in np.argsort(-freq) if i != 0][:6]
+    sample_cond = text_emb[torch.tensor(top, device=dev)].repeat_interleave(6, dim=0)   # 6 prompt x 6 örnek
+    OUT.mkdir(parents=True, exist_ok=True)
+    CKPT.mkdir(parents=True, exist_ok=True)
+    (OUT / "sample_prompts.txt").write_text("\n".join(captions[i] for i in top))
+
+    seed_all(SEED + start_step)
+    B = args.batch
+    step = start_step
+    step_t = torch.tensor(float(start_step), device=dev)
+    run = torch.zeros((), device=dev)
+    gn_sum = torch.zeros((), device=dev)
+    gn_max = torch.zeros((), device=dev)
+    n_bad = torch.zeros((), device=dev)
+    unet.train()
+
+    print("İlk adımda XLA derlemesi yapılır (birkaç dakika sürebilir; önbellekten sonraki çalıştırmalarda hızlıdır)...")
+    t_start = time.time()
+    t_log = time.time()
+
+    while step < args.steps:
+        # ---- batch (hepsi cihazda)
+        idx = torch.randint(0, N, (B,), device=dev)
+        x0 = images[idx].float().div_(127.5).sub_(1.0)
+        flip = torch.rand(B, device=dev) < 0.5
+        x0 = torch.where(flip[:, None, None, None], x0.flip(3), x0)
+
+        cond = text_emb[labels[idx]]
+        drop = torch.rand(B, device=dev) < P_UNCOND
+        cond = torch.where(drop[:, None, None], null[None], cond)
+
+        noise = torch.randn_like(x0)
+        t = torch.randint(0, T_STEPS, (B,), device=dev)
+        a_t = ac[t].sqrt()[:, None, None, None]
+        s_t = (1.0 - ac[t]).sqrt()[:, None, None, None]
+        xt = a_t * x0 + s_t * noise
+
+        # ---- ileri / geri
+        with amp(dev):
+            pred = unet(xt, t, encoder_hidden_states=cond).sample
+        loss = (((pred.float() - noise) ** 2).mean(dim=(1, 2, 3)) * minsnr_w[t]).mean()
+        loss.backward()
+
+        # ---- optimizer (cihazda, senkronizasyonsuz)
+        gn = clip_grads(params, 1.0)
+        ok = torch.isfinite(gn) & torch.isfinite(loss)
+        lr = lr_tensor(step_t, args.lr, args.steps)
+        opt.step(lr, ok)
+        opt.zero_grad()
+        ema.update(params)
+
+        # ---- istatistikler (cihazda birikir)
+        okf = ok.float()
+        run += torch.where(ok, loss.detach(), torch.zeros_like(loss))
+        gn_ok = torch.where(ok, gn, torch.zeros_like(gn))
+        gn_sum += gn_ok
+        gn_max = torch.maximum(gn_max, gn_ok)
+        n_bad += 1.0 - okf
+        step_t += 1.0
+
+        sync()                       # adım sınırı: grafiği gönder, host bir sonraki adımı izlemeye devam eder
+        step += 1
+
+        if step == start_step + 1:
+            loss.item()              # ilk derlemenin bitmesini bekle, süreyi göster
+            print(f"ilk adım (derleme dahil): {time.time() - t_start:.1f} sn")
+            t_log = time.time()
+
+        if step % LOG_EVERY == 0:
+            bad = int(n_bad.item())
+            good = max(1, LOG_EVERY - bad)
+            avg = run.item() / good
+            gn_avg = gn_sum.item() / good
+            ips = LOG_EVERY / (time.time() - t_log)
+            t_log = time.time()
+            print(f"step {step}/{args.steps} | loss {avg:.4f} | lr {lr.item():.2e} | {ips:.2f} it/s "
+                  f"| {(time.time() - t_start) / 60:.1f} dk | gn {gn_avg:.2f} (maks {gn_max.item():.1f}) "
+                  f"| atlanan {bad}")
+            run.zero_(); gn_sum.zero_(); gn_max.zero_(); n_bad.zero_()
+            if bad >= LOG_EVERY:
+                print("Tüm adımlar NaN/inf! --lr düşür veya --resume ile son sağlam checkpoint'ten devam et.")
+                break
+            if step - start_step >= 1000 and avg > args.diverge:
+                print(f"Loss {avg:.3f} > {args.diverge}: model çöktü, duruyorum. "
+                      f"Son sağlam checkpoint'ten daha düşük --lr ile --resume yap.")
+                break
+
+        if step % args.sample_every == 0 or step == args.steps:
+            ema.copy_to(ema_model)
+            export_ema(ema, args.base_ch, CKPT / "ema_unet")
+            if step % args.ckpt_every == 0 or step == args.steps:
+                save_train_state(CKPT / f"train_state_{step}.pt", unet, ema, opt, step)
+            imgs = generate(ema_model, sample_cond, null[None], steps=args.sample_steps, cfg=4.0, seed=0)
+            save_grid(imgs, OUT / f"step_{step:06d}.png", nrow=6)
+            print(f"[kaydedildi] step {step} -> {OUT / f'step_{step:06d}.png'}")
+            t_log = time.time()
+            if (time.time() - t_start) / 3600.0 > args.max_hours:
+                print(f"--max_hours ({args.max_hours} sa) doldu, durduruluyor. --resume ile devam edebilirsin.")
+                break
+
+    print("Bitti. Model:", CKPT / "ema_unet")
+
+
+# ------------------------------------------------------------------ 3) ÖRNEK ÜRETME
+def cmd_sample(args):
+    from diffusers import UNet2DConditionModel
+    dev = get_device()
+    unet = set_attn(UNet2DConditionModel.from_pretrained(str(CKPT / "ema_unet")), args.attn).to(dev).eval()
+    tok, txt = load_clip_text()
+    cond = encode_prompts(tok, txt, args.prompt).repeat_interleave(args.n, dim=0).to(dev)
+    null = encode_prompts(tok, txt, [""]).to(dev)
+    imgs = generate(unet, cond, null, steps=args.steps, cfg=args.cfg, seed=args.seed)
+    OUT.mkdir(parents=True, exist_ok=True)
+    path = OUT / f"sample_{int(time.time())}.png"
+    save_grid(imgs, path, nrow=args.n)
+    print("Kaydedildi:", path, "| satırlar =", args.prompt)
+
+
+# ------------------------------------------------------------------ CLI
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    sub = ap.add_subparsers(dest="cmd", required=True)
+
+    p1 = sub.add_parser("prepare")
+    p1.add_argument("--image_dir", default=None, help="Boşsa huggan/anime-faces indirilir")
+
+    p2 = sub.add_parser("train")
+    p2.add_argument("--steps", type=int, default=40000)
+    p2.add_argument("--batch", type=int, default=128, help="Batch (TPU v5e-1 tek çip). OOM olursa 64; rahatsa 256")
+    p2.add_argument("--lr", type=float, default=1e-4)
+    p2.add_argument("--base_ch", type=int, default=64, help="UNet taban kanal (64 -> ~17M, 128 -> MXU dostu ~4x param)")
+    p2.add_argument("--attn", choices=["chunked", "sdpa"], default="chunked",
+                    help="chunked: bellek-verimli özel self-attention; sdpa: standart (büyük batch'te OOM riski)")
+    p2.add_argument("--ckpt_every", type=int, default=2000)
+    p2.add_argument("--sample_every", type=int, default=500)
+    p2.add_argument("--sample_steps", type=int, default=50, help="Periyodik örneklerde DDIM adımı")
+    p2.add_argument("--max_hours", type=float, default=11.0)
+    p2.add_argument("--diverge", type=float, default=0.4, help="100 adım ort. loss bunu aşarsa dur")
+    p2.add_argument("--resume", default=None)
+
+    p3 = sub.add_parser("sample")
+    p3.add_argument("--prompt", nargs="+", default=["anime face, blue long hair, red eyes, smiling"])
+    p3.add_argument("--n", type=int, default=6, help="prompt başına görsel sayısı")
+    p3.add_argument("--cfg", type=float, default=4.0)
+    p3.add_argument("--steps", type=int, default=50)
+    p3.add_argument("--seed", type=int, default=0)
+    p3.add_argument("--attn", choices=["chunked", "sdpa"], default="chunked")
+
+    args = ap.parse_args()
+    init_xla_cache()
+    {"prepare": cmd_prepare, "train": cmd_train, "sample": cmd_sample}[args.cmd](args)