app.py

# app.py (Corrected Version)

import os
import math
import pickle
import shutil
import subprocess
import sys
import textwrap
import time
from dataclasses import dataclass
from typing import Optional

import spaces
import gradio as gr
import numpy as np
import torch
import torch.nn as nn
from torch.nn import functional as F

# --- One-Time Setup Function ---

def setup_data():
    """
    Checks for dataset metadata and prepares it if missing.
    This involves cloning a repo, running a script, and cleaning up.
    """
    data_dir = 'shakespeare_char'
    meta_path = os.path.join(data_dir, 'meta.pkl')

    if os.path.exists(meta_path):
        print("Dataset metadata found. Skipping setup.")
        return

    print("Dataset metadata not found. Starting one-time setup...")
    print("This may take a minute...")

    repo_url = "https://github.com/karpathy/nanoGPT"
    repo_dir = "nanoGPT"

    try:
        print(f"Cloning {repo_url}...")
        subprocess.run(["git", "clone", repo_url], check=True, capture_output=True)

        source_data_dir = os.path.join(repo_dir, 'data', 'shakespeare_char')
        print(f"Copying data from {source_data_dir} to {data_dir}...")
        shutil.copytree(source_data_dir, data_dir)

        prepare_script_path = os.path.join(data_dir, 'prepare.py')
        print(f"Running {prepare_script_path} to generate metadata...")
        subprocess.run([sys.executable, prepare_script_path], check=True, capture_output=True)

        print("Setup successful. 'meta.pkl' has been created.")

    except subprocess.CalledProcessError as e:
        print(f"An error occurred during setup: {e}", file=sys.stderr)
        print(f"Stdout: {e.stdout.decode()}", file=sys.stderr)
        print(f"Stderr: {e.stderr.decode()}", file=sys.stderr)
        sys.exit("Setup failed. Please check your git installation and internet connection.")
    except Exception as e:
        print(f"An unexpected error occurred: {e}", file=sys.stderr)
        sys.exit("Setup failed.")
    finally:
        if os.path.exists(repo_dir):
            print(f"Cleaning up by removing '{repo_dir}' directory...")
            shutil.rmtree(repo_dir)

# --- Run Setup and Load Data ---
setup_data()

# Load metadata for character mappings
data_dir = './shakespeare_char/'
meta_path = os.path.join(data_dir, 'meta.pkl')
with open(meta_path, 'rb') as f:
    meta = pickle.load(f)

itos = meta['itos']
stoi = meta['stoi']
vocab_size = meta['vocab_size']
CONTEXT_LENGTH = 256

def decode(indices_tensor: torch.Tensor):
    if indices_tensor.dim() == 2:
        indices_tensor = indices_tensor[0]
    indices = indices_tensor.cpu().numpy()
    return ''.join([itos[i] for i in indices])

def wrap_text(long_text, width=80):
    paragraphs = long_text.splitlines()
    wrapped = [textwrap.fill(p, width=width) if p else '' for p in paragraphs]
    return "\n".join(wrapped)

# --- Model Architecture ---

class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
        self.gelu    = nn.GELU()
        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
        self.dropout = nn.Dropout(config.dropout)
    def forward(self, x):
        return self.dropout(self.c_proj(self.gelu(self.c_fc(x))))

class SelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.dropout = config.dropout
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
    def forward(self, x):
        B, T, C = x.size()
        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        if self.flash:
            y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=False)
        else:
            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
            att = F.softmax(att, dim=-1)
            att = self.attn_dropout(att)
            y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.c_proj(y))

def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    return x * (1 + scale) + shift

def bias_add_scale(x: torch.Tensor, bias: Optional[torch.Tensor], scale: torch.Tensor, residual: Optional[torch.Tensor]) -> torch.Tensor:
    out = scale * (x + bias) if bias is not None else scale * x
    return residual + out if residual is not None else out

class DDiTBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd, bias=config.bias)
        self.attn = SelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd, bias=config.bias)
        self.mlp = MLP(config)
        self.adaLN_modulation = nn.Linear(config.cond_dim, 6 * config.n_embd, bias=True)
        self.adaLN_modulation.weight.data.zero_()
        self.adaLN_modulation.bias.data.zero_()
    def forward(self, x, c):
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c)[:, None].chunk(6, dim=2)
        x_skip = x
        x = modulate(self.ln_1(x), shift_msa, scale_msa)
        x = self.attn(x)
        x = bias_add_scale(x, None, gate_msa, x_skip)
        x = bias_add_scale(self.mlp(modulate(self.ln_2(x), shift_mlp, scale_mlp)), None, gate_mlp, x)
        return x

class DDitFinalLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.norm_final = nn.LayerNorm(config.n_embd, bias=config.bias)
        self.linear = nn.Linear(config.n_embd, config.vocab_size)
        self.linear.weight.data.zero_()
        self.linear.bias.data.zero_()
        self.adaLN_modulation = nn.Linear(config.cond_dim, 2 * config.n_embd)
        self.adaLN_modulation.weight.data.zero_()
        self.adaLN_modulation.bias.data.zero_()
    def forward(self, x, c):
        shift, scale = self.adaLN_modulation(c)[:, None].chunk(2, dim=2)
        x = modulate(self.norm_final(x), shift, scale)
        return self.linear(x)

class TimestepEmbedder(nn.Module):
    def __init__(self, hidden_size, frequency_embedding_size=256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size, bias=True),
        )
        self.frequency_embedding_size = frequency_embedding_size
    @staticmethod
    def timestep_embedding(t, dim, max_period=10000):
        half = dim // 2
        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(device=t.device)
        args = t[:, None].float() * freqs[None]
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        if dim % 2:
            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
        return embedding
    def forward(self, t):
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
        return self.mlp(t_freq)

class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.sigma_map = TimestepEmbedder(config.cond_dim)
        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd),
            wpe = nn.Embedding(config.block_size, config.n_embd),
            drop = nn.Dropout(config.dropout),
            h = nn.ModuleList([DDiTBlock(config) for _ in range(config.n_layer)]),
            ln_f = nn.LayerNorm(config.n_embd, bias=config.bias), # <<< FIX 1: ADDED THIS LAYER
        ))
        self.lm_head = DDitFinalLayer(config)
        self.apply(self._init_weights)

        # Apply special scaled init to the residual projections
        for pn, p in self.named_parameters():
            if pn.endswith('c_proj.weight'):
                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, sigma):
        sigma = sigma.reshape(-1)
        b, t = idx.size()
        c = F.silu(self.sigma_map(sigma))
        pos = torch.arange(0, t, dtype=torch.long, device=idx.device)
        tok_emb = self.transformer.wte(idx)
        pos_emb = self.transformer.wpe(pos)
        x = self.transformer.drop(tok_emb + pos_emb)
        for block in self.transformer.h:
            x = block(x, c)
        x = self.transformer.ln_f(x) # <<< FIX 2: CALLED THE LAYER HERE
        x = self.lm_head(x, c)
        return torch.scatter(x, -1, idx[..., None], torch.zeros_like(x[..., :1]))

@dataclass
class GPTConfig:
    block_size: int = 1024
    vocab_size: int = 50304
    n_layer: int = 12
    n_head: int = 12
    n_embd: int = 768
    cond_dim: int = 64
    dropout: float = 0.0
    bias: bool = False

# --- Noise Schedule & Sampling Logic ---

class GeometricNoise:
    def __init__(self, sigma_min=1e-4, sigma_max=20):
        self.sigmas = 1.0 * torch.tensor([sigma_min, sigma_max])
    def total_noise(self, t):
        return self.sigmas[0] ** (1 - t) * self.sigmas[1] ** t
    def __call__(self, t):
        return self.total_noise(t), None # Rate not needed for sampling

def transition(x_t: torch.Tensor, delta_sigma: torch.Tensor) -> torch.Tensor:
    base_prob = (1 - torch.exp(-delta_sigma[..., None])) / vocab_size
    trans = torch.ones(*x_t.shape, vocab_size, device=x_t.device) * base_prob
    trans = trans.scatter(-1, x_t[..., None], torch.zeros_like(trans))
    diag_fill = 1 - trans.sum(dim=-1, keepdim=True)
    return trans.scatter(-1, x_t[..., None], diag_fill)

