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/NanoDiffusion/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()