Upload 6 files

config.json

@@ -0,0 +1,23 @@
+{
+  "hidden_size": 768,
+  "intermediate_size": 2112,
+  "num_hidden_layers": 16,
+  "num_attention_heads": 12,
+  "num_key_value_heads": 6,
+  "head_dim": 64,
+  "vocab_size": 32000,
+  "mask_vocab_size": 32001,
+  "max_position_embeddings": 5120,
+  "rope_theta": 10000.0,
+  "rms_norm_eps": 1e-06,
+  "hidden_act": "silu",
+  "timestep_emb_hidden": 768,
+  "mask_token_id": 32000,
+  "pad_token_id": 1,
+  "mask_ratio_min": 0.0,
+  "mask_ratio_max": 1.0,
+  "model_type": "metadiffusion",
+  "architectures": [
+    "MetaDiffusionForCausalLM"
+  ]
+}

generation_config.json

@@ -0,0 +1,13 @@
+{
+  "bos_token_id": 0,
+  "eos_token_id": 2,
+  "pad_token_id": 1,
+  "mask_token_id": 32000,
+  "temperature": 0.6,
+  "repetition_penalty": 1.5,
+  "re_mask": 0.1,
+  "num_steps": 512,
+  "max_new_tokens": 512,
+  "use_cache": false,
+  "transformers_version": "4.40.0"
+}

inference.py