def staggered_score(score, delta_sigma):
    exp_factor = torch.exp(-delta_sigma)[..., None]
    correction = ((exp_factor - 1) / (vocab_size * exp_factor)) * score.sum(dim=-1, keepdim=True)
    return correction + score / exp_factor

def sample_categorical(probs: torch.Tensor) -> torch.Tensor:
    eps = 1e-10
    gumbel_noise = -torch.log(-torch.log(torch.rand_like(probs) + eps) + eps)
    return torch.argmax(torch.log(probs + eps) + gumbel_noise, dim=-1)

# --- Global Model Loading ---

print("Setting up model and device...")
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"===================================")
print(f"Using device: {DEVICE}")
print(f"===================================")

model_args = dict(n_layer=6, n_head=6, n_embd=384, cond_dim=64,
                  bias=False, vocab_size=vocab_size, block_size=CONTEXT_LENGTH, dropout=0.2)
config = GPTConfig(**model_args)
model = GPT(config)

print("Loading pre-trained model weights...")
model.load_state_dict(
    torch.hub.load_state_dict_from_url(
        'https://raw.githubusercontent.com/ash80/diffusion-gpt/master/pretrained_model/model_epoch_25.pth',
        map_location=DEVICE
    )
)
model.to(DEVICE)
model.eval()

NOISE = GeometricNoise(sigma_min=1e-4, sigma_max=20)
print("Model setup complete. Launching Gradio demo...")

# --- Gradio Generation Function ---
@spaces.GPU
def generate_text(steps):
    steps = int(steps)
    eps = 1e-5
    timesteps = torch.linspace(1, eps, steps + 1, device=DEVICE)
    step_size = (1 - eps) / steps

    x = torch.randint(0, vocab_size, (1, CONTEXT_LENGTH), device=DEVICE)

    initial_text = decode(x)
    yield f"Step 0/{steps} (Initial Noise):\n\n{wrap_text(initial_text)}"
    time.sleep(0.5)

    with torch.no_grad():
        for i in range(steps):
            progress(i / steps, desc=f"Denoising Step {i+1}/{steps}")

            t = timesteps[i] * torch.ones(x.shape[0], 1, device=DEVICE)
            curr_sigma_bar, _ = NOISE(t)

            next_t = t - step_size
            next_sigma_bar, _ = NOISE(next_t)
            delta_sigma = curr_sigma_bar - next_sigma_bar

            log_score = model(x, curr_sigma_bar)
            score = torch.exp(log_score)

            stag_score = staggered_score(score, delta_sigma)
            probs = stag_score * transition(x, delta_sigma)
            x = sample_categorical(probs)

            decoded_text = decode(x)
            yield f"Step {i+1}/{steps}:\n\n{wrap_text(decoded_text)}"

    final_text = decode(x)
    yield f"Final Result (Step {steps}/{steps}):\n\n{wrap_text(final_text)}"

# --- Gradio Interface ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # The Annotated Discrete Diffusion Model: Live Demo
        This demo visualizes the denoising process of a character-level discrete diffusion model.
        Start with pure random noise and watch as coherent text, in the style of Shakespeare, emerges over several steps.
        """
    )
    with gr.Row():
        steps_slider = gr.Slider(
            minimum=10,
            maximum=200,
            value=128,
            step=1,
            label="Number of Denoising Steps",
            info="More steps can lead to better quality but take longer."
        )
        generate_button = gr.Button("Generate", variant="primary")

    output_textbox = gr.Textbox(
        label="Denoising Process",
        lines=15,
        interactive=False,
        show_copy_button=True,
        placeholder="The denoising process will appear here..."
    )

    generate_button.click(
        fn=generate_text,
        inputs=[steps_slider],
        outputs=[output_textbox]
    )

if __name__ == "__main__":
    demo.launch()