Update app.py

app.py

@@ -1,798 +1,798 @@
-import gradio as gr
-import spaces
-import torch
-import torch.nn as nn
-from torch.nn import functional as F
-import numpy as np
-import math
-import os
-import pickle
-import requests
-import textwrap
-import subprocess
-import shutil
-import time
-from dataclasses import dataclass
-from typing import Optional
-from transformers import AutoTokenizer
-
-# ==============================================================================
-# ------------------------- VERSION 1: SHARED SETUP ----------------------------
-# ==============================================================================
-
-def setup_environment():
-    """Checks for and sets up the necessary data for V1."""
-    nano_gpt_repo_path = 'nanoGPT'
-    data_dir_path = 'shakespeare_char'
-    meta_path = os.path.join(data_dir_path, 'meta.pkl')
-
-    if os.path.exists(meta_path):
-        return
-
-    print("Required data not found. Starting one-time setup...")
-    if not os.path.exists(nano_gpt_repo_path):
-        try:
-            subprocess.run(['git', 'clone', 'https://github.com/karpathy/nanoGPT.git'], check=True, capture_output=True, text=True)
-        except subprocess.CalledProcessError as e:
-            print(f"Error cloning repository: {e.stderr}")
-            pass 
-
-    source_data_dir = os.path.join(nano_gpt_repo_path, 'data', 'shakespeare_char')
-    if not os.path.exists(data_dir_path) and os.path.exists(source_data_dir):
-        shutil.copytree(source_data_dir, data_dir_path)
-    
-    # Check if we can run prepare
-    prepare_script_path = os.path.join(data_dir_path, 'prepare.py')
-    if os.path.exists(prepare_script_path) and not os.path.exists(meta_path):
-        subprocess.run(['python', 'prepare.py'], check=True, cwd=data_dir_path, capture_output=True, text=True)
-
-setup_environment()
-
-def download_file(url, filename):
-    if os.path.exists(filename):
-        return
-    print(f"Downloading '{filename}'...")
-    try:
-        response = requests.get(url, stream=True)
-        response.raise_for_status()
-        with open(filename, 'wb') as f:
-            for chunk in response.iter_content(chunk_size=8192):
-                f.write(chunk)
-    except requests.exceptions.RequestException as e:
-        print(f"Error downloading {url}: {e}")
-
-# ==============================================================================
-# ---------------------- VERSION 1: ARCHITECTURE & LOGIC -----------------------
-# ==============================================================================
-
-# V1 Constants and Meta Loading
-v1_data_dir = './shakespeare_char/'
-v1_meta_url = 'https://huggingface.co/spaces/thejagstudio/diffusion-gpt/resolve/main/meta.pkl'
-v1_meta_path = 'meta.pkl'
-download_file(v1_meta_url, v1_meta_path)
-
-v1_vocab_size = 65 # Fallback
-v1_itos = {}
-v1_stoi = {}
-
-if os.path.exists(v1_meta_path):
-    with open(v1_meta_path, 'rb') as f:
-        meta = pickle.load(f)
-        v1_vocab_size = meta['vocab_size']
-        v1_itos = meta['itos']
-        v1_stoi = meta['stoi']
-
-v1_context_length = 256
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-def v1_decode(indices_tensor: torch.Tensor):
-    if indices_tensor.dim() > 1:
-        indices_tensor = indices_tensor.squeeze(0)
-    indices = indices_tensor.cpu().numpy()
-    return ''.join([v1_itos.get(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)
-
-@dataclass
-class V1_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
-
-class V1_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):
-        x = self.c_fc(x)
-        x = self.gelu(x)
-        x = self.c_proj(x)
-        x = self.dropout(x)
-        return x
-
-class V1_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 = torch.nn.functional.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)
-        y = self.resid_dropout(self.c_proj(y))
-        return y
-
-def v1_modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
-    return x * (1 + scale) + shift
-
-def v1_bias_add_scale(x: torch.Tensor, bias: Optional[torch.Tensor], scale: torch.Tensor, residual: Optional[torch.Tensor]) -> torch.Tensor:
-    if bias is not None:
-        out = scale * (x + bias)
-    else:
-        out = scale * x
-    if residual is not None:
-        out = residual + out
-    return out
-
-class V1_DDiTBlock(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.ln_1 = nn.LayerNorm(config.n_embd, bias=config.bias)
-        self.attn = V1_SelfAttention(config)
-        self.ln_2 = nn.LayerNorm(config.n_embd, bias=config.bias)
-        self.mlp = V1_MLP(config)
-        self.adaLN_modulation = nn.Linear(config.cond_dim, 6 * config.n_embd)
-        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 = v1_modulate(self.ln_1(x), shift_msa, scale_msa)
-        x = self.attn(x)
-        x = v1_bias_add_scale(self.attn(self.ln_1(x)), None, gate_msa, x_skip)
-        x = v1_bias_add_scale(self.mlp(v1_modulate(self.ln_2(x), shift_mlp, scale_mlp)), None, gate_mlp, x)
-        return x
-
-class V1_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 = v1_modulate(self.norm_final(x), shift, scale)
-        x = self.linear(x)
-        return x
-
-class V1_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)
-        t_emb = self.mlp(t_freq)
-        return t_emb
-
-class V1_GPT(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        assert config.vocab_size is not None
-        assert config.block_size is not None
-        self.config = config
-        self.sigma_map = V1_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([V1_DDiTBlock(config) for _ in range(config.n_layer)]),
-            ln_f = nn.LayerNorm(config.n_embd, bias=config.bias),
-        ))
-        self.lm_head = V1_DDitFinalLayer(config)
-        self.apply(self._init_weights)
-        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))
-        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
-        pos = torch.arange(0, t, dtype=torch.long, device=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)
-        x = self.lm_head(x, c)
-        x = torch.scatter(x, -1, idx[..., None], torch.zeros_like(x[..., :1]))
-        return x
-
-class V1_GeometricNoise:
-    def __init__(self, sigma_min=1e-4, sigma_max=20):
-        self.sigmas = 1.0 * torch.tensor([sigma_min, sigma_max]).to(device)
-    def rate_noise(self, t):
-        return self.sigmas[0] ** (1 - t) * self.sigmas[1] ** t * (self.sigmas[1].log() - self.sigmas[0].log())
-    def total_noise(self, t):
-        return self.sigmas[0] ** (1 - t) * self.sigmas[1] ** t
-    def __call__(self, t):
-        return self.total_noise(t), self.rate_noise(t)
-
-# --- V1 Inference Logic ---
-def v1_transition(x_t: torch.Tensor, delta_sigma: torch.Tensor) -> torch.Tensor:
-    base_prob = (1 - torch.exp(-delta_sigma[..., None])) / v1_vocab_size
-    trans = torch.ones(*x_t.shape, v1_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)
-    trans = trans.scatter(-1, x_t[..., None], diag_fill)
-    return trans
-
-def v1_staggered_score(score, delta_sigma):
-    exp_factor = torch.exp(-delta_sigma)[..., None]
-    correction = ((exp_factor - 1) / (v1_vocab_size * exp_factor)) * score.sum(dim=-1, keepdim=True)
-    return correction + score / exp_factor
-
-def v1_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)
-
-# --- V1 Model Loading ---
-print("Initializing V1 Model...")
-v1_model_args = dict(n_layer=6, n_head=6, n_embd=384, cond_dim=64,
-                 bias=False, vocab_size=v1_vocab_size, block_size=v1_context_length, dropout=0.2)
-v1_config = V1_GPTConfig(**v1_model_args)
-v1_model = V1_GPT(v1_config)
-try:
-    v1_model.load_state_dict(
-        torch.hub.load_state_dict_from_url(
-            'https://huggingface.co/spaces/thejagstudio/diffusion-gpt/resolve/main/final_model.pth?download=true',
-            map_location=device
-        )
-    )
-    v1_model.to(device)
-    v1_model.eval()
-    print("V1 Model loaded successfully.")
-except Exception as e:
-    print(f"Failed to load V1 model: {e}")
-    v1_model = None
-
-v1_noise = V1_GeometricNoise(sigma_min=1e-4, sigma_max=20)
-
-
-def v1_generate_stream(steps, speed):
-    """
-    Generator function for V1 that yields frames directly.
-    Combined logic of generation and replay to allow for immediate stopping.
-    """
-    if v1_model is None:
-        yield "Error: V1 Model not loaded"
-        return
-        
-    steps = int(steps)
-    speed = float(speed)
-    eps = 1e-5
-    
-    # Calculate delay based on speed slider (similar to V2)
-    # 0.5 is base constant, speed scales it down
-    delay = 0.5 / max(speed, 0.1)
-
-    x = torch.randint(0, v1_vocab_size, (1, v1_context_length), device=device)
-    initial_text = f"--- Initial Random Noise ---\n\n{wrap_text(v1_decode(x[0]))}"
-    yield initial_text
-    time.sleep(delay)
-
-    timesteps = torch.linspace(1, eps, steps + 1, device=device)
-    step_size = (1 - eps) / steps
-    
-    with torch.no_grad():
-        for i in range(steps):
-            t = timesteps[i] * torch.ones(x.shape[0], 1, device=device)
-            curr_sigma_bar = v1_noise(t)[0]
-            
-            next_sigma_bar = v1_noise(t - step_size)[0]
-            delta_sigma = curr_sigma_bar - next_sigma_bar
-
-            log_score = v1_model(x, curr_sigma_bar)
-            score = torch.exp(log_score)
-
-            stag_score = v1_staggered_score(score, delta_sigma)
-            probs = stag_score * v1_transition(x, delta_sigma)
-            x = v1_sample_categorical(probs)
-            
-            progress_text = f"--- Denoising Step {i + 1}/{steps} ---\n\n{wrap_text(v1_decode(x[0]))}"
-            yield progress_text
-            
-            # Artificial delay for visualization
-            if speed < 20:
-                time.sleep(delay)
-            
-        t = timesteps[steps] * torch.ones(x.shape[0], 1, device=device)
-        curr_sigma_bar = v1_noise(t)[0]
-        delta_sigma = curr_sigma_bar
-
-        log_score = v1_model(x, curr_sigma_bar)
-        score = torch.exp(log_score)
-        stag_score = v1_staggered_score(score, delta_sigma)
-        probs = stag_score * v1_transition(x, delta_sigma)
-        x = v1_sample_categorical(probs)
-
-    final_text = f"--- Final Denoised Text (Step {steps}) ---\n\n{wrap_text(v1_decode(x[0]))}"
-    yield final_text
-
-# ==============================================================================
-# ---------------------- VERSION 2: ARCHITECTURE & LOGIC -----------------------
-# ==============================================================================
-
-# PLEASE UPDATE THIS PATH TO YOUR ACTUAL LOCAL FILE OR URL
-V2_MODEL_PATH = "checkpoints/model_fp32.pt"
-
-class V2_RMSNorm(nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-    
-    def forward(self, x):
-        var = x.pow(2).mean(-1, keepdim=True)
-        x = x * torch.rsqrt(var + self.eps)
-        return self.weight * x
-
-class V2_RotaryEmbedding(nn.Module):
-    def __init__(self, dim, max_position_embeddings=16384, base=100000, scaling_factor=1.0):
-        super().__init__()
-        self.scaling_factor = scaling_factor
-        self.dim = dim
-        self.base = base
-        self.max_position_embeddings = max_position_embeddings
-        self.inv_freq = None
-        self._cache = {}
-    
-    def _update_freqs(self, device):
-        base = self.base * (self.scaling_factor ** (self.dim / (self.dim - 2)))
-        inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
-        self.inv_freq = inv_freq
-    
-    def forward(self, x, seq_len=None):
-        if seq_len is None:
-            seq_len = x.shape[-2]
-        
-        if self.inv_freq is None or self.inv_freq.device != x.device:
-            self._update_freqs(x.device)
-        
-        cache_key = (seq_len, x.device, x.dtype)
-        if cache_key in self._cache:
-            return self._cache[cache_key]
-        
-        t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype)
-        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
-        emb = torch.cat((freqs, freqs), dim=-1)
-        
-        cos = emb.cos()[None, None, :, :]
-        sin = emb.sin()[None, None, :, :]
-        
-        self._cache[cache_key] = (cos, sin)
-        if len(self._cache) > 10:
-            self._cache.pop(next(iter(self._cache)))
-        
-        return cos, sin
-
-def v2_apply_rotary_pos_emb(q, k, cos, sin):
-    def rotate_half(x):
-        x1 = x[..., : x.shape[-1] // 2]
-        x2 = x[..., x.shape[-1] // 2 :]
-        return torch.cat((-x2, x1), dim=-1)
-    q_embed = (q * cos) + (rotate_half(q) * sin)
-    k_embed = (k * cos) + (rotate_half(k) * sin)
-    return q_embed, k_embed
-
-class V2_DiffusionAttention(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-        self.num_heads = config.num_attention_heads
-        self.head_dim = self.hidden_size // self.num_heads
-        self.num_key_value_heads = config.num_key_value_heads
-        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
-        self.use_flash_attn = config.use_flash_attn
-        
-        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
-        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
-        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
-        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
-    
-    def forward(self, hidden_states, freqs_cis, attention_mask=None, past_kv=None):
-        bsz, q_len, _ = hidden_states.size()
-        
-        q = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        k = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-        v = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-        
-        cos, sin = freqs_cis
-        cos = cos[:, :, :q_len, :]
-        sin = sin[:, :, :q_len, :]
-        q, k = v2_apply_rotary_pos_emb(q, k, cos, sin)
-        
-        if past_kv is not None:
-            cache_k, cache_v = past_kv
-            k = torch.cat([cache_k, k], dim=2)
-            v = torch.cat([cache_v, v], dim=2)
-        
-        current_kv = (k, v)
-        k = k.repeat_interleave(self.num_key_value_groups, dim=1)
-        v = v.repeat_interleave(self.num_key_value_groups, dim=1)
-        
-        attn_mask = None
-        if attention_mask is not None:
-            attn_mask = attention_mask[:, None, None, :].to(dtype=q.dtype)
-            attn_mask = (1.0 - attn_mask) * torch.finfo(q.dtype).min
-        
-        output = F.scaled_dot_product_attention(
-            q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False
-        )
-        
-        output = output.transpose(1, 2).contiguous().view(bsz, q_len, self.hidden_size)
-        return self.o_proj(output), current_kv
-
-class V2_MLP(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
-        self.act_fn = nn.SiLU()
-    
-    def forward(self, x):
-        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
-
-class V2_BlockDiffusionBlock(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.self_attn = V2_DiffusionAttention(config)
-        self.mlp = V2_MLP(config)
-        self.input_layernorm = V2_RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = V2_RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.use_activation_checkpointing = config.use_activation_checkpointing
-    
-    def forward(self, hidden_states, freqs_cis, attention_mask, past_kv):
-        return self._forward(hidden_states, freqs_cis, attention_mask, past_kv)
-    
-    def _forward(self, hidden_states, freqs_cis, attention_mask, past_kv):
-        residual = hidden_states
-        hidden_states = self.input_layernorm(hidden_states)
-        attn_out, new_kv = self.self_attn(hidden_states, freqs_cis, attention_mask, past_kv)
-        hidden_states = residual + attn_out
-        
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = residual + self.mlp(hidden_states)
-        return hidden_states, new_kv
-
-@dataclass
-class V2_ModelConfig:
-    vocab_size: int = 151936
-    hidden_size: int = 1024
-    intermediate_size: int = 2816
-    num_hidden_layers: int = 16
-    num_attention_heads: int = 16
-    num_key_value_heads: int = 4
-    max_position_embeddings: int = 16384
-    rms_norm_eps: float = 1e-6
-    rope_theta: float = 100000.0
-    pad_token_id: int = 0
-    mask_token_id: int = 1
-    use_flash_attn: bool = True
-    use_activation_checkpointing: bool = False
-    attention_dropout: float = 0.0
-    hidden_dropout: float = 0.0
-
-ModelConfig = V2_ModelConfig
-
-class V2_DiffusionLLM(nn.Module):
-    def __init__(self, config: V2_ModelConfig):
-        super().__init__()
-        self.config = config
-        pad_idx = config.pad_token_id if config.pad_token_id < config.vocab_size else None
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=pad_idx)
-        
-        self.layers = nn.ModuleList([V2_BlockDiffusionBlock(config) for _ in range(config.num_hidden_layers)])
-        self.norm = V2_RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self.rotary_emb = V2_RotaryEmbedding(
-            config.hidden_size // config.num_attention_heads,
-            config.max_position_embeddings
-        )
-        self.lm_head.weight = self.embed_tokens.weight
-    
-    def forward(self, input_ids, attention_mask=None, past_key_values=None):
-        bsz, seqlen = input_ids.shape
-        hidden_states = self.embed_tokens(input_ids)
-        freqs_cis = self.rotary_emb(hidden_states, seq_len=seqlen)
-        
-        if past_key_values is None:
-            past_key_values = [None] * len(self.layers)
-        
-        new_kvs = []
-        for i, layer in enumerate(self.layers):
-            hidden_states, kv = layer(hidden_states, freqs_cis, attention_mask, past_key_values[i])
-            new_kvs.append(kv)
-        
-        hidden_states = self.norm(hidden_states)
-        logits = self.lm_head(hidden_states)
-        return logits, new_kvs
-
-DiffusionLLM = V2_DiffusionLLM
-
-# --- V2 Loading Logic ---
-print("Initializing V2 components...")
-v2_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
-if v2_tokenizer.pad_token is None:
-    v2_tokenizer.pad_token = v2_tokenizer.eos_token
-
-v2_model = None
-v2_config = None
-
-if os.path.exists(V2_MODEL_PATH):
-    print(f"Loading V2 model from {V2_MODEL_PATH}...")
-    try:
-        checkpoint = torch.load(V2_MODEL_PATH, map_location=device, weights_only=False)
-        v2_config = checkpoint['config']
-        v2_model = V2_DiffusionLLM(v2_config)
-        state_dict = checkpoint['model_state']
-        state_dict = {k: v.float() for k, v in state_dict.items()}
-        v2_model.load_state_dict(state_dict)
-        v2_model = v2_model.to(device)
-        v2_model.eval()
-        print("V2 Model loaded.")
-    except Exception as e:
-        print(f"Error loading V2 model: {e}")
-else:
-    print(f"V2 Model file not found at {V2_MODEL_PATH}. Version 2 tab will not work without it.")
-
-
-@torch.no_grad()
-def v2_generate_block_diffusion(prompt, steps, block_size, max_new_tokens, replay_speed):
-    """
-    Refactored to yield frames for real-time streaming.
-    """
-    if v2_model is None:
-        yield "Error: V2 Model not found. Check path."
-        return
-
-    v2_model.eval()
-    # Handle inputs
-    steps = int(steps)
-    block_size = int(block_size)
-    max_new_tokens = int(max_new_tokens)
-    speed = float(replay_speed)
-
-    prompt_ids = v2_tokenizer.encode(prompt, return_tensors="pt").to(device)
-    config = v2_model.config
-    num_blocks = max_new_tokens // block_size
-
-    context_ids = prompt_ids
-
-    # Helper params
-    temperature = 1.0
-    top_k = 40
-    top_p = 0.9
-    repetition_penalty = 1.2
-    
-    # Calculate delay based on speed slider
-    delay = 0.5 / max(speed, 0.1) 
-
-    for block_idx in range(num_blocks):
-        mask_block = torch.full((1, block_size), config.mask_token_id, device=device)
-        is_masked = torch.ones(1, block_size, dtype=torch.bool, device=device)
-        
-        for step_idx in range(steps):
-            # --- SNAPSHOT & YIELD ---
-            # Decode context
-            ctx_str = v2_tokenizer.decode(context_ids[0], skip_special_tokens=True)
-            
-            # Decode block with masking visual
-            block_tokens = mask_block[0].tolist()
-            block_vis = []
-            for i, tid in enumerate(block_tokens):
-                if is_masked[0, i]:
-                    block_vis.append("░") # Mask symbol
-                else:
-                    block_vis.append(v2_tokenizer.decode([tid], skip_special_tokens=False))
-            
-            block_str = "".join(block_vis)
-            
-            frame_text = (f"--- Generating Block {block_idx+1}/{num_blocks} | Step {step_idx+1}/{steps} ---\n\n"
-                          f"{ctx_str}{block_str}")
-            
-            yield frame_text
-            
-            # Artificial delay to visualize the step
-            if speed < 20: # If max speed, skip sleep
-                time.sleep(delay)
-            # ------------------------
-
-            full_input = torch.cat([context_ids, mask_block], dim=1)
-            attention_mask = torch.ones_like(full_input, dtype=torch.float32)
-            
-            logits, _ = v2_model(full_input, attention_mask=attention_mask)
-            block_logits = logits[:, -block_size:, :]
-            
-            # Repetition penalty
-            if repetition_penalty != 1.0:
-                seen_tokens = set(context_ids[0].tolist())
-                for i in range(block_size):
-                    if not is_masked[0, i]:
-                        seen_tokens.add(mask_block[0, i].item())
-                for token_id in seen_tokens:
-                    if token_id < block_logits.shape[-1]:
-                        if block_logits[0, :, token_id].mean() > 0:
-                            block_logits[:, :, token_id] /= repetition_penalty
-                        else:
-                            block_logits[:, :, token_id] *= repetition_penalty
-            
-            block_logits = block_logits / temperature
-            
-            # Top-K
-            if top_k > 0:
-                top_k_logits, top_k_indices = torch.topk(block_logits, top_k, dim=-1)
-                block_logits = torch.full_like(block_logits, float('-inf'))
-                block_logits.scatter_(-1, top_k_indices, top_k_logits)
-            
-            # Top-P
-            if top_p < 1.0:
-                sorted_logits, sorted_indices = torch.sort(block_logits, descending=True, dim=-1)
-                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
-                sorted_indices_to_remove = cumulative_probs > top_p
-                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-                sorted_indices_to_remove[..., 0] = 0
-                indices_to_remove = sorted_indices_to_remove.scatter(-1, sorted_indices, sorted_indices_to_remove)
-                block_logits[indices_to_remove] = float('-inf')
-            
-            probs = F.softmax(block_logits, dim=-1)
-            probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0)
-            probs = probs.clamp(min=1e-10)
-            probs = probs / probs.sum(dim=-1, keepdim=True)
-            
-            sampled_tokens = torch.multinomial(probs.view(-1, probs.size(-1)), num_samples=1)
-            sampled_tokens = sampled_tokens.view(1, block_size)
-            
-            confidence = probs.gather(-1, sampled_tokens.unsqueeze(-1)).squeeze(-1)
-            
-            tokens_to_unmask = max(1, block_size // steps)
-            if step_idx == steps - 1:
-                tokens_to_unmask = is_masked.sum().item()
-            
-            if tokens_to_unmask > 0 and is_masked.sum() > 0:
-                masked_confidence = confidence.clone()
-                masked_confidence[~is_masked] = -1.0
-                num_to_unmask = min(tokens_to_unmask, is_masked.sum().item())
-                _, top_indices = torch.topk(masked_confidence.view(-1), num_to_unmask)
-                
-                for idx in top_indices:
-                    mask_block[0, idx] = sampled_tokens[0, idx]
-                    is_masked[0, idx] = False
-        
-        context_ids = torch.cat([context_ids, mask_block], dim=1)
-    
-    generated_ids = context_ids[0].tolist()
-    final_text = v2_tokenizer.decode(generated_ids, skip_special_tokens=True)
-    yield final_text
-
-
-# ==============================================================================
-# ------------------------------- GRADIO UI ------------------------------------
-# ==============================================================================
-
-css = '''.gradio-container > .fillable {max-width: 900px !important}
-h3{margin-top: 1em}
-p{margin-top: 0}
-textarea{font-family: monospace; background-color: #1a1b1e; color: #e0e0e0}
-'''
-
-with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
-    gr.Markdown("# Diffusion Language Models Playground")
-    
-    with gr.Tabs():
-        
-        # --- TAB 1: VERSION 1 (CHAR DIFFUSION) ---
-        with gr.Tab("Version 1: Character Diffusion (NanoGPT)"):
-            gr.Markdown("### Tiny 11M parameter character-based continuous diffusion.")
-            with gr.Row():
-                with gr.Column(scale=1):
-                    v1_steps = gr.Slider(64, 512, 128, step=1, label="Denoising Steps")
-                    v1_speed = gr.Slider(1, 20, 10, step=1, label="Generation/Replay Speed")
-                    with gr.Row():
-                        v1_btn = gr.Button("Generate", variant="primary")
-                        v1_stop = gr.Button("Stop", variant="stop")
-                with gr.Column(scale=2):
-                    v1_out = gr.Textbox(label="Generated Text", lines=15, interactive=False)
-            
-            # V1 Logic: Merged generation and replay for proper stopping
-            v1_event = v1_btn.click(v1_generate_stream, inputs=[v1_steps, v1_speed], outputs=[v1_out])
-            v1_stop.click(fn=None, inputs=None, outputs=None, cancels=[v1_event])
-
-        # --- TAB 2: VERSION 2 (BLOCK DIFFUSION) ---
-        with gr.Tab("Version 2: Block Diffusion (LLaDA-style)"):
-            gr.Markdown("### Block-based diffusion using Qwen tokenizer.")
-            if v2_model is None:
-                gr.Warning(f"V2 Model not loaded. Please check path: {V2_MODEL_PATH}")
-            
-            with gr.Row():
-                with gr.Column(scale=1):
-                    v2_prompt = gr.Textbox(label="Prompt", value="The king went to the")
-                    v2_steps = gr.Slider(4, 64, 16, step=1, label="Steps per Block")
-                    v2_block_size = gr.Slider(8, 126, 32, step=8, label="Block Size")
-                    v2_max_tokens = gr.Slider(32, 1024, 128, step=32, label="Total New Tokens")
-                    v2_speed = gr.Slider(1, 20, 1, step=1, label="Generation/Replay Speed") 
-                    with gr.Row():
-                        v2_btn = gr.Button("Generate", variant="primary")
-                        v2_stop = gr.Button("Stop", variant="stop")
-                with gr.Column(scale=2):
-                    v2_out = gr.Textbox(label="Generated Text", lines=15, interactive=False)
-
-            # V2 Logic
-            v2_event = v2_btn.click(
-                v2_generate_block_diffusion, 
-                inputs=[v2_prompt, v2_steps, v2_block_size, v2_max_tokens, v2_speed], 
-                outputs=[v2_out]
-            )
-            v2_stop.click(fn=None, inputs=None, outputs=None, cancels=[v2_event])
-
-if __name__ == "__main__":
+import gradio as gr
+import spaces
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import numpy as np
+import math
+import os
+import pickle
+import requests
+import textwrap
+import subprocess
+import shutil
+import time
+from dataclasses import dataclass
+from typing import Optional
+from transformers import AutoTokenizer
+
+# ==============================================================================
+# ------------------------- VERSION 1: SHARED SETUP ----------------------------
+# ==============================================================================
+
+def setup_environment():
+    """Checks for and sets up the necessary data for V1."""
+    nano_gpt_repo_path = 'nanoGPT'
+    data_dir_path = 'shakespeare_char'
+    meta_path = os.path.join(data_dir_path, 'meta.pkl')
+
+    if os.path.exists(meta_path):
+        return
+
+    print("Required data not found. Starting one-time setup...")
+    if not os.path.exists(nano_gpt_repo_path):
+        try:
+            subprocess.run(['git', 'clone', 'https://github.com/karpathy/nanoGPT.git'], check=True, capture_output=True, text=True)
+        except subprocess.CalledProcessError as e:
+            print(f"Error cloning repository: {e.stderr}")
+            pass 
+
+    source_data_dir = os.path.join(nano_gpt_repo_path, 'data', 'shakespeare_char')
+    if not os.path.exists(data_dir_path) and os.path.exists(source_data_dir):
+        shutil.copytree(source_data_dir, data_dir_path)
+    
+    # Check if we can run prepare
+    prepare_script_path = os.path.join(data_dir_path, 'prepare.py')
+    if os.path.exists(prepare_script_path) and not os.path.exists(meta_path):
+        subprocess.run(['python', 'prepare.py'], check=True, cwd=data_dir_path, capture_output=True, text=True)
+
+setup_environment()
+
+def download_file(url, filename):
+    if os.path.exists(filename):
+        return
+    print(f"Downloading '{filename}'...")
+    try:
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+        with open(filename, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+    except requests.exceptions.RequestException as e:
+        print(f"Error downloading {url}: {e}")
+
+# ==============================================================================
+# ---------------------- VERSION 1: ARCHITECTURE & LOGIC -----------------------
+# ==============================================================================
+
+# V1 Constants and Meta Loading
+v1_data_dir = './shakespeare_char/'
+v1_meta_url = 'https://huggingface.co/spaces/thejagstudio/NanoDiffusion/resolve/main/meta.pkl'
+v1_meta_path = 'meta.pkl'
+download_file(v1_meta_url, v1_meta_path)
+
+v1_vocab_size = 65 # Fallback
+v1_itos = {}
+v1_stoi = {}
+
+if os.path.exists(v1_meta_path):
+    with open(v1_meta_path, 'rb') as f:
+        meta = pickle.load(f)
+        v1_vocab_size = meta['vocab_size']
+        v1_itos = meta['itos']
+        v1_stoi = meta['stoi']
+
+v1_context_length = 256
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+def v1_decode(indices_tensor: torch.Tensor):
+    if indices_tensor.dim() > 1:
+        indices_tensor = indices_tensor.squeeze(0)
+    indices = indices_tensor.cpu().numpy()
+    return ''.join([v1_itos.get(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)
+
+@dataclass
+class V1_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
+
+class V1_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):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        return x
+
+class V1_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 = torch.nn.functional.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)
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+def v1_modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return x * (1 + scale) + shift
+
+def v1_bias_add_scale(x: torch.Tensor, bias: Optional[torch.Tensor], scale: torch.Tensor, residual: Optional[torch.Tensor]) -> torch.Tensor:
+    if bias is not None:
+        out = scale * (x + bias)
+    else:
+        out = scale * x
+    if residual is not None:
+        out = residual + out
+    return out
+
+class V1_DDiTBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(config.n_embd, bias=config.bias)
+        self.attn = V1_SelfAttention(config)
+        self.ln_2 = nn.LayerNorm(config.n_embd, bias=config.bias)
+        self.mlp = V1_MLP(config)
+        self.adaLN_modulation = nn.Linear(config.cond_dim, 6 * config.n_embd)
+        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 = v1_modulate(self.ln_1(x), shift_msa, scale_msa)
+        x = self.attn(x)
+        x = v1_bias_add_scale(self.attn(self.ln_1(x)), None, gate_msa, x_skip)
+        x = v1_bias_add_scale(self.mlp(v1_modulate(self.ln_2(x), shift_mlp, scale_mlp)), None, gate_mlp, x)
+        return x
+
+class V1_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 = v1_modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+class V1_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)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+class V1_GPT(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        assert config.vocab_size is not None
+        assert config.block_size is not None
+        self.config = config
+        self.sigma_map = V1_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([V1_DDiTBlock(config) for _ in range(config.n_layer)]),
+            ln_f = nn.LayerNorm(config.n_embd, bias=config.bias),
+        ))
+        self.lm_head = V1_DDitFinalLayer(config)
+        self.apply(self._init_weights)
+        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))
+        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+        pos = torch.arange(0, t, dtype=torch.long, device=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)
+        x = self.lm_head(x, c)
+        x = torch.scatter(x, -1, idx[..., None], torch.zeros_like(x[..., :1]))
+        return x
+
+class V1_GeometricNoise:
+    def __init__(self, sigma_min=1e-4, sigma_max=20):
+        self.sigmas = 1.0 * torch.tensor([sigma_min, sigma_max]).to(device)
+    def rate_noise(self, t):
+        return self.sigmas[0] ** (1 - t) * self.sigmas[1] ** t * (self.sigmas[1].log() - self.sigmas[0].log())
+    def total_noise(self, t):
+        return self.sigmas[0] ** (1 - t) * self.sigmas[1] ** t
+    def __call__(self, t):
+        return self.total_noise(t), self.rate_noise(t)
+
+# --- V1 Inference Logic ---
+def v1_transition(x_t: torch.Tensor, delta_sigma: torch.Tensor) -> torch.Tensor:
+    base_prob = (1 - torch.exp(-delta_sigma[..., None])) / v1_vocab_size
+    trans = torch.ones(*x_t.shape, v1_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)
+    trans = trans.scatter(-1, x_t[..., None], diag_fill)
+    return trans
+
+def v1_staggered_score(score, delta_sigma):
+    exp_factor = torch.exp(-delta_sigma)[..., None]
+    correction = ((exp_factor - 1) / (v1_vocab_size * exp_factor)) * score.sum(dim=-1, keepdim=True)
+    return correction + score / exp_factor
+
+def v1_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)
+
+# --- V1 Model Loading ---
+print("Initializing V1 Model...")
+v1_model_args = dict(n_layer=6, n_head=6, n_embd=384, cond_dim=64,
+                 bias=False, vocab_size=v1_vocab_size, block_size=v1_context_length, dropout=0.2)
+v1_config = V1_GPTConfig(**v1_model_args)
+v1_model = V1_GPT(v1_config)
+try:
+    v1_model.load_state_dict(
+        torch.hub.load_state_dict_from_url(
+            'https://huggingface.co/spaces/thejagstudio/diffusion-gpt/resolve/main/final_model.pth?download=true',
+            map_location=device
+        )
+    )
+    v1_model.to(device)
+    v1_model.eval()
+    print("V1 Model loaded successfully.")
+except Exception as e:
+    print(f"Failed to load V1 model: {e}")
+    v1_model = None
+
+v1_noise = V1_GeometricNoise(sigma_min=1e-4, sigma_max=20)
+
+
+def v1_generate_stream(steps, speed):
+    """
+    Generator function for V1 that yields frames directly.
+    Combined logic of generation and replay to allow for immediate stopping.
+    """
+    if v1_model is None:
+        yield "Error: V1 Model not loaded"
+        return
+        
+    steps = int(steps)
+    speed = float(speed)
+    eps = 1e-5
+    
+    # Calculate delay based on speed slider (similar to V2)
+    # 0.5 is base constant, speed scales it down
+    delay = 0.5 / max(speed, 0.1)
+
+    x = torch.randint(0, v1_vocab_size, (1, v1_context_length), device=device)
+    initial_text = f"--- Initial Random Noise ---\n\n{wrap_text(v1_decode(x[0]))}"
+    yield initial_text
+    time.sleep(delay)
+
+    timesteps = torch.linspace(1, eps, steps + 1, device=device)
+    step_size = (1 - eps) / steps
+    
+    with torch.no_grad():
+        for i in range(steps):
+            t = timesteps[i] * torch.ones(x.shape[0], 1, device=device)
+            curr_sigma_bar = v1_noise(t)[0]
+            
+            next_sigma_bar = v1_noise(t - step_size)[0]
+            delta_sigma = curr_sigma_bar - next_sigma_bar
+
+            log_score = v1_model(x, curr_sigma_bar)
+            score = torch.exp(log_score)
+
+            stag_score = v1_staggered_score(score, delta_sigma)
+            probs = stag_score * v1_transition(x, delta_sigma)
+            x = v1_sample_categorical(probs)
+            
+            progress_text = f"--- Denoising Step {i + 1}/{steps} ---\n\n{wrap_text(v1_decode(x[0]))}"
+            yield progress_text
+            
+            # Artificial delay for visualization
+            if speed < 20:
+                time.sleep(delay)
+            
+        t = timesteps[steps] * torch.ones(x.shape[0], 1, device=device)
+        curr_sigma_bar = v1_noise(t)[0]
+        delta_sigma = curr_sigma_bar
+
+        log_score = v1_model(x, curr_sigma_bar)
+        score = torch.exp(log_score)
+        stag_score = v1_staggered_score(score, delta_sigma)
+        probs = stag_score * v1_transition(x, delta_sigma)
+        x = v1_sample_categorical(probs)
+
+    final_text = f"--- Final Denoised Text (Step {steps}) ---\n\n{wrap_text(v1_decode(x[0]))}"
+    yield final_text
+
+# ==============================================================================
+# ---------------------- VERSION 2: ARCHITECTURE & LOGIC -----------------------
+# ==============================================================================
+
+# PLEASE UPDATE THIS PATH TO YOUR ACTUAL LOCAL FILE OR URL
+V2_MODEL_PATH = "checkpoints/model_fp32.pt"
+
+class V2_RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+    
+    def forward(self, x):
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.eps)
+        return self.weight * x
+
+class V2_RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=16384, base=100000, scaling_factor=1.0):
+        super().__init__()
+        self.scaling_factor = scaling_factor
+        self.dim = dim
+        self.base = base
+        self.max_position_embeddings = max_position_embeddings
+        self.inv_freq = None
+        self._cache = {}
+    
+    def _update_freqs(self, device):
+        base = self.base * (self.scaling_factor ** (self.dim / (self.dim - 2)))
+        inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+        self.inv_freq = inv_freq
+    
+    def forward(self, x, seq_len=None):
+        if seq_len is None:
+            seq_len = x.shape[-2]
+        
+        if self.inv_freq is None or self.inv_freq.device != x.device:
+            self._update_freqs(x.device)
+        
+        cache_key = (seq_len, x.device, x.dtype)
+        if cache_key in self._cache:
+            return self._cache[cache_key]
+        
+        t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype)
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        
+        cos = emb.cos()[None, None, :, :]
+        sin = emb.sin()[None, None, :, :]
+        
+        self._cache[cache_key] = (cos, sin)
+        if len(self._cache) > 10:
+            self._cache.pop(next(iter(self._cache)))
+        
+        return cos, sin
+
+def v2_apply_rotary_pos_emb(q, k, cos, sin):
+    def rotate_half(x):
+        x1 = x[..., : x.shape[-1] // 2]
+        x2 = x[..., x.shape[-1] // 2 :]
+        return torch.cat((-x2, x1), dim=-1)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+class V2_DiffusionAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.use_flash_attn = config.use_flash_attn
+        
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+    
+    def forward(self, hidden_states, freqs_cis, attention_mask=None, past_kv=None):
+        bsz, q_len, _ = hidden_states.size()
+        
+        q = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        
+        cos, sin = freqs_cis
+        cos = cos[:, :, :q_len, :]
+        sin = sin[:, :, :q_len, :]
+        q, k = v2_apply_rotary_pos_emb(q, k, cos, sin)
+        
+        if past_kv is not None:
+            cache_k, cache_v = past_kv
+            k = torch.cat([cache_k, k], dim=2)
+            v = torch.cat([cache_v, v], dim=2)
+        
+        current_kv = (k, v)
+        k = k.repeat_interleave(self.num_key_value_groups, dim=1)
+        v = v.repeat_interleave(self.num_key_value_groups, dim=1)
+        
+        attn_mask = None
+        if attention_mask is not None:
+            attn_mask = attention_mask[:, None, None, :].to(dtype=q.dtype)
+            attn_mask = (1.0 - attn_mask) * torch.finfo(q.dtype).min
+        
+        output = F.scaled_dot_product_attention(
+            q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False
+        )
+        
+        output = output.transpose(1, 2).contiguous().view(bsz, q_len, self.hidden_size)
+        return self.o_proj(output), current_kv
+
+class V2_MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.act_fn = nn.SiLU()
+    
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+class V2_BlockDiffusionBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.self_attn = V2_DiffusionAttention(config)
+        self.mlp = V2_MLP(config)
+        self.input_layernorm = V2_RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = V2_RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.use_activation_checkpointing = config.use_activation_checkpointing
+    
+    def forward(self, hidden_states, freqs_cis, attention_mask, past_kv):
+        return self._forward(hidden_states, freqs_cis, attention_mask, past_kv)
+    
+    def _forward(self, hidden_states, freqs_cis, attention_mask, past_kv):
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        attn_out, new_kv = self.self_attn(hidden_states, freqs_cis, attention_mask, past_kv)
+        hidden_states = residual + attn_out
+        
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = residual + self.mlp(hidden_states)
+        return hidden_states, new_kv
+
+@dataclass
+class V2_ModelConfig:
+    vocab_size: int = 151936
+    hidden_size: int = 1024
+    intermediate_size: int = 2816
+    num_hidden_layers: int = 16
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 4
+    max_position_embeddings: int = 16384
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 100000.0
+    pad_token_id: int = 0
+    mask_token_id: int = 1
+    use_flash_attn: bool = True
+    use_activation_checkpointing: bool = False
+    attention_dropout: float = 0.0
+    hidden_dropout: float = 0.0
+
+ModelConfig = V2_ModelConfig
+
+class V2_DiffusionLLM(nn.Module):
+    def __init__(self, config: V2_ModelConfig):
+        super().__init__()
+        self.config = config
+        pad_idx = config.pad_token_id if config.pad_token_id < config.vocab_size else None
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=pad_idx)
+        
+        self.layers = nn.ModuleList([V2_BlockDiffusionBlock(config) for _ in range(config.num_hidden_layers)])
+        self.norm = V2_RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.rotary_emb = V2_RotaryEmbedding(
+            config.hidden_size // config.num_attention_heads,
+            config.max_position_embeddings
+        )
+        self.lm_head.weight = self.embed_tokens.weight
+    
+    def forward(self, input_ids, attention_mask=None, past_key_values=None):
+        bsz, seqlen = input_ids.shape
+        hidden_states = self.embed_tokens(input_ids)
+        freqs_cis = self.rotary_emb(hidden_states, seq_len=seqlen)
+        
+        if past_key_values is None:
+            past_key_values = [None] * len(self.layers)
+        
+        new_kvs = []
+        for i, layer in enumerate(self.layers):
+            hidden_states, kv = layer(hidden_states, freqs_cis, attention_mask, past_key_values[i])
+            new_kvs.append(kv)
+        
+        hidden_states = self.norm(hidden_states)
+        logits = self.lm_head(hidden_states)
+        return logits, new_kvs
+
+DiffusionLLM = V2_DiffusionLLM
+
+# --- V2 Loading Logic ---
+print("Initializing V2 components...")
+v2_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
+if v2_tokenizer.pad_token is None:
+    v2_tokenizer.pad_token = v2_tokenizer.eos_token
+
+v2_model = None
+v2_config = None
+
+if os.path.exists(V2_MODEL_PATH):
+    print(f"Loading V2 model from {V2_MODEL_PATH}...")
+    try:
+        checkpoint = torch.load(V2_MODEL_PATH, map_location=device, weights_only=False)
+        v2_config = checkpoint['config']
+        v2_model = V2_DiffusionLLM(v2_config)
+        state_dict = checkpoint['model_state']
+        state_dict = {k: v.float() for k, v in state_dict.items()}
+        v2_model.load_state_dict(state_dict)
+        v2_model = v2_model.to(device)
+        v2_model.eval()
+        print("V2 Model loaded.")
+    except Exception as e:
+        print(f"Error loading V2 model: {e}")
+else:
+    print(f"V2 Model file not found at {V2_MODEL_PATH}. Version 2 tab will not work without it.")
+
+
+@torch.no_grad()
+def v2_generate_block_diffusion(prompt, steps, block_size, max_new_tokens, replay_speed):
+    """
+    Refactored to yield frames for real-time streaming.
+    """
+    if v2_model is None:
+        yield "Error: V2 Model not found. Check path."
+        return
+
+    v2_model.eval()
+    # Handle inputs
+    steps = int(steps)
+    block_size = int(block_size)
+    max_new_tokens = int(max_new_tokens)
+    speed = float(replay_speed)
+
+    prompt_ids = v2_tokenizer.encode(prompt, return_tensors="pt").to(device)
+    config = v2_model.config
+    num_blocks = max_new_tokens // block_size
+
+    context_ids = prompt_ids
+
+    # Helper params
+    temperature = 1.0
+    top_k = 40
+    top_p = 0.9
+    repetition_penalty = 1.2
+    
+    # Calculate delay based on speed slider
+    delay = 0.5 / max(speed, 0.1) 
+
+    for block_idx in range(num_blocks):
+        mask_block = torch.full((1, block_size), config.mask_token_id, device=device)
+        is_masked = torch.ones(1, block_size, dtype=torch.bool, device=device)
+        
+        for step_idx in range(steps):
+            # --- SNAPSHOT & YIELD ---
+            # Decode context
+            ctx_str = v2_tokenizer.decode(context_ids[0], skip_special_tokens=True)
+            
+            # Decode block with masking visual
+            block_tokens = mask_block[0].tolist()
+            block_vis = []
+            for i, tid in enumerate(block_tokens):
+                if is_masked[0, i]:
+                    block_vis.append("░") # Mask symbol
+                else:
+                    block_vis.append(v2_tokenizer.decode([tid], skip_special_tokens=False))
+            
+            block_str = "".join(block_vis)
+            
+            frame_text = (f"--- Generating Block {block_idx+1}/{num_blocks} | Step {step_idx+1}/{steps} ---\n\n"
+                          f"{ctx_str}{block_str}")
+            
+            yield frame_text
+            
+            # Artificial delay to visualize the step
+            if speed < 20: # If max speed, skip sleep
+                time.sleep(delay)
+            # ------------------------
+
+            full_input = torch.cat([context_ids, mask_block], dim=1)
+            attention_mask = torch.ones_like(full_input, dtype=torch.float32)
+            
+            logits, _ = v2_model(full_input, attention_mask=attention_mask)
+            block_logits = logits[:, -block_size:, :]
+            
+            # Repetition penalty
+            if repetition_penalty != 1.0:
+                seen_tokens = set(context_ids[0].tolist())
+                for i in range(block_size):
+                    if not is_masked[0, i]:
+                        seen_tokens.add(mask_block[0, i].item())
+                for token_id in seen_tokens:
+                    if token_id < block_logits.shape[-1]:
+                        if block_logits[0, :, token_id].mean() > 0:
+                            block_logits[:, :, token_id] /= repetition_penalty
+                        else:
+                            block_logits[:, :, token_id] *= repetition_penalty
+            
+            block_logits = block_logits / temperature
+            
+            # Top-K
+            if top_k > 0:
+                top_k_logits, top_k_indices = torch.topk(block_logits, top_k, dim=-1)
+                block_logits = torch.full_like(block_logits, float('-inf'))
+                block_logits.scatter_(-1, top_k_indices, top_k_logits)
+            
+            # Top-P
+            if top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(block_logits, descending=True, dim=-1)
+                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                sorted_indices_to_remove[..., 0] = 0
+                indices_to_remove = sorted_indices_to_remove.scatter(-1, sorted_indices, sorted_indices_to_remove)
+                block_logits[indices_to_remove] = float('-inf')
+            
+            probs = F.softmax(block_logits, dim=-1)
+            probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0)
+            probs = probs.clamp(min=1e-10)
+            probs = probs / probs.sum(dim=-1, keepdim=True)
+            
+            sampled_tokens = torch.multinomial(probs.view(-1, probs.size(-1)), num_samples=1)
+            sampled_tokens = sampled_tokens.view(1, block_size)
+            
+            confidence = probs.gather(-1, sampled_tokens.unsqueeze(-1)).squeeze(-1)
+            
+            tokens_to_unmask = max(1, block_size // steps)
+            if step_idx == steps - 1:
+                tokens_to_unmask = is_masked.sum().item()
+            
+            if tokens_to_unmask > 0 and is_masked.sum() > 0:
+                masked_confidence = confidence.clone()
+                masked_confidence[~is_masked] = -1.0
+                num_to_unmask = min(tokens_to_unmask, is_masked.sum().item())
+                _, top_indices = torch.topk(masked_confidence.view(-1), num_to_unmask)
+                
+                for idx in top_indices:
+                    mask_block[0, idx] = sampled_tokens[0, idx]
+                    is_masked[0, idx] = False
+        
+        context_ids = torch.cat([context_ids, mask_block], dim=1)
+    
+    generated_ids = context_ids[0].tolist()
+    final_text = v2_tokenizer.decode(generated_ids, skip_special_tokens=True)
+    yield final_text
+
+
+# ==============================================================================
+# ------------------------------- GRADIO UI ------------------------------------
+# ==============================================================================
+
+css = '''.gradio-container > .fillable {max-width: 900px !important}
+h3{margin-top: 1em}
+p{margin-top: 0}
+textarea{font-family: monospace; background-color: #1a1b1e; color: #e0e0e0}
+'''
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
+    gr.Markdown("# Diffusion Language Models Playground")
+    
+    with gr.Tabs():
+        
+        # --- TAB 1: VERSION 1 (CHAR DIFFUSION) ---
+        with gr.Tab("Version 1: Character Diffusion (NanoGPT)"):
+            gr.Markdown("### Tiny 11M parameter character-based continuous diffusion.")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    v1_steps = gr.Slider(64, 512, 128, step=1, label="Denoising Steps")
+                    v1_speed = gr.Slider(1, 20, 10, step=1, label="Generation/Replay Speed")
+                    with gr.Row():
+                        v1_btn = gr.Button("Generate", variant="primary")
+                        v1_stop = gr.Button("Stop", variant="stop")
+                with gr.Column(scale=2):
+                    v1_out = gr.Textbox(label="Generated Text", lines=15, interactive=False)
+            
+            # V1 Logic: Merged generation and replay for proper stopping
+            v1_event = v1_btn.click(v1_generate_stream, inputs=[v1_steps, v1_speed], outputs=[v1_out])
+            v1_stop.click(fn=None, inputs=None, outputs=None, cancels=[v1_event])
+
+        # --- TAB 2: VERSION 2 (BLOCK DIFFUSION) ---
+        with gr.Tab("Version 2: Block Diffusion (LLaDA-style)"):
+            gr.Markdown("### Block-based diffusion using Qwen tokenizer.")
+            if v2_model is None:
+                gr.Warning(f"V2 Model not loaded. Please check path: {V2_MODEL_PATH}")
+            
+            with gr.Row():
+                with gr.Column(scale=1):
+                    v2_prompt = gr.Textbox(label="Prompt", value="The king went to the")
+                    v2_steps = gr.Slider(4, 64, 16, step=1, label="Steps per Block")
+                    v2_block_size = gr.Slider(8, 126, 32, step=8, label="Block Size")
+                    v2_max_tokens = gr.Slider(32, 1024, 128, step=32, label="Total New Tokens")
+                    v2_speed = gr.Slider(1, 20, 1, step=1, label="Generation/Replay Speed") 
+                    with gr.Row():
+                        v2_btn = gr.Button("Generate", variant="primary")
+                        v2_stop = gr.Button("Stop", variant="stop")
+                with gr.Column(scale=2):
+                    v2_out = gr.Textbox(label="Generated Text", lines=15, interactive=False)
+
+            # V2 Logic
+            v2_event = v2_btn.click(
+                v2_generate_block_diffusion, 
+                inputs=[v2_prompt, v2_steps, v2_block_size, v2_max_tokens, v2_speed], 
+                outputs=[v2_out]
+            )
+            v2_stop.click(fn=None, inputs=None, outputs=None, cancels=[v2_event])
+
+if __name__ == "__main__":
     demo.launch()