@@ -0,0 +1,391 @@
+#!/usr/bin/env python3
+
+import argparse
+import json
+import os
+import sys
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from safetensors.torch import load_file
+from transformers import AutoTokenizer, AutoConfig
+
+
+# ---------------------------------------------------------------------------
+# Model definition
+# ---------------------------------------------------------------------------
+
+class MetaDiffusionConfig:
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x):
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.eps)
+        return self.weight * x
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=5120, base=10000.0):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(
+            position_ids.shape[0], -1, 1
+        )
+        position_ids_expanded = position_ids[:, None, :].float()
+        freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return emb.cos().to(dtype=x.dtype), emb.sin().to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin):
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size * 4),
+            nn.SiLU(),
+            nn.Linear(hidden_size * 4, hidden_size),
+        )
+
+    def forward(self, t):
+        half_dim = self.mlp[0].in_features // 2
+        emb = math.log(10000.0) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=t.device) * -emb)
+        emb = t[:, None].float() * emb[None, :]
+        emb = torch.cat([emb.sin(), emb.cos()], dim=-1)
+        return self.mlp(emb)
+
+
+class TimestepResidual(nn.Module):
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.proj = nn.Linear(hidden_size, hidden_size)
+
+    def forward(self, x, emb):
+        return x + self.proj(emb)[:, None, :]
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
+        self.rotary_emb = RotaryEmbedding(
+            config.head_dim,
+            max_position_embeddings=config.max_position_embeddings,
+            base=config.rope_theta,
+        )
+
+    def forward(self, x, position_ids):
+        batch, seq, _ = x.shape
+        q = self.q_proj(x).view(batch, seq, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(batch, seq, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(batch, seq, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        cos, sin = self.rotary_emb(x, position_ids)
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        if self.num_kv_groups > 1:
+            k = k.repeat_interleave(self.num_kv_groups, dim=1)
+            v = v.repeat_interleave(self.num_kv_groups, dim=1)
+
+        out = F.scaled_dot_product_attention(q, k, v)
+        out = out.transpose(1, 2).contiguous().view(batch, seq, -1)
+        return self.o_proj(out)
+
+
+class 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)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.self_attn = SelfAttention(config)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp = MLP(config)
+        self.timestep_residual = TimestepResidual(config.hidden_size)
+
+    def forward(self, x, timestep_emb, position_ids):
+        residual = x
+        x = self.input_layernorm(x)
+        x = self.self_attn(x, position_ids)
+        x = residual + x
+        x = self.timestep_residual(x, timestep_emb)
+
+        residual = x
+        x = self.post_attention_layernorm(x)
+        x = self.mlp(x)
+        x = residual + x
+        x = self.timestep_residual(x, timestep_emb)
+
+        return x
+
+
+class MetaDiffusionLM(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.mask_token_id = getattr(config, "mask_token_id", config.vocab_size)
+
+        self.embed_tokens = nn.Embedding(
+            config.mask_vocab_size, config.hidden_size,
+            padding_idx=getattr(config, "pad_token_id", 1)
+        )
+        self.timestep_emb = TimestepEmbedding(getattr(config, "timestep_emb_hidden", config.hidden_size))
+        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.lm_head = nn.Linear(config.hidden_size, config.mask_vocab_size, bias=False)
+
+    def forward(self, input_ids, timesteps):
+        batch, seq = input_ids.shape
+        position_ids = torch.arange(seq, device=input_ids.device).unsqueeze(0).expand(batch, -1)
+        x = self.embed_tokens(input_ids)
+        t_emb = self.timestep_emb(timesteps)
+
+        for layer in self.layers:
+            x = layer(x, t_emb, position_ids)
+
+        x = self.norm(x)
+        logits = self.lm_head(x)
+        return logits
+
+
+# ---------------------------------------------------------------------------
+# Generation
+# ---------------------------------------------------------------------------
+
+def cumulative_unmask_frac(i, N, schedule="cosine"):
+    if schedule == "cosine":
+        return 0.5 * (1 - math.cos(math.pi * i / N))
+    return i / N
+
+
+def generate(model, tokenizer, prompt, seq_len=256, num_steps=64, device="cuda",
+             temperature=0.6, repetition_penalty=1.5, watch=False,
+             watch_every=1, mask_token_id=32000):
+    model.eval()
+
+    # Tokenize prompt
+    prompt_ids = tokenizer.encode(prompt, add_special_tokens=False)
+    prompt_ids = torch.tensor([prompt_ids], device=device)
+
+    # Build input: prompt + [MASK] tokens
+    total_len = seq_len
+    gen_len = max(total_len - prompt_ids.shape[1], 0)
+
+    input_ids = torch.full((1, total_len), mask_token_id, device=device, dtype=torch.long)
+    input_ids[0, :prompt_ids.shape[1]] = prompt_ids
+
+    for i in range(num_steps):
+        frac_now = cumulative_unmask_frac(i, num_steps)
+        frac_next = cumulative_unmask_frac(i + 1, num_steps)
+
+        # How many tokens to unmask this step
+        n_masked = (input_ids == mask_token_id).sum().item()
+        n_total_to_unmask = int((frac_next - frac_now) * (total_len - prompt_ids.shape[1]) + 0.5)
+        if i == num_steps - 1:
+            n_unmask = n_masked
+        else:
+            n_unmask = max(n_total_to_unmask, 1) if n_masked > 0 else 0
+
+        t = 1.0 - frac_now
+        t_batch = torch.full((1,), t, device=device)
+
+        with torch.no_grad():
+            logits = model(input_ids, t_batch)
+
+            # Prevent model from predicting [MASK] token
+            logits[:, :, mask_token_id] = -1e9
+
+            if repetition_penalty != 1.0:
+                for tok in input_ids[0].unique():
+                    tok_idx = tok.item()
+                    logits[0, :, tok_idx] = torch.where(
+                        logits[0, :, tok_idx] < 0,
+                        logits[0, :, tok_idx] * repetition_penalty,
+                        logits[0, :, tok_idx] / repetition_penalty
+                    )
+
+            # Sample at masked positions
+            mask_positions = (input_ids == mask_token_id)
+            mask_logits = logits[mask_positions]
+
+            probs = F.softmax(mask_logits / temperature, dim=-1)
+            sampled = torch.multinomial(probs, 1).squeeze(-1)
+
+            # Select which masks to fill (by confidence)
+            if n_unmask < mask_positions.sum():
+                # Get entropy/confidence for each mask
+                log_probs = F.log_softmax(mask_logits, dim=-1)
+                confidence, _ = log_probs.max(dim=-1)
+                _, top_indices = confidence.topk(n_unmask)
+
+                # Only fill top-confidence positions
+                mask_flat = mask_positions.nonzero(as_tuple=False)
+                fill_positions = mask_flat[top_indices]
+                for idx, tok in zip(fill_positions, sampled[top_indices]):
+                    input_ids[idx[0], idx[1]] = tok
+            else:
+                # Fill all remaining masks
+                input_ids[mask_positions] = sampled
+
+        if watch and i % watch_every == 0:
+            text = tokenizer.decode(input_ids[0], skip_special_tokens=True)
+            n_remaining = (input_ids == mask_token_id).sum().item()
+            print(f"Step {i+1}/{num_steps} | LR={t:.3f} | Masks remaining: {n_remaining}")
+            print(text[:200])
+            print()
+
+    # Decode
+    return tokenizer.decode(input_ids[0], skip_special_tokens=False)
+
+def load_model(model_path, device="cuda"):
+    """Load model from safetensors file, directory, or HuggingFace Hub."""
+    # Check if it's a local path or HF hub id
+    is_file = os.path.isfile(model_path) and model_path.endswith(".safetensors")
+    is_dir = os.path.isdir(model_path)
+    is_local = is_file or is_dir
+
+    if is_local:
+        if is_file:
+            safetensors_path = model_path
+            config_path = os.path.join(os.path.dirname(model_path), "config.json")
+        else:
+            config_path = os.path.join(model_path, "config.json")
+            safetensors_path = os.path.join(model_path, "model.safetensors")
+
+        if not os.path.isfile(safetensors_path):
+            print(f"ERROR: model.safetensors not found in {model_path}")
+            sys.exit(1)
+        if not os.path.isfile(config_path):
+            print(f"ERROR: config.json not found next to {safetensors_path}")
+            sys.exit(1)
+
+        with open(config_path) as f:
+            config_dict = json.load(f)
+    else:
+        # Load from HuggingFace Hub
+        from huggingface_hub import hf_hub_download
+        config_path = hf_hub_download(model_path, "config.json")
+        safetensors_path = hf_hub_download(model_path, "model.safetensors")
+
+        with open(config_path) as f:
+            config_dict = json.load(f)
+
+    # Build config
+    config = MetaDiffusionConfig(**config_dict)
+    model = MetaDiffusionLM(config)
+    model = model.to(device)
+
+    # Load weights (remap HF names to model names)
+    state_dict = load_file(safetensors_path)
+    
+    # Remap from HF naming to model naming
+    new_state_dict = {}
+    for key, value in state_dict.items():
+        if key.startswith("model."):
+            new_key = key[len("model."):]
+        else:
+            new_key = key
+        new_state_dict[new_key] = value
+    
+    result = model.load_state_dict(new_state_dict, strict=False)
+    if result.missing_keys:
+        print(f"  Warning: missing keys: {result.missing_keys[:5]}...")
+    if result.unexpected_keys:
+        print(f"  Warning: unexpected keys: {result.unexpected_keys[:5]}...")
+
+    model = model.to(device)
+    print(f"  Model loaded: {sum(p.numel() for p in model.parameters())/1e6:.1f}M params")
+
+    return model, config
+
+
+def main():
+    parser = argparse.ArgumentParser(description="MetaDiffusion inference")
+    parser.add_argument("--model-path", required=True, help="Path to model directory or HF Hub ID")
+    parser.add_argument("--prompt", default="The cat sat on the", help="Input prompt")
+    parser.add_argument("--seq-len", type=int, default=256, help="Sequence length")
+    parser.add_argument("--num-steps", type=int, default=512, help="Denoising steps")
+    parser.add_argument("--temperature", type=float, default=0.6, help="Sampling temperature")
+    parser.add_argument("--repetition-penalty", type=float, default=1.5, help="Repetition penalty")
+    parser.add_argument("--device", default="cuda", help="Device (cuda/cpu)")
+    parser.add_argument("--watch", action="store_true", help="Show denoising progress")
+    parser.add_argument("--watch-every", type=int, default=4, help="Show progress every N steps")
+    parser.add_argument("--base-model", default="SupraLabs/Supra-1.5-50M-Base-exp",
+                        help="HuggingFace model for tokenizer")
+    args = parser.parse_args()
+
+    if "cpu" in args.device:
+        device = torch.device("cpu")
+    else:
+        device = torch.device(args.device if torch.cuda.is_available() else "cpu")
+
+    model, config = load_model(args.model_path, device)
+
+    # Load tokenizer from base model
+    tokenizer = AutoTokenizer.from_pretrained(args.base_model)
+    mask_token_id = getattr(config, "mask_token_id", config.vocab_size)
+
+    print(f"\nPrompt: '{args.prompt}'")
+    print(f"Steps: {args.num_steps} | Temp: {args.temperature}")
+    print()
+
+    output = generate(
+        model, tokenizer, args.prompt,
+        seq_len=args.seq_len, num_steps=args.num_steps,
+        device=device, temperature=args.temperature,
+        repetition_penalty=args.repetition_penalty,
+        watch=args.watch,
+        watch_every=args.watch_every, mask_token_id=mask_token_id
+    )
+
+    print("Output:")
+    print(output)
+
+
+if __name__ == "__main__":
+    main()

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5444f4b6b87d0db2bc93be91747ecd80fa2d306f5c793a52e6deabb5f6796cf3
+size 339058560

tokenizer_config.json

@@ -0,0 +1,12 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 5120,
+  "pad_token": "<pad>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<unk>"
+}