Lou (Custom Text-to-Image Diffusion Model)

Modèle de génération d'images 512x512 entièrement codé de zéro en PyTorch (VAE, U-Net, DDPM, Cross-Attention) exploitant un encodeur textuel T5-small.

Gradio Interface & Inference Code

Copiez ce script dans une cellule pour lancer l'interface de génération autonome :

import sys, subprocess, math, torch, torch.nn as nn, torch.nn.functional as F
from PIL import Image

print("Installs..."); subprocess.run([sys.executable, "-m", "pip", "install", "-q", "transformers", "safetensors", "huggingface_hub", "gradio"], check=True)
import gradio as gr
from transformers import T5Tokenizer, T5EncoderModel
from safetensors.torch import load_model
from huggingface_hub import hf_hub_download

class VAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.enc = nn.Sequential(nn.Conv2d(3, 64, 3, 1, 1), nn.Conv2d(64, 128, 4, 2, 1), nn.SiLU(), nn.Conv2d(128, 256, 4, 2, 1), nn.SiLU(), nn.Conv2d(256, 256, 4, 2, 1), nn.SiLU(), nn.Conv2d(256, 8, 3, 1, 1))
        self.dec = nn.Sequential(nn.Conv2d(4, 256, 3, 1, 1), nn.ConvTranspose2d(256, 256, 4, 2, 1), nn.SiLU(), nn.ConvTranspose2d(256, 128, 4, 2, 1), nn.SiLU(), nn.ConvTranspose2d(128, 64, 4, 2, 1), nn.SiLU(), nn.Conv2d(64, 3, 3, 1, 1), nn.Tanh())
    def forward(self, x):
        m, l = torch.chunk(self.enc(x), 2, dim=1)
        return self.dec(m + torch.randn_like(m) * torch.exp(0.5 * torch.clamp(l, -30, 20))), m, l

class CrossAttn(nn.Module):
    def __init__(self, dim, ctx_dim):
        super().__init__()
        self.q, self.k, self.v, self.o = nn.Linear(dim, dim), nn.Linear(ctx_dim, dim), nn.Linear(ctx_dim, dim), nn.Linear(dim, dim)
    def forward(self, x, ctx):
        b, c, h, w = x.shape
        xf = x.permute(0, 2, 3, 1).view(b, h*w, c)
        q, k, v = self.q(xf), self.k(ctx), self.v(ctx)
        a = F.softmax(torch.matmul(q, k.transpose(-1, -2)) * (1.0 / math.sqrt(max(c // 4, 1))), dim=-1)
        out = torch.matmul(a, v).view(b, h, w, c).permute(0, 3, 1, 2)
        return self.o(out.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)

class UNetBlk(nn.Module):
    def __init__(self, inp, out):
        super().__init__()
        self.c1, self.mlp, self.c2, self.at = nn.Conv2d(inp, out, 3, 1, 1), nn.Linear(256, out), nn.Conv2d(out, out, 3, 1, 1), CrossAttn(out, 512)
    def forward(self, x, t, ctx):
        h = F.silu(self.c1(x)) + self.mlp(t)[:, :, None, None]
        return h + self.at(F.silu(self.c2(h)), ctx)

class UNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.te = nn.Sequential(nn.Linear(64, 256), nn.SiLU(), nn.Linear(256, 256))
        self.d1, self.d2, self.u1, self.u2 = UNetBlk(4, 64), UNetBlk(64, 128), UNetBlk(128 + 64, 64), UNetBlk(64 + 64, 4)
    def forward(self, x, t, ctx):
        f = torch.arange(32, device=x.device) * -(math.log(10000) / 31)
        emb = torch.cat([(t[:, None] * f.exp()[None, :]).sin(), (t[:, None] * f.exp()[None, :]).cos()], dim=-1)
        t_emb = self.te(emb)
        x1 = self.d1(x, t_emb, ctx)
        x2 = self.d2(F.max_pool2d(x1, 2), t_emb, ctx)
        x3 = self.u1(torch.cat([F.interpolate(x2, size=x1.shape[-2:], mode="nearest"), x1], dim=1), t_emb, ctx)
        return self.u2(torch.cat([x3, x1], dim=1), t_emb, ctx)

class Lou(nn.Module):
    def __init__(self, tokenizer, text_encoder, device="cuda"):
        super().__init__()
        self.vae, self.unet, self.tok, self.enc, self.dev = VAE(), UNet(), tokenizer, text_encoder, device
    def generate(self, prompt, seed=42, num_of_step=50, num_of_image=1, width=512, height=512, guidance_scale=7.5, negative_prompt=""):
        self.eval(); torch.manual_seed(seed)
        betas = torch.linspace(0.0001, 0.02, num_of_step, device=self.dev)
        alphas, cum = 1.0 - betas, torch.cumprod(1.0 - betas, dim=0)
        prev = F.pad(cum[:-1], (1, 0), value=1.0)
        tk_p = self.tok([prompt], padding="max_length", max_length=32, truncation=True, return_tensors="pt").input_ids.to(self.dev)
        tk_n = self.tok([negative_prompt], padding="max_length", max_length=32, truncation=True, return_tensors="pt").input_ids.to(self.dev)
        with torch.no_grad(): ctx_p, ctx_n = self.enc(tk_p).last_hidden_state, self.enc(tk_n).last_hidden_state
        res = []
        for _ in range(num_of_image):
            z_t = torch.randn((1, 4, height // 8, width // 8), device=self.dev)
            for t in reversed(range(num_of_step)):
                ts = torch.tensor([t], device=self.dev).long()
                with torch.no_grad(): pred_noise = self.unet(z_t, ts, ctx_n) + guidance_scale * (self.unet(z_t, ts, ctx_p) - self.unet(z_t, ts, ctx_n))
                clean = (z_t - (1.0 - cum[t]).sqrt() * pred_noise) / cum[t].sqrt()
                z_t = (alphas[t].sqrt() * (1.0 - prev[t]) / (1.0 - cum[t])) * z_t + (prev[t].sqrt() * betas[t] / (1.0 - cum[t])) * clean + (betas[t] * (1.0 - prev[t]) / (1.0 - cum[t])).sqrt() * (torch.randn_like(z_t) if t > 0 else 0)
            with torch.no_grad(): out_im = self.vae.dec(z_t / 0.18215)
            out_im = torch.clamp((out_im.squeeze(0) + 1.0) / 2.0, 0.0, 1.0).cpu().permute(1, 2, 0).numpy()
            res.append(Image.fromarray((out_im * 255).astype("uint8")))
        return res if num_of_image == 1 else res

dev = "cuda" if torch.cuda.is_available() else "cpu"
tok = T5Tokenizer.from_pretrained("t5-small", legacy=False)
enc = T5EncoderModel.from_pretrained("t5-small").to(dev).eval()
lou = Lou(tok, enc, device=dev).to(dev)
load_model(lou, hf_hub_download(repo_id="lea97338/Lou", filename="model.safetensors"), strict=True)

def predict(p, np, s, g, sd, w, h):
    return lou.generate(p, int(sd), int(s), 1, int(w), int(h), g, np)

gr.Interface(
    fn=predict,
    inputs=[gr.Textbox(value="A beautiful pink flower"), gr.Textbox(value="blurry"), gr.Slider(10, 150, 50), gr.Slider(1.0, 20.0, 7.5), gr.Number(42), gr.Dropdown([256, 512], value=512), gr.Dropdown([256, 512], value=512)],
    outputs=gr.Image(type="pil"),
    title="Lou Text-To-Image"
).launch(share=True)