26.6-dars: Diffusion — sampling va boshqaruv

26.5 DAGI TESKARI JARAYON (ajdodli, "ancestral"):
  x_T ~ N(0, I)
  t = T-1, ..., 0:
    eps = model(x_t, t)                           <- tarmoq chaqiruvi
    x0_hat = (x_t - sqrt(1 - ab_t) * eps) / sqrt(ab_t)
    x_{t-1} = mu_t(x_t, x0_hat) + sigma_t * z,  z ~ N(0, I)

NARX:
  NFE (network function evaluations) = T
  bizda: T = 100, model 173 120 parametr
    bitta chaqiruv ~ 2 * 173 120 = 0.35 MFLOP / rasm
    butun sampling ~ 35 MFLOP / rasm
  qadamlar KETMA-KET: x_{t-1} x_t ga bog'liq - vaqt bo'yicha parallellab bo'lmaydi
  (batch bo'yicha parallel - mumkin; lekin kechikish (latency) T ga proporsional)

TAQQOSLASH:
  GAN, VAE dekoderi:  1 chaqiruv
  DDPM (T = 1000 odatiy qog'ozlarda): 1000 chaqiruv

G'OYA (Song va boshq., 2020):
  o'quv maqsadi o'sha-o'sha (eps ni bashorat qilish) - qayta o'rgatish YO'Q
  faqat sampler o'zgaradi: x_t dan istalgan oldingi s < t ga sakrash

BIR QADAM (t -> s):
  eps    = model(x_t, t)
  x0_hat = (x_t - sqrt(1 - ab_t) * eps) / sqrt(ab_t)       (26.5 dagi kabi)
  x_s    = sqrt(ab_s) * x0_hat + sqrt(1 - ab_s - sigma^2) * eps + sigma * z

  sigma = eta * sqrt((1 - ab_s) / (1 - ab_t) * (1 - ab_t / ab_s))
    eta = 0: sigma = 0 -> DETERMINISTIK (bir xil x_T -> bir xil rasm)
    eta = 1: DDPM bilan bir xil shovqin darajasi

TALQIN:
  "x0 ni taxmin qil -> uni s darajasiga qayta shovqinla, lekin tasodifiy
   shovqin o'rniga model bashorat qilgan eps ni ishlat"
  eta = 0 da jarayon - oddiy differensial tenglama (ODE) ning diskret yechimi

QADAMLAR JADVALI:
  S ta teng oraliqli t: linspace(T-1, 0, S)
  S = 10 -> t = 99, 88, 77, 66, 55, 44, 33, 22, 11, 0

  usul          qadam     FD          izoh
  DDIM eta=0     100     21.55
  DDIM eta=0      20     21.81
  DDIM eta=0      10     22.17       <- 10 barobar kam chaqiruv, FD +0.62
  DDIM eta=0       5     25.87
  DDIM eta=0       3     52.49
  DDIM eta=0       2    269.04       <- qulash
  o'quv rasmlarining o'zi   2.23     (pastki chegara)

1-MISOL:
  DDPM (100 qadam)          FD 38.61
  DDIM eta=1, S = 100       FD 38.61   <- AYNAN bir xil (bir xil urug'lar)
  DDIM eta=0, S = 100       FD 21.55

  eta = 1, S = T da DDIM formulasi DDPM ning o'zi (sigma^2 = posterior dispersiya)

  eta=1 da qadamlar kamaysa sifat tezroq tushadi:
    S = 20: 41.36,  S = 10: 47.33,  S = 5: 56.91

NEGA BIZDA STOXASTIK SAMPLING YOMONROQ:
  har qadamda yangi shovqin qo'shiladi -> model uni keyingi qadamlarda
  to'liq "tozalay" olishi kerak; kichik, kam o'rgatilgan modelda
  eps xatolari yig'iladi (ehtimoliy sabab; katta modellarda farq kichikroq,
  ba'zan stoxastik sampling hatto yaxshi)

SHARTSIZ:  eps(x_t, t)          -> p(x)
SHARTLI:   eps(x_t, t, y)       -> p(x | y)

QANDAY BERILADI (bizning MLP):
  h = W_x x_t + emb_vaqt(t) + emb_sinf(y)       - yashirin qatlamga qo'shiladi
  emb_sinf = nn.Embedding(11, 256)               - 10 raqam + 1 "shart yo'q"

KATTA MODELLARDA:
  sinf -> embedding, vaqt embeddingiga qo'shiladi (U-Net ning har blokiga)
  matn -> matn encoder -> cross-attention (24.3 dagi mexanizm; 26.8-dars)

O'RGATISH: odatdagi eps-MSE, faqat (x0, y) juftliklari bilan

G'OYA (Ho va Salimans, 2022):
  bitta modelni IKKI rejimda o'rgatamiz:
    ehtimol 1 - p: shart bilan         eps(x_t, t, y)
    ehtimol p:     shart "tashlanadi"  eps(x_t, t, BOSH)     (bizda p = 0.15)

  samplingda ikkalasini aralashtiramiz:
    eps_w = eps_u + w * (eps_c - eps_u)
          = (1 - w) * eps_u + w * eps_c

  w = 0: shartsiz        w = 1: oddiy shartli
  w > 1: "shartsizdan uzoqroq, shart tomonga" - ekstrapolyatsiya

NEGA ISHLAYDI (skor talqini):
  eps ~ -grad log p(x_t) (masshtabgacha)
  eps_c - eps_u ~ -grad log p(y | x_t)   (Bayes: p(x|y) ~ p(y|x) p(x))
  w - "yashirin klassifikator" gradientining og'irligi
  -> p(x) * p(y | x)^w  ga o'xshash taqsimotdan namuna

NARX: har qadamda 2 chaqiruv (yoki 2 barobar batch)

    w   to'g'ri sinf    FD     xilma/haqiqiy
    0      0.097       38.51      1.63       <- shartsiz: tasodifiy (1/10)
    1      0.830       27.90      1.08
    2      0.977       16.41      0.81       <- eng past FD
    3      0.985       19.73      0.75
    5      0.985       34.37      0.74
    8      0.976       61.68      0.76       <- FD yana yomon, piksellar "kuyadi"

CLASSIFIER GUIDANCE (Dhariwal va Nichol, 2021):
  alohida klassifikator p(y | x_t) ni SHOVQINLI rasmlarda o'rgatish
  eps ga uning gradientini qo'shish
  kamchilik: qo'shimcha model, shovqinli rasmda o'rgatish kerak

CLASSIFIER-FREE (bu dars):
  qo'shimcha model YO'Q - shartni tashlab yuborish kifoya
  bugungi matn-rasm tizimlarining odatiy usuli

SALBIY PROMPT (amaliyotda):
  eps_u o'rniga "istalmagan" shartdagi bashorat ishlatiladi:
    eps_w = eps_neg + w * (eps_c - eps_neg)

VAZIFA: niqob m (1 - ma'lum, 0 - noma'lum), ma'lum qism x_malum
  noma'lum qismni "mos" qilib generatsiya qilish

1. SODDA: shartsiz namuna -> oxirida ma'lum qismni yopishtirish
   noma'lum qism ma'lum qismni "bilmaydi" -> mos kelmaydi

2. HAR QADAMDA ALMASHTIRISH (qayta o'rgatishsiz):
   har t da:  x_t[ma'lum] = sqrt(ab_t) * x_malum + sqrt(1 - ab_t) * z
              (ma'lum qismning t-darajadagi shovqinli nusxasi)
   keyin odatdagi DDPM qadami
   muammo: model ma'lum qismni "ko'radi", lekin uyg'unlashtirishga vaqt kam

3. QAYTA NAMUNA (RePaint g'oyasi, Lugmayr va boshq., 2022):
   t -> t-1 qadamdan keyin yana t ga shovqin qo'shib, qadamni r marta takrorlash
   noma'lum qism ma'lum qismga "moslashishga" ko'proq imkon oladi
   narx: ~r * T chaqiruv (bizda r = 5: 496 chaqiruv)

4. MAXSUS O'RGATILGAN: niqob va ma'lum qism modelga kirish sifatida beriladi
   (amaliy tizimlarda keng tarqalgan; qayta o'rgatish kerak)

  usul                         MSE(past)  to'g'ri sinf
  o'rtacha bilan                 0.304       0.685
  eng yaqin qo'shni (3 ta)       0.239       0.882     <- eng yaxshi!
  diffusion: oxirida yopishtir   0.473       0.500
  diffusion: har qadamda         0.354       0.663
  har qadamda + qayta (r=5)      0.270       0.760

MUAMMO: katta rasmda piksel fazosi ulkan
  512 x 512 x 3 = 786 432 son; har qadam U-Net shu o'lchamda ishlaydi

G'OYA (Rombach va boshq., 2022 - "Latent Diffusion Models"):
  1. avtoenkoder (VAE) o'rgatish: rasm -> kichik latent -> rasm
     Stable Diffusion da: 512x512x3 -> 64x64x4 (8 barobar kichraytirish)
     = 16 384 son, ~48 barobar kam
  2. diffusion LATENTDA o'rgatiladi
  3. sampling: latentda S qadam -> dekoder BIR marta

AFZALLIK: har qadam arzon; dekoder mayda detallarni "biladi"
CHEKLOV: sifat avtoenkoder rekonstruksiyasidan yaxshi bo'la olmaydi

4-MISOL (digits, 64 -> 8 o'lchamli latent, bir xil 2000 o'quv qadami):
                 parametr   FLOP/qadam   sampling (S = 20)    FD
  piksel          173 120     346 K        6.92 M           20.66
  latent           39 432      79 K        1.60 M           35.48
  AE rekonstruksiya chegarasi:                                4.21
  -> latent 4.4 barobar arzon, lekin shu byudjetda sezilarli yomonroq

import torch

# DDIM qadami (t -> s), eta = 0 - deterministik
eps = model(x, t)
x0 = ((x - (1 - ab[t]).sqrt() * eps) / ab[t].sqrt()).clamp(-1, 1)
eps = (x - ab[t].sqrt() * x0) / (1 - ab[t]).sqrt()
sigma = eta * ((1 - ab[s]) / (1 - ab[t]) * (1 - ab[t] / ab[s])).sqrt()
x = ab[s].sqrt() * x0 + (1 - ab[s] - sigma ** 2).sqrt() * eps + sigma * torch.randn_like(x)

# qadamlar jadvali
qadamlar = torch.linspace(T - 1, 0, S).round().long().tolist()

# CFG: o'rgatishda shartni tashlash
y[torch.rand(B, generator=g) < 0.15] = BOSH
# samplingda aralashtirish
eps = eps_u + w * (eps_c - eps_u)

# inpainting: har qadamda ma'lum qismni almashtirish
x = niqob * (ab[t].sqrt() * x_malum + (1 - ab[t]).sqrt() * z) + (1 - niqob) * x

# latent diffusion
z = (ae.enc(X) - mu) / sd           # latentni standartlash
rasm = ae.dec(ddim(model_z, torch.randn(n, 8), S) * sd + mu)

DDPM: T ta ketma-ket chaqiruv; qadam tashlab bo'lmaydi
DDIM: qayta o'rgatishsiz, S << T qadam; eta = 0 deterministik
shartli: emb(y) qo'shiladi; CFG: eps_u + w (eps_c - eps_u)
w oshsa: to'g'rilik oshadi, xilma-xillik kamayadi, FD U-shaklida
inpainting: har qadamda ma'lum qismni shovqinli nusxa bilan almashtirish
latent diffusion: AE latentida diffusion; chegara - AE sifati

"""DDPM sampling narxi va DDIM noldan: qadamlar soni va sifat egri chizig'i."""

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import load_digits

T = 100
BELGI = " .:-=+*#%@"


def kosinus_jadval(T, s=0.008):
    t = torch.arange(T + 1, dtype=torch.float64) / T
    f = torch.cos((t + s) / (1 + s) * math.pi / 2) ** 2
    return (f / f[0])[1:].clamp(1e-5, 0.9999).float()


AB = kosinus_jadval(T)                            # alfa_bar_t, t = 0..T-1
ALFA = torch.cat([AB[:1], AB[1:] / AB[:-1]])
BETA = 1 - ALFA


def vaqt_emb(t, d=32):
    k = torch.exp(-math.log(1000.0) * torch.arange(d // 2) / (d // 2))
    a = t.float()[:, None] * k[None]
    return torch.cat([a.sin(), a.cos()], 1)


class Shovqin(nn.Module):
    """eps(x_t, t): 64 piksel + vaqt embeddingi -> MLP -> 64 (26.5-darsdagi model)."""

    def __init__(self, D=64, h=256):
        super().__init__()
        self.vaqt = nn.Sequential(nn.Linear(32, h), nn.SiLU())
        self.kir = nn.Linear(D, h)
        self.tana = nn.Sequential(nn.SiLU(), nn.Linear(h, h), nn.SiLU(), nn.Linear(h, h),
                                  nn.SiLU(), nn.Linear(h, D))

    def forward(self, x, t):
        return self.tana(self.kir(x) + self.vaqt(vaqt_emb(t)))


class Klassifikator(nn.Module):
    """Baholovchi: 64 -> 128 -> 32 (xususiyatlar) -> 10."""

    def __init__(self):
        super().__init__()
        self.xus = nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 32), nn.ReLU())
        self.bosh = nn.Linear(32, 10)

    def forward(self, x):
        return self.bosh(self.xus(x))


def malumot():
    d = load_digits()
    X = torch.tensor(d.data / 8.0 - 1.0, dtype=torch.float32)      # [-1, 1]
    p = torch.tensor(np.random.default_rng(0).permutation(len(X)))
    return X[p], torch.tensor(d.target)[p]


def orgat_diffusion(X, qadamlar=2000, B=128, seed=0):
    torch.manual_seed(seed)
    m = Shovqin()
    opt = torch.optim.Adam(m.parameters(), lr=2e-3)
    jadval = torch.optim.lr_scheduler.CosineAnnealingLR(opt, qadamlar)
    g = torch.Generator().manual_seed(seed)
    for _ in range(qadamlar):
        x0 = X[torch.randint(0, len(X), (B,), generator=g)]
        t = torch.randint(0, T, (B,), generator=g)
        eps = torch.randn(x0.shape, generator=g)
        ab = AB[t][:, None]
        loss = F.mse_loss(m(ab.sqrt() * x0 + (1 - ab).sqrt() * eps, t), eps)
        opt.zero_grad()
        loss.backward()
        opt.step()
        jadval.step()
    m.eval()
    return m


def orgat_klassifikator(X, y, seed=0):
    torch.manual_seed(seed)
    k = Klassifikator()
    opt = torch.optim.Adam(k.parameters(), lr=3e-3, weight_decay=1e-4)
    g = torch.Generator().manual_seed(seed)
    for _ in range(800):
        i = torch.randint(0, len(X), (128,), generator=g)
        loss = F.cross_entropy(k(X[i] + 0.1 * torch.randn(X[i].shape, generator=g)), y[i])
        opt.zero_grad()
        loss.backward()
        opt.step()
    k.eval()
    return k


def frechet(a, b):
    """||mu1 - mu2||^2 + tr(S1 + S2 - 2 sqrt(S1 S2)) - klassifikator xususiyatlarida."""
    a, b = a.double().numpy(), b.double().numpy()
    s1, s2 = np.cov(a, rowvar=False), np.cov(b, rowvar=False)
    w, v = np.linalg.eigh(s1)
    ild = (v * np.sqrt(np.clip(w, 0, None))) @ v.T
    oz = np.linalg.eigvalsh(ild @ s2 @ ild)
    return float(((a.mean(0) - b.mean(0)) ** 2).sum() + np.trace(s1) + np.trace(s2)
                 - 2 * np.sqrt(np.clip(oz, 0, None)).sum())


@torch.no_grad()
def ddpm(m, x, g):
    """Ajdodli (ancestral) DDPM: T ta qadam, har qadamda yangi shovqin."""
    for t in range(T - 1, -1, -1):
        eps = m(x, torch.full((len(x),), t))
        x0 = ((x - (1 - AB[t]).sqrt() * eps) / AB[t].sqrt()).clamp(-1, 1)
        ab_old = AB[t - 1] if t > 0 else torch.tensor(1.0)
        orta = (ab_old.sqrt() * BETA[t] * x0 + ALFA[t].sqrt() * (1 - ab_old) * x) / (1 - AB[t])
        if t > 0:
            x = orta + (BETA[t] * (1 - ab_old) / (1 - AB[t])).sqrt() * torch.randn(x.shape, generator=g)
        else:
            x = orta
    return x


@torch.no_grad()
def ddim(m, x, S, eta=0.0, g=None):
    """DDIM: S ta qadam. eta = 0 - deterministik, eta = 1 - DDPM ga o'xshash shovqinli."""
    qadamlar = torch.linspace(T - 1, 0, S).round().long().tolist()
    for i, t in enumerate(qadamlar):
        eps = m(x, torch.full((len(x),), t))
        x0 = ((x - (1 - AB[t]).sqrt() * eps) / AB[t].sqrt()).clamp(-1, 1)
        if i + 1 == len(qadamlar):
            return x0
        ab, ab_old = AB[t], AB[qadamlar[i + 1]]
        eps = (x - ab.sqrt() * x0) / (1 - ab).sqrt()
        sigma = eta * ((1 - ab_old) / (1 - ab) * (1 - ab / ab_old)).sqrt()
        x = ab_old.sqrt() * x0 + (1 - ab_old - sigma ** 2).sqrt() * eps
        if eta > 0:
            x = x + sigma * torch.randn(x.shape, generator=g)
    return x


def ascii_qator(rasmlar):
    satrlar = []
    for r in range(8):
        qism = []
        for x in rasmlar:
            q = ((x.view(8, 8)[r] + 1) / 2).clamp(0, 1)
            qism.append("".join(BELGI[min(9, int(v * 10))] for v in q.tolist()))
        satrlar.append("    " + "   ".join(qism))
    return "\n".join(satrlar)


def main() -> None:
    torch.set_num_threads(1)
    X, y = malumot()
    Xo, yo, Xt = X[:1200], y[:1200], X[1200:]
    m = orgat_diffusion(Xo)
    k = orgat_klassifikator(Xo, yo)
    with torch.no_grad():
        f_test = k.xus(Xt)
    print(f"  test aniqligi (baholovchi): {(k(Xt).argmax(1) == y[1200:]).float().mean():.3f}")

    def baho(s):
        with torch.no_grad():
            f = k.xus(s)
            p = torch.softmax(k.bosh(f), 1)
        sinf = torch.bincount(p.argmax(1), minlength=10).float() / len(s)
        return frechet(f, f_test), p.max(1).values.mean().item(), int((sinf >= 0.05).sum())

    n_par = sum(p.numel() for p in m.parameters())
    print("\n=== 1. DDPM narxi: T ta ketma-ket tarmoq chaqiruvi ===")
    print(f"  model parametrlari: {n_par}, bitta chaqiruv ~{2 * n_par / 1e6:.2f} MFLOP/rasm")
    print(f"  DDPM: T = {T} qadam -> {T} chaqiruv -> ~{T * 2 * n_par / 1e6:.0f} MFLOP/rasm")
    print("  qadamlar KETMA-KET: t qadam t+1 natijasini kutadi - parallellab bo'lmaydi")

    print("\n=== 2. Sifat va qadamlar soni (500 namuna, 3 sampling urug'i) ===")
    n = 500
    fd_o, *_ = baho(Xo[:n])
    print(f"  pastki chegara - o'quv rasmlarining o'zi: FD {fd_o:.2f}")
    print("  usul              qadam   FD (o'rt +- SE)   ishonch  sinflar>=5%")
    natija = {}
    for nom, S, eta in ([("DDPM (ajdodli)", T, None)] +
                        [(f"DDIM eta=0", S, 0.0) for S in (100, 50, 20, 10, 5, 3, 2)] +
                        [(f"DDIM eta=1", S, 1.0) for S in (100, 20, 10, 5)]):
        qiymat = []
        for u in range(3):
            x = torch.randn((n, 64), generator=torch.Generator().manual_seed(100 + u))
            g = torch.Generator().manual_seed(200 + u)
            s = ddpm(m, x, g) if eta is None else ddim(m, x, S, eta, g)
            qiymat.append(baho(s))
        q = np.array(qiymat)
        natija[(nom, S)] = q[:, 0]
        print(f"  {nom:<16} {S:>5}   {q[:, 0].mean():>6.2f} +- {q[:, 0].std(ddof=1) / math.sqrt(3):<5.2f}"
              f"   {q[:, 1].mean():>7.3f} {q[:, 2].mean():>8.1f}")

    print("\n=== 3. Qaror: eng kam qadam, eng yaxshisidan sezilarli yomon emas ===")
    ddim0 = {S: v for (nom, S), v in natija.items() if nom == "DDIM eta=0"}
    eng = min(ddim0, key=lambda S: ddim0[S].mean())
    print(f"  juftlashgan farq (bir xil x_T) eng yaxshisidan (S = {eng}); amaliy chegara - FD ning 5%")
    maqbul = []
    for S in sorted(ddim0):
        farq = ddim0[S] - ddim0[eng]
        se = farq.std(ddof=1) / math.sqrt(len(farq))
        stat = farq.mean() > 2 * se + 1e-9
        amaliy = farq.mean() > 0.05 * ddim0[eng].mean()
        if not amaliy:
            maqbul.append(S)
        print(f"  S = {S:>3}: {farq.mean():+7.2f} (SE {se:.2f})  statistik: "
              f"{'sezilarli' if stat else 'sezilarli emas':<14}  amaliy: {'yomon' if amaliy else 'maqbul'}")
    tanlov = min(maqbul)
    print(f"  tanlov: DDIM S = {tanlov} -> {T // tanlov} barobar kam chaqiruv "
          f"(~{tanlov * 2 * n_par / 1e6:.1f} MFLOP/rasm)")
    d_ddpm = natija[("DDPM (ajdodli)", T)]
    print(f"  DDPM (100 qadam) FD {d_ddpm.mean():.2f}, DDIM S = {tanlov} FD {ddim0[tanlov].mean():.2f}")

    print("\n=== 4. DDIM deterministik: bir xil x_T -> bir xil rasm ===")
    x = torch.randn((200, 64), generator=torch.Generator().manual_seed(7))
    a1, a2 = ddim(m, x.clone(), 50), ddim(m, x.clone(), 50)
    print(f"  ikki marta, bir xil x_T: max farq {float((a1 - a2).abs().max()):.1e}")
    b10 = ddim(m, x.clone(), 10)
    boshqa = ddim(m, torch.randn((200, 64), generator=torch.Generator().manual_seed(8)), 50)
    print(f"  S = 50 va S = 10, bir xil x_T:   o'rtacha piksel farqi {float((a1 - b10).abs().mean()):.3f}")
    print(f"  S = 50, boshqa x_T:               o'rtacha piksel farqi {float((a1 - boshqa).abs().mean()):.3f}")
    with torch.no_grad():
        mos = (k(a1).argmax(1) == k(b10).argmax(1)).float().mean().item()
    print(f"  S = 50 va S = 10 da bir xil sinf: {mos:.1%}")
    print("  bir xil x_T, S = 100 / 20 / 5 / 2 (chapdan o'ngga):")
    x1 = x[:1].clone()
    print(ascii_qator([ddim(m, x1.clone(), S)[0] for S in (100, 20, 5, 2)]))


if __name__ == "__main__":
    main()

  test aniqligi (baholovchi): 0.975

=== 1. DDPM narxi: T ta ketma-ket tarmoq chaqiruvi ===
  model parametrlari: 173120, bitta chaqiruv ~0.35 MFLOP/rasm
  DDPM: T = 100 qadam -> 100 chaqiruv -> ~35 MFLOP/rasm
  qadamlar KETMA-KET: t qadam t+1 natijasini kutadi - parallellab bo'lmaydi

=== 2. Sifat va qadamlar soni (500 namuna, 3 sampling urug'i) ===
  pastki chegara - o'quv rasmlarining o'zi: FD 2.23
  usul              qadam   FD (o'rt +- SE)   ishonch  sinflar>=5%
  DDPM (ajdodli)     100    38.61 +- 0.69      0.855      9.0
  DDIM eta=0         100    21.55 +- 0.99      0.898     10.0
  DDIM eta=0          50    21.62 +- 0.99      0.898     10.0
  DDIM eta=0          20    21.81 +- 1.02      0.897     10.0
  DDIM eta=0          10    22.17 +- 1.03      0.895     10.0
  DDIM eta=0           5    25.87 +- 0.96      0.891      9.7
  DDIM eta=0           3    52.49 +- 1.27      0.876      9.7
  DDIM eta=0           2   269.04 +- 1.00      0.826      8.7
  DDIM eta=1         100    38.61 +- 0.69      0.855      9.0
  DDIM eta=1          20    41.36 +- 1.40      0.857      9.0
  DDIM eta=1          10    47.33 +- 2.14      0.852      9.0
  DDIM eta=1           5    56.91 +- 0.48      0.840      9.0

=== 3. Qaror: eng kam qadam, eng yaxshisidan sezilarli yomon emas ===
  juftlashgan farq (bir xil x_T) eng yaxshisidan (S = 100); amaliy chegara - FD ning 5%
  S =   2: +247.50 (SE 0.50)  statistik: sezilarli       amaliy: yomon
  S =   3:  +30.94 (SE 0.44)  statistik: sezilarli       amaliy: yomon
  S =   5:   +4.32 (SE 0.22)  statistik: sezilarli       amaliy: yomon
  S =  10:   +0.62 (SE 0.05)  statistik: sezilarli       amaliy: maqbul
  S =  20:   +0.27 (SE 0.03)  statistik: sezilarli       amaliy: maqbul
  S =  50:   +0.08 (SE 0.01)  statistik: sezilarli       amaliy: maqbul
  S = 100:   +0.00 (SE 0.00)  statistik: sezilarli emas  amaliy: maqbul
  tanlov: DDIM S = 10 -> 10 barobar kam chaqiruv (~3.5 MFLOP/rasm)
  DDPM (100 qadam) FD 38.61, DDIM S = 10 FD 22.17

=== 4. DDIM deterministik: bir xil x_T -> bir xil rasm ===
  ikki marta, bir xil x_T: max farq 0.0e+00
  S = 50 va S = 10, bir xil x_T:   o'rtacha piksel farqi 0.027
  S = 50, boshqa x_T:               o'rtacha piksel farqi 0.485
  S = 50 va S = 10 da bir xil sinf: 95.5%
  bir xil x_T, S = 100 / 20 / 5 / 2 (chapdan o'ngga):
      =:**       =:**       +-**.      @@ @@
     .**%@..    .**%@..    :##%@:    @@@@@@
     .@..%-     .@..%-     :@..#-     @@   @
       ==#        ==#        =+#         @@ @
     .%@@-      .%@@-      .%@@-.      @@@ @
     =@::-.     =@-:-.     =@+::      @@@   @
     .%  :..    .%  ...    .#  ...   @ @ @  =
    : - %+:    . = #+:    . + *+-    @@@  @@

"""Shartli diffusion va classifier-free guidance: w oshsa aniqlik va xilma-xillik qanday o'zgaradi."""

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import load_digits

T = 100
BOSH = 10                                          # "shartsiz" belgisi (null sinf)
BELGI = " .:-=+*#%@"


def kosinus_jadval(T, s=0.008):
    t = torch.arange(T + 1, dtype=torch.float64) / T
    f = torch.cos((t + s) / (1 + s) * math.pi / 2) ** 2
    return (f / f[0])[1:].clamp(1e-5, 0.9999).float()


AB = kosinus_jadval(T)


def vaqt_emb(t, d=32):
    k = torch.exp(-math.log(1000.0) * torch.arange(d // 2) / (d // 2))
    a = t.float()[:, None] * k[None]
    return torch.cat([a.sin(), a.cos()], 1)


class ShartliShovqin(nn.Module):
    """eps(x_t, t, y): vaqt va SINF embeddingi yashirin qatlamga qo'shiladi; y = 10 - shartsiz."""

    def __init__(self, D=64, h=256):
        super().__init__()
        self.vaqt = nn.Sequential(nn.Linear(32, h), nn.SiLU())
        self.sinf = nn.Embedding(11, h)
        self.kir = nn.Linear(D, h)
        self.tana = nn.Sequential(nn.SiLU(), nn.Linear(h, h), nn.SiLU(), nn.Linear(h, h),
                                  nn.SiLU(), nn.Linear(h, D))

    def forward(self, x, t, y):
        return self.tana(self.kir(x) + self.vaqt(vaqt_emb(t)) + self.sinf(y))


class Klassifikator(nn.Module):
    def __init__(self):
        super().__init__()
        self.xus = nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 32), nn.ReLU())
        self.bosh = nn.Linear(32, 10)

    def forward(self, x):
        return self.bosh(self.xus(x))


def malumot():
    d = load_digits()
    X = torch.tensor(d.data / 8.0 - 1.0, dtype=torch.float32)
    p = torch.tensor(np.random.default_rng(0).permutation(len(X)))
    return X[p], torch.tensor(d.target)[p]


def orgat_diffusion(X, y, p_shartsiz=0.15, qadamlar=2000, B=128, seed=0):
    """Bitta model ikki ish qiladi: shartli va (p_shartsiz ulushda) shartsiz."""
    torch.manual_seed(seed)
    m = ShartliShovqin()
    opt = torch.optim.Adam(m.parameters(), lr=2e-3)
    jadval = torch.optim.lr_scheduler.CosineAnnealingLR(opt, qadamlar)
    g = torch.Generator().manual_seed(seed)
    for _ in range(qadamlar):
        i = torch.randint(0, len(X), (B,), generator=g)
        x0, yb = X[i], y[i].clone()
        yb[torch.rand(B, generator=g) < p_shartsiz] = BOSH       # shartni tashlab yuborish
        t = torch.randint(0, T, (B,), generator=g)
        eps = torch.randn(x0.shape, generator=g)
        ab = AB[t][:, None]
        loss = F.mse_loss(m(ab.sqrt() * x0 + (1 - ab).sqrt() * eps, t, yb), eps)
        opt.zero_grad()
        loss.backward()
        opt.step()
        jadval.step()
    m.eval()
    return m


def orgat_klassifikator(X, y, seed=0):
    torch.manual_seed(seed)
    k = Klassifikator()
    opt = torch.optim.Adam(k.parameters(), lr=3e-3, weight_decay=1e-4)
    g = torch.Generator().manual_seed(seed)
    for _ in range(800):
        i = torch.randint(0, len(X), (128,), generator=g)
        loss = F.cross_entropy(k(X[i] + 0.1 * torch.randn(X[i].shape, generator=g)), y[i])
        opt.zero_grad()
        loss.backward()
        opt.step()
    k.eval()
    return k


def frechet(a, b):
    a, b = a.double().numpy(), b.double().numpy()
    s1, s2 = np.cov(a, rowvar=False), np.cov(b, rowvar=False)
    w, v = np.linalg.eigh(s1)
    ild = (v * np.sqrt(np.clip(w, 0, None))) @ v.T
    oz = np.linalg.eigvalsh(ild @ s2 @ ild)
    return float(((a.mean(0) - b.mean(0)) ** 2).sum() + np.trace(s1) + np.trace(s2)
                 - 2 * np.sqrt(np.clip(oz, 0, None)).sum())


@torch.no_grad()
def cfg_sampling(m, x, y, w, S=20):
    """DDIM (eta = 0) + classifier-free guidance: eps = eps_u + w * (eps_c - eps_u)."""
    qadamlar = torch.linspace(T - 1, 0, S).round().long().tolist()
    bosh = torch.full_like(y, BOSH)
    for i, t in enumerate(qadamlar):
        tt = torch.full((len(x),), t)
        eps_c = m(x, tt, y)
        eps_u = m(x, tt, bosh)
        eps = eps_u + w * (eps_c - eps_u)
        x0 = ((x - (1 - AB[t]).sqrt() * eps) / AB[t].sqrt()).clamp(-1, 1)
        if i + 1 == len(qadamlar):
            return x0
        ab, ab_old = AB[t], AB[qadamlar[i + 1]]
        eps = (x - ab.sqrt() * x0) / (1 - ab).sqrt()
        x = ab_old.sqrt() * x0 + (1 - ab_old).sqrt() * eps
    return x


def sinf_ichi_xilma(f, y):
    """Har sinf ichida juft-juft o'rtacha masofa (xususiyatlar fazosida), sinflar bo'yicha o'rtacha."""
    q = []
    for c in range(10):
        fc = f[y == c]
        d = torch.cdist(fc, fc)
        q.append(d.sum().item() / (len(fc) * (len(fc) - 1)))
    return float(np.mean(q))


def ascii_qator(rasmlar):
    satrlar = []
    for r in range(8):
        qism = []
        for x in rasmlar:
            q = ((x.view(8, 8)[r] + 1) / 2).clamp(0, 1)
            qism.append("".join(BELGI[min(9, int(v * 10))] for v in q.tolist()))
        satrlar.append("    " + "  ".join(qism))
    return "\n".join(satrlar)


def main() -> None:
    torch.set_num_threads(1)
    X, y = malumot()
    Xo, yo, Xt, yt = X[:1200], y[:1200], X[1200:], y[1200:]
    m = orgat_diffusion(Xo, yo)
    k = orgat_klassifikator(Xo, yo)
    with torch.no_grad():
        f_test = k.xus(Xt)
    print("=== 1. Shartli model: sinf embeddingi + shartsiz rejim ===")
    print(f"  sinf embeddingi: 11 ta (0-9 raqamlar, 10 - 'shart yo'q'), o'quvda 15% shart tashlanadi")
    print(f"  baholovchi klassifikator test aniqligi: {(k(Xt).argmax(1) == yt).float().mean():.3f}")
    xil_real = sinf_ichi_xilma(f_test, yt)
    print(f"  haqiqiy test: sinf ichi xilma-xillik {xil_real:.2f} (xususiyatlar fazosida), "
          f"to'yingan piksel (|x| > 0.99) {(Xt.abs() > 0.99).float().mean():.1%}")

    print("\n=== 2. Guidance og'irligi w (har sinfdan 50, 3 urug', DDIM 20 qadam) ===")
    print("  w = 0 - shartsiz, w = 1 - oddiy shartli, w > 1 - kuchaytirilgan")
    print("    w   to'g'ri sinf      FD     xilma-xillik  xilma/haqiqiy  to'yingan piksel")
    yb = torch.arange(10).repeat_interleave(50)
    olchov = {}
    for w in (0.0, 1.0, 2.0, 3.0, 5.0, 8.0):
        q = []
        for u in range(3):
            x = torch.randn((500, 64), generator=torch.Generator().manual_seed(10 + u))
            s = cfg_sampling(m, x, yb, w)
            with torch.no_grad():
                f = k.xus(s)
                tog = (k.bosh(f).argmax(1) == yb).float().mean().item()
            q.append((tog, frechet(f, f_test), sinf_ichi_xilma(f, yb),
                      (s.abs() > 0.99).float().mean().item()))
        q = np.array(q)
        olchov[w] = q
        se = q.std(0, ddof=1) / math.sqrt(3)
        print(f"  {w:>3.0f}   {q[:, 0].mean():.3f}+-{se[0]:.3f}  {q[:, 1].mean():>6.2f}+-{se[1]:.2f}"
              f"  {q[:, 2].mean():>8.2f}    {q[:, 2].mean() / xil_real:>9.2f}  {q[:, 3].mean():>12.1%}")

    print("\n=== 3. Murosa: to'g'rilik oshadi, xilma-xillik kamayadi ===")
    wlar = sorted(olchov)
    for a, b in zip(wlar[:-1], wlar[1:]):
        d_tog = olchov[b][:, 0] - olchov[a][:, 0]
        d_xil = olchov[b][:, 2] - olchov[a][:, 2]
        se_t = d_tog.std(ddof=1) / math.sqrt(3)
        print(f"  w {a:.0f} -> {b:.0f}: to'g'ri sinf {d_tog.mean():+.3f} "
              f"({'sezilarli' if abs(d_tog.mean()) > 2 * se_t else 'sezilarli emas'}), "
              f"xilma-xillik {d_xil.mean():+.2f}")
    eng_fd = min(wlar, key=lambda w: olchov[w][:, 1].mean())
    eng_tog = max(wlar, key=lambda w: olchov[w][:, 0].mean())
    print(f"  eng past FD: w = {eng_fd:.0f};  eng yuqori to'g'ri sinf: w = {eng_tog:.0f}")
    if eng_fd != eng_tog:
        print("  bitta 'eng yaxshi' w yo'q - maqsadga qarab tanlanadi")

    print("\n=== 4. Bir xil x_T, '3' sinfi, w = 0 / 1 / 3 / 8 ===")
    x = torch.randn((4, 64), generator=torch.Generator().manual_seed(5))
    uch = torch.full((4,), 3)
    birinchi = []
    for w in (0.0, 1.0, 3.0, 8.0):
        s = cfg_sampling(m, x.clone(), uch, w)
        with torch.no_grad():
            bashorat = k(s).argmax(1).tolist()
        print(f"  w = {w:.0f}: klassifikator ko'rgan sinflar {bashorat}")
        birinchi.append(s[0])
    print("  1-x_T, chapdan o'ngga w = 0 / 1 / 3 / 8:")
    print(ascii_qator(birinchi))
    print("  hisob narxi: guidance har qadamda 2 chaqiruv (shartli + shartsiz)")


if __name__ == "__main__":
    main()

=== 1. Shartli model: sinf embeddingi + shartsiz rejim ===
  sinf embeddingi: 11 ta (0-9 raqamlar, 10 - 'shart yo'q'), o'quvda 15% shart tashlanadi
  baholovchi klassifikator test aniqligi: 0.975
  haqiqiy test: sinf ichi xilma-xillik 12.23 (xususiyatlar fazosida), to'yingan piksel (|x| > 0.99) 58.0%

=== 2. Guidance og'irligi w (har sinfdan 50, 3 urug', DDIM 20 qadam) ===
  w = 0 - shartsiz, w = 1 - oddiy shartli, w > 1 - kuchaytirilgan
    w   to'g'ri sinf      FD     xilma-xillik  xilma/haqiqiy  to'yingan piksel
    0   0.097+-0.007   38.51+-0.87     19.98         1.63         26.1%
    1   0.830+-0.008   27.90+-0.13     13.21         1.08         24.4%
    2   0.977+-0.003   16.41+-0.40      9.97         0.81         26.4%
    3   0.985+-0.001   19.73+-0.51      9.16         0.75         28.4%
    5   0.985+-0.001   34.37+-1.04      9.05         0.74         32.0%
    8   0.976+-0.006   61.68+-1.59      9.32         0.76         36.4%

=== 3. Murosa: to'g'rilik oshadi, xilma-xillik kamayadi ===
  w 0 -> 1: to'g'ri sinf +0.733 (sezilarli), xilma-xillik -6.76
  w 1 -> 2: to'g'ri sinf +0.147 (sezilarli), xilma-xillik -3.24
  w 2 -> 3: to'g'ri sinf +0.009 (sezilarli), xilma-xillik -0.81
  w 3 -> 5: to'g'ri sinf -0.001 (sezilarli emas), xilma-xillik -0.10
  w 5 -> 8: to'g'ri sinf -0.009 (sezilarli emas), xilma-xillik +0.27
  eng past FD: w = 2;  eng yuqori to'g'ri sinf: w = 3
  bitta 'eng yaxshi' w yo'q - maqsadga qarab tanlanadi

=== 4. Bir xil x_T, '3' sinfi, w = 0 / 1 / 3 / 8 ===
  w = 0: klassifikator ko'rgan sinflar [0, 0, 6, 4]
  w = 1: klassifikator ko'rgan sinflar [3, 3, 3, 3]
  w = 3: klassifikator ko'rgan sinflar [3, 3, 3, 3]
  w = 8: klassifikator ko'rgan sinflar [3, 3, 3, 3]
  1-x_T, chapdan o'ngga w = 0 / 1 / 3 / 8:
    :..=@-    ..-*@:    ..=%@= .  *=-@@@ +
     --*#:     =++*=    :+*-*@-   *#%.#@@
    -#+ .-.   :*- =+ .  :=  %# :  -# .@% %
     #= *=-.   += @#..   .- @* :    % @* *
      @#%#-     %.%#-     . @#.   ::  @*
     -*  #.    ::  *:    .   @=   +-   @=
      -*+#      **:#:    .#+ #*   +--+ =#:
     .+*@+=.   .%#@@+.   .%%@@=     %@@@=.
  hisob narxi: guidance har qadamda 2 chaqiruv (shartli + shartsiz)

"""Inpainting: raqamning yuqori yarmi ma'lum, pastki yarmini diffusion bilan to'ldirish."""

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import load_digits

T = 100
BELGI = " .:-=+*#%@"


def kosinus_jadval(T, s=0.008):
    t = torch.arange(T + 1, dtype=torch.float64) / T
    f = torch.cos((t + s) / (1 + s) * math.pi / 2) ** 2
    return (f / f[0])[1:].clamp(1e-5, 0.9999).float()


AB = kosinus_jadval(T)
ALFA = torch.cat([AB[:1], AB[1:] / AB[:-1]])
BETA = 1 - ALFA


def vaqt_emb(t, d=32):
    k = torch.exp(-math.log(1000.0) * torch.arange(d // 2) / (d // 2))
    a = t.float()[:, None] * k[None]
    return torch.cat([a.sin(), a.cos()], 1)


class Shovqin(nn.Module):
    def __init__(self, D=64, h=256):
        super().__init__()
        self.vaqt = nn.Sequential(nn.Linear(32, h), nn.SiLU())
        self.kir = nn.Linear(D, h)
        self.tana = nn.Sequential(nn.SiLU(), nn.Linear(h, h), nn.SiLU(), nn.Linear(h, h),
                                  nn.SiLU(), nn.Linear(h, D))

    def forward(self, x, t):
        return self.tana(self.kir(x) + self.vaqt(vaqt_emb(t)))


class Klassifikator(nn.Module):
    def __init__(self):
        super().__init__()
        self.xus = nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 32), nn.ReLU())
        self.bosh = nn.Linear(32, 10)

    def forward(self, x):
        return self.bosh(self.xus(x))


def malumot():
    d = load_digits()
    X = torch.tensor(d.data / 8.0 - 1.0, dtype=torch.float32)
    p = torch.tensor(np.random.default_rng(0).permutation(len(X)))
    return X[p], torch.tensor(d.target)[p]


def orgat_diffusion(X, qadamlar=2000, B=128, seed=0):
    """Oddiy SHARTSIZ model - inpainting uchun qayta o'rgatish kerak emas."""
    torch.manual_seed(seed)
    m = Shovqin()
    opt = torch.optim.Adam(m.parameters(), lr=2e-3)
    jadval = torch.optim.lr_scheduler.CosineAnnealingLR(opt, qadamlar)
    g = torch.Generator().manual_seed(seed)
    for _ in range(qadamlar):
        x0 = X[torch.randint(0, len(X), (B,), generator=g)]
        t = torch.randint(0, T, (B,), generator=g)
        eps = torch.randn(x0.shape, generator=g)
        ab = AB[t][:, None]
        loss = F.mse_loss(m(ab.sqrt() * x0 + (1 - ab).sqrt() * eps, t), eps)
        opt.zero_grad()
        loss.backward()
        opt.step()
        jadval.step()
    m.eval()
    return m


def orgat_klassifikator(X, y, seed=0):
    torch.manual_seed(seed)
    k = Klassifikator()
    opt = torch.optim.Adam(k.parameters(), lr=3e-3, weight_decay=1e-4)
    g = torch.Generator().manual_seed(seed)
    for _ in range(800):
        i = torch.randint(0, len(X), (128,), generator=g)
        loss = F.cross_entropy(k(X[i] + 0.1 * torch.randn(X[i].shape, generator=g)), y[i])
        opt.zero_grad()
        loss.backward()
        opt.step()
    k.eval()
    return k


@torch.no_grad()
def inpainting(m, x_malum, niqob, g, har_qadamda=True, r=1):
    """DDPM sampling; har qadamda ma'lum piksellar x_malum ning t-darajali shovqinli nusxasi bilan almashtiriladi.

    har_qadamda=False - 'sodda' usul: shartsiz sampling, faqat oxirida ma'lum qismni yopishtirish.
    r > 1 - qayta namuna olish (RePaint g'oyasi): qadamni qilib, yana shovqin qo'shib, r marta takrorlash.
    """
    x = torch.randn(x_malum.shape, generator=g)
    for t in range(T - 1, -1, -1):
        for j in range(r):
            if har_qadamda:
                shovqinli = AB[t].sqrt() * x_malum + (1 - AB[t]).sqrt() * torch.randn(x.shape, generator=g)
                x = niqob * shovqinli + (1 - niqob) * x
            eps = m(x, torch.full((len(x),), t))
            x0 = ((x - (1 - AB[t]).sqrt() * eps) / AB[t].sqrt()).clamp(-1, 1)
            ab_old = AB[t - 1] if t > 0 else torch.tensor(1.0)
            orta = (ab_old.sqrt() * BETA[t] * x0 + ALFA[t].sqrt() * (1 - ab_old) * x) / (1 - AB[t])
            if t == 0:
                x = orta
                break
            x = orta + (BETA[t] * (1 - ab_old) / (1 - AB[t])).sqrt() * torch.randn(x.shape, generator=g)
            if j + 1 < r:                                  # t-1 -> t: yana bir qadam shovqin
                x = ALFA[t].sqrt() * x + BETA[t].sqrt() * torch.randn(x.shape, generator=g)
    return niqob * x_malum + (1 - niqob) * x


def ascii_qator(rasmlar):
    satrlar = []
    for r in range(8):
        qism = []
        for x in rasmlar:
            q = ((x.view(8, 8)[r] + 1) / 2).clamp(0, 1)
            qism.append("".join(BELGI[min(9, int(v * 10))] for v in q.tolist()))
        satrlar.append("    " + "  ".join(qism))
    return "\n".join(satrlar)


def main() -> None:
    torch.set_num_threads(1)
    X, y = malumot()
    Xo, yo, Xt, yt = X[:1200], y[:1200], X[1200:1400], y[1200:1400]
    m = orgat_diffusion(Xo)
    k = orgat_klassifikator(Xo, yo)
    niqob = torch.zeros(64)
    niqob[:32] = 1.0                                       # yuqori 4 qator ma'lum
    past = niqob == 0

    print("=== 1. Vazifa: 200 test raqami, yuqori 4 qator ma'lum, pastki 4 qator yashirin ===")
    with torch.no_grad():
        print(f"  to'liq rasm: klassifikator aniqligi {(k(Xt).argmax(1) == yt).float().mean():.3f}")
        yarim = Xt * niqob - (1 - niqob)                   # pastki yarim "bo'sh" (-1)
        print(f"  pastki yarim o'chirilgan: aniqlik {(k(yarim).argmax(1) == yt).float().mean():.3f}")

    print("\n=== 2. To'ldirish usullari (har biri 3 marta, turli urug') ===")
    print("  usul                        MSE(past)  to'g'ri sinf  xilma-xillik")
    natija = {}

    def baho(nom, varlar):
        v = torch.stack(varlar)                                # (3, 200, 64)
        mse = ((v[:, :, past] - Xt[None][:, :, past]) ** 2).mean(dim=(1, 2)).numpy()
        with torch.no_grad():
            tog = np.array([(k(s).argmax(1) == yt).float().mean().item() for s in v])
        xil = v[:, :, past].std(0).mean().item()
        natija[nom] = (mse, tog)
        print(f"  {nom:<27} {mse.mean():>8.3f}   {tog.mean():>10.3f}   {xil:>10.3f}")

    ortacha = Xo.mean(0)
    baho("o'rtacha bilan", [niqob * Xt + (1 - niqob) * ortacha] * 3)
    d = torch.cdist(Xt[:, :32], Xo[:, :32])                    # yuqori yarim bo'yicha qo'shni
    yaqin = d.topk(3, largest=False).indices                   # 3 ta eng yaqin
    baho("eng yaqin qo'shni (3 ta)", [niqob * Xt + (1 - niqob) * Xo[yaqin[:, j]] for j in range(3)])
    baho("diffusion: oxirida yopishtir",
         [inpainting(m, Xt, niqob, torch.Generator().manual_seed(u), False) for u in range(3)])
    baho("diffusion: har qadamda",
         [inpainting(m, Xt, niqob, torch.Generator().manual_seed(u)) for u in range(3)])
    varlar = [inpainting(m, Xt, niqob, torch.Generator().manual_seed(u), r=5) for u in range(3)]
    baho("har qadamda + qayta (r=5)", varlar)

    print("\n=== 3. Juftlashgan taqqoslash (3 urug' bo'yicha) ===")
    asos = "har qadamda + qayta (r=5)"
    print(f"  asos: '{asos}' minus boshqa usul")
    for nom in ("o'rtacha bilan", "eng yaqin qo'shni (3 ta)", "diffusion: oxirida yopishtir",
                "diffusion: har qadamda"):
        for j, olchov in ((0, "MSE"), (1, "to'g'ri sinf")):
            f = natija[asos][j] - natija[nom][j]
            se = f.std(ddof=1) / math.sqrt(len(f))
            xulosa = "sezilarli" if abs(f.mean()) > 2 * se + 1e-9 else "sezilarli emas"
            print(f"    {nom:<29} {olchov:<13} {f.mean():+.3f} (SE {se:.3f}) -> {xulosa}")
    print(f"  tarmoq chaqiruvlari: har qadamda {T}, qayta namuna bilan {5 * (T - 1) + 1}")

    print("\n=== 4. Bitta raqam: asl | yarimi | 3 xil to'ldirish (r=5) ===")
    i = int((yt == 5).nonzero()[0])
    with torch.no_grad():
        sinflar = [int(k(v[i:i + 1]).argmax()) for v in varlar[:3]]
    print(f"  asl sinf {int(yt[i])}, to'ldirilganlar klassifikator bo'yicha: {sinflar}")
    print(ascii_qator([Xt[i], Xt[i] * niqob - (1 - niqob)] + [v[i] for v in varlar[:3]]))
    print("  ⭐ Yuqori yarmi bir xil, pastki yarmi har urug'da boshqacha - lekin mos")


if __name__ == "__main__":
    main()

=== 1. Vazifa: 200 test raqami, yuqori 4 qator ma'lum, pastki 4 qator yashirin ===
  to'liq rasm: klassifikator aniqligi 0.975
  pastki yarim o'chirilgan: aniqlik 0.250

=== 2. To'ldirish usullari (har biri 3 marta, turli urug') ===
  usul                        MSE(past)  to'g'ri sinf  xilma-xillik
  o'rtacha bilan                 0.304        0.685        0.000
  eng yaqin qo'shni (3 ta)       0.239        0.882        0.234
  diffusion: oxirida yopishtir    0.473        0.500        0.312
  diffusion: har qadamda         0.354        0.663        0.268
  har qadamda + qayta (r=5)      0.270        0.760        0.208

=== 3. Juftlashgan taqqoslash (3 urug' bo'yicha) ===
  asos: 'har qadamda + qayta (r=5)' minus boshqa usul
    o'rtacha bilan                MSE           -0.034 (SE 0.003) -> sezilarli
    o'rtacha bilan                to'g'ri sinf  +0.075 (SE 0.010) -> sezilarli
    eng yaqin qo'shni (3 ta)      MSE           +0.031 (SE 0.020) -> sezilarli emas
    eng yaqin qo'shni (3 ta)      to'g'ri sinf  -0.122 (SE 0.016) -> sezilarli
    diffusion: oxirida yopishtir  MSE           -0.203 (SE 0.004) -> sezilarli
    diffusion: oxirida yopishtir  to'g'ri sinf  +0.260 (SE 0.010) -> sezilarli
    diffusion: har qadamda        MSE           -0.084 (SE 0.007) -> sezilarli
    diffusion: har qadamda        to'g'ri sinf  +0.097 (SE 0.009) -> sezilarli
  tarmoq chaqiruvlari: har qadamda 100, qayta namuna bilan 496

=== 4. Bitta raqam: asl | yarimi | 3 xil to'ldirish (r=5) ===
  asl sinf 5, to'ldirilganlar klassifikator bo'yicha: [5, 5, 5]
      %@##-     %@##-     %@##-     %@##-     %@##-
     :@++-     :@++-     :@++-     :@++-     :@++-
     =%        =%        =%        =%        =%
     +@%@=     +@%@=     +@%@=     +@%@=     +@%@=
      --+#               .++-+.     -.*=     :#+**.
        #*                 :.=:    .. -#     .. -+:
      .%%                 -=:=-    .#-*#.   . :++=.
     .%*                 .@@=.:     %%+:      #@+:
  ⭐ Yuqori yarmi bir xil, pastki yarmi har urug'da boshqacha - lekin mos

"""Latent diffusion g'oyasi: avtoenkoder latentida diffusion va piksel fazosida diffusion."""

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import load_digits

T = 100


def kosinus_jadval(T, s=0.008):
    t = torch.arange(T + 1, dtype=torch.float64) / T
    f = torch.cos((t + s) / (1 + s) * math.pi / 2) ** 2
    return (f / f[0])[1:].clamp(1e-5, 0.9999).float()


AB = kosinus_jadval(T)


def vaqt_emb(t, d=32):
    k = torch.exp(-math.log(1000.0) * torch.arange(d // 2) / (d // 2))
    a = t.float()[:, None] * k[None]
    return torch.cat([a.sin(), a.cos()], 1)


class Shovqin(nn.Module):
    def __init__(self, D, h):
        super().__init__()
        self.vaqt = nn.Sequential(nn.Linear(32, h), nn.SiLU())
        self.kir = nn.Linear(D, h)
        self.tana = nn.Sequential(nn.SiLU(), nn.Linear(h, h), nn.SiLU(), nn.Linear(h, h),
                                  nn.SiLU(), nn.Linear(h, D))

    def forward(self, x, t):
        return self.tana(self.kir(x) + self.vaqt(vaqt_emb(t)))


class AE(nn.Module):
    """64 piksel -> z (8 o'lcham) -> 64 piksel (26.2-darsdagi avtoenkoder)."""

    def __init__(self, z=8):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, z))
        self.dec = nn.Sequential(nn.Linear(z, 128), nn.ReLU(), nn.Linear(128, 64), nn.Tanh())


class Klassifikator(nn.Module):
    def __init__(self):
        super().__init__()
        self.xus = nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 32), nn.ReLU())
        self.bosh = nn.Linear(32, 10)

    def forward(self, x):
        return self.bosh(self.xus(x))


def malumot():
    d = load_digits()
    X = torch.tensor(d.data / 8.0 - 1.0, dtype=torch.float32)
    p = torch.tensor(np.random.default_rng(0).permutation(len(X)))
    return X[p], torch.tensor(d.target)[p]


def orgat_ae(X, qadamlar=1500, B=128, seed=0):
    torch.manual_seed(seed)
    ae = AE()
    opt = torch.optim.Adam(ae.parameters(), lr=2e-3)
    g = torch.Generator().manual_seed(seed)
    for _ in range(qadamlar):
        x = X[torch.randint(0, len(X), (B,), generator=g)]
        loss = F.mse_loss(ae.dec(ae.enc(x)), x)
        opt.zero_grad()
        loss.backward()
        opt.step()
    ae.eval()
    return ae


def orgat_diffusion(X, h, qadamlar=2000, B=128, seed=0):
    torch.manual_seed(seed)
    m = Shovqin(X.shape[1], h)
    opt = torch.optim.Adam(m.parameters(), lr=2e-3)
    jadval = torch.optim.lr_scheduler.CosineAnnealingLR(opt, qadamlar)
    g = torch.Generator().manual_seed(seed)
    for _ in range(qadamlar):
        x0 = X[torch.randint(0, len(X), (B,), generator=g)]
        t = torch.randint(0, T, (B,), generator=g)
        eps = torch.randn(x0.shape, generator=g)
        ab = AB[t][:, None]
        loss = F.mse_loss(m(ab.sqrt() * x0 + (1 - ab).sqrt() * eps, t), eps)
        opt.zero_grad()
        loss.backward()
        opt.step()
        jadval.step()
    m.eval()
    return m


def orgat_klassifikator(X, y, seed=0):
    torch.manual_seed(seed)
    k = Klassifikator()
    opt = torch.optim.Adam(k.parameters(), lr=3e-3, weight_decay=1e-4)
    g = torch.Generator().manual_seed(seed)
    for _ in range(800):
        i = torch.randint(0, len(X), (128,), generator=g)
        loss = F.cross_entropy(k(X[i] + 0.1 * torch.randn(X[i].shape, generator=g)), y[i])
        opt.zero_grad()
        loss.backward()
        opt.step()
    k.eval()
    return k


def frechet(a, b):
    a, b = a.double().numpy(), b.double().numpy()
    s1, s2 = np.cov(a, rowvar=False), np.cov(b, rowvar=False)
    w, v = np.linalg.eigh(s1)
    ild = (v * np.sqrt(np.clip(w, 0, None))) @ v.T
    oz = np.linalg.eigvalsh(ild @ s2 @ ild)
    return float(((a.mean(0) - b.mean(0)) ** 2).sum() + np.trace(s1) + np.trace(s2)
                 - 2 * np.sqrt(np.clip(oz, 0, None)).sum())


@torch.no_grad()
def ddim(m, x, S=20, chegara=None):
    qadamlar = torch.linspace(T - 1, 0, S).round().long().tolist()
    for i, t in enumerate(qadamlar):
        eps = m(x, torch.full((len(x),), t))
        x0 = (x - (1 - AB[t]).sqrt() * eps) / AB[t].sqrt()
        if chegara is not None:
            x0 = x0.clamp(-chegara, chegara)
        if i + 1 == len(qadamlar):
            return x0
        ab, ab_old = AB[t], AB[qadamlar[i + 1]]
        eps = (x - ab.sqrt() * x0) / (1 - ab).sqrt()
        x = ab_old.sqrt() * x0 + (1 - ab_old).sqrt() * eps
    return x


def main() -> None:
    torch.set_num_threads(1)
    X, y = malumot()
    Xo, yo, Xt = X[:1200], y[:1200], X[1200:]
    k = orgat_klassifikator(Xo, yo)
    with torch.no_grad():
        f_test = k.xus(Xt)

    def baho(s):
        with torch.no_grad():
            f = k.xus(s)
            p = torch.softmax(k.bosh(f), 1)
        sinf = torch.bincount(p.argmax(1), minlength=10).float() / len(s)
        return frechet(f, f_test), p.max(1).values.mean().item(), int((sinf >= 0.05).sum())

    print("=== 1. Avtoenkoder: 64 piksel -> 8 o'lchamli latent ===")
    ae = orgat_ae(Xo)
    with torch.no_grad():
        z = ae.enc(Xo)
        mu, sd = z.mean(0), z.std(0)
        zn = (z - mu) / sd                                   # latentni standartlash
        qayta = ae.dec(ae.enc(Xo[:500]))
    print(f"  test rekonstruksiya MSE: {F.mse_loss(ae.dec(ae.enc(Xt)), Xt).item():.4f}")
    fd_q, ish_q, _ = baho(qayta)
    print(f"  enc -> dec qilingan 500 o'quv rasmi: FD {fd_q:.2f}, ishonch {ish_q:.3f}")
    print("  (latent diffusion bundan yaxshi bo'la olmaydi - dekoder 'shifti')")

    print("\n=== 2. Ikki diffusion, bir xil o'quv qadamlari (2000 x 128) ===")
    piksel = orgat_diffusion(Xo, h=256)
    latent = orgat_diffusion(zn, h=128)
    n_pik = sum(p.numel() for p in piksel.parameters())
    n_lat = sum(p.numel() for p in latent.parameters())
    n_dec = sum(p.numel() for p in ae.dec.parameters())
    S = 20
    print("  model        fazo  parametr  FLOP/qadam   sampling FLOP (S = 20, 1 rasm)")
    print(f"  piksel       64    {n_pik:>8}  {2 * n_pik / 1e3:>7.0f} K   {S * 2 * n_pik / 1e6:>6.2f} M")
    print(f"  latent        8    {n_lat:>8}  {2 * n_lat / 1e3:>7.0f} K   "
          f"{(S * 2 * n_lat + 2 * n_dec) / 1e6:>6.2f} M (dekoder bilan)")
    print(f"  o'quv FLOP nisbati (piksel / latent): ~{n_pik / n_lat:.1f} barobar")

    print("\n=== 3. Sifat (500 namuna, 3 urug') ===")
    fd_o, ish_o, _ = baho(Xo[:500])
    print(f"  o'quv rasmlarining o'zi:  FD {fd_o:.2f}, ishonch {ish_o:.3f}")
    natija = {}
    for nom in ("piksel", "latent"):
        q = []
        for u in range(3):
            g = torch.Generator().manual_seed(50 + u)
            if nom == "piksel":
                s = ddim(piksel, torch.randn((500, 64), generator=g), S, chegara=1.0)
            else:
                zs = ddim(latent, torch.randn((500, 8), generator=g), S)
                with torch.no_grad():
                    s = ae.dec(zs * sd + mu)
            q.append(baho(s))
        q = np.array(q)
        natija[nom] = q
        se = q[:, 0].std(ddof=1) / math.sqrt(3)
        print(f"  {nom:<7} diffusion: FD {q[:, 0].mean():6.2f} +- {se:.2f}, "
              f"ishonch {q[:, 1].mean():.3f}, sinflar>=5%: {q[:, 2].mean():.1f}")
    f = natija["latent"][:, 0] - natija["piksel"][:, 0]
    se = f.std(ddof=1) / math.sqrt(3)
    xulosa = ("latent sezilarli yaxshi" if f.mean() < -2 * se else
              "piksel sezilarli yaxshi" if f.mean() > 2 * se else "sezilarli farq yo'q")
    print(f"  latent - piksel FD: {f.mean():+.2f} (SE {se:.2f}) -> {xulosa}")

    print("\n=== 4. Xulosa (natijadan) ===")
    print(f"  latent model {n_pik / n_lat:.1f} barobar kichik, har qadam shuncha arzon")
    print(f"  latent sifat chegarasi - AE rekonstruksiyasi: FD {fd_q:.2f}")
    if natija["latent"][:, 0].mean() > 2 * fd_q:
        print("  latent diffusion bu chegaradan hali ancha uzoq - xato asosan diffusion qismida")
    else:
        print("  latent diffusion chegaraga yaqin - yaxshilash uchun avval AE ni yaxshilash kerak")


if __name__ == "__main__":
    main()

=== 1. Avtoenkoder: 64 piksel -> 8 o'lchamli latent ===
  test rekonstruksiya MSE: 0.0474
  enc -> dec qilingan 500 o'quv rasmi: FD 4.21, ishonch 0.986
  (latent diffusion bundan yaxshi bo'la olmaydi - dekoder 'shifti')

=== 2. Ikki diffusion, bir xil o'quv qadamlari (2000 x 128) ===
  model        fazo  parametr  FLOP/qadam   sampling FLOP (S = 20, 1 rasm)
  piksel       64      173120      346 K     6.92 M
  latent        8       39432       79 K     1.60 M (dekoder bilan)
  o'quv FLOP nisbati (piksel / latent): ~4.4 barobar

=== 3. Sifat (500 namuna, 3 urug') ===
  o'quv rasmlarining o'zi:  FD 2.23, ishonch 0.999
  piksel  diffusion: FD  20.66 +- 0.21, ishonch 0.897, sinflar>=5%: 9.0
  latent  diffusion: FD  35.48 +- 2.91, ishonch 0.947, sinflar>=5%: 8.3
  latent - piksel FD: +14.82 (SE 2.93) -> piksel sezilarli yaxshi

=== 4. Xulosa (natijadan) ===
  latent model 4.4 barobar kichik, har qadam shuncha arzon
  latent sifat chegarasi - AE rekonstruksiyasi: FD 4.21
  latent diffusion bu chegaradan hali ancha uzoq - xato asosan diffusion qismida

for t in range(T - 1, -1, -10):                       # DDPM formulasi bilan  # ⚠️
    x = ddpm_qadam(model, x, t)
qadamlar = torch.linspace(T - 1, 0, 10).round().long().tolist()                # ✅
x = ddim(model, x, S=10)

loss = F.mse_loss(model(xt, t, y), eps)               # eps_u o'rganilmaydi   # ⚠️
y = y.clone()
y[torch.rand(len(y), generator=g) < 0.15] = BOSH                              # ✅
loss = F.mse_loss(model(xt, t, y), eps)

eps = eps_c + w * (eps_u - eps_c)                     # shartdan UZOQlashadi  # ⚠️
eps = eps_u + w * (eps_c - eps_u)                                             # ✅

w = 15.0                                              # "kuchliroq bo'lsin"   # ⚠️
for w in (1, 2, 3, 5, 8):                                                     # ✅
    olchov[w] = (togri_sinf(w), fd(w), xilma_xillik(w))

x = niqob * x_malum + (1 - niqob) * x                 # t-darajaga mos emas   # ⚠️
x = niqob * (ab[t].sqrt() * x_malum + (1 - ab[t]).sqrt() * z) + (1 - niqob) * x  # ✅

z = ae.enc(X)                                         # masshtab N(0, I) emas # ⚠️
z = (ae.enc(X) - mu) / sd                                                     # ✅
rasm = ae.dec(z_namuna * sd + mu)

fd_10 = fd(ddim(model, torch.randn(500, 64), 10))    # bitta tasodifiy x_T    # ⚠️
farq = [fd(ddim(m, x_u, 10)) - fd(ddim(m, x_u, 100)) for x_u in x_urug]       # ✅

1.  # T = 100 DDPM, model 0.35 MFLOP/chaqiruv - bitta rasm necha MFLOP?
2.  # DDIM uchun modelni qayta o'rgatish kerakmi?
3.  # eta = 0 DDIM da bir xil x_T ikki marta berilsa?
4.  # eta = 1, S = T bo'lgan DDIM nimaga teng?
5.  # CFG formulasi va w = 0, w = 1 ning ma'nosi?
6.  # CFG bilan bitta sampling qadami necha chaqiruv?
7.  # w juda katta bo'lsa xilma-xillik va FD bilan nima bo'ladi?
8.  # inpaintingda ma'lum qism t-qadamda nima bilan almashtiriladi?
9.  # qayta namuna (r = 5, T = 100) necha chaqiruv?
10. # Stable Diffusion latenti 512x512x3 rasmdan necha barobar kichik?
11. # latent diffusion sifatining yuqori chegarasi nima?
12. # nega juftlashgan farqda +0.08 ham "sezilarli" chiqishi mumkin?

1.  eps = model(x, t)
    x0 = (x - (1 - ab[t]).sqrt() * eps) / ab[t].sqrt()
    x = ab[t - 10].sqrt() * x0 + (1 - ab[t - 10]).sqrt() * torch.randn_like(x)

2.  eps = eps_c + w * (eps_c - eps_u)

3.  for t in reversed(range(T)):
        x = niqob * x_malum + (1 - niqob) * x
        x = ddpm_qadam(model, x, t)

4.  z = ae.enc(X)
    model_z = orgat_diffusion(z)
    rasm = ae.dec(ddim(model_z, torch.randn(n, 8), 20))

5.  print("S = 10 yaxshi" if fd(ddim(m, torch.randn(500, 64), 10)) < 20 else "yomon")