Create app.py

app.py

@@ -0,0 +1,277 @@
+#!/usr/bin/env python3
+"""
+RND1 Diffusion Model Demo for Hugging Face Spaces with ZeroGPU
+"""
+
+import torch
+import gradio as gr
+import spaces
+import random
+import numpy as np
+from transformers import AutoTokenizer
+
+# Global model and tokenizer
+model = None
+tokenizer = None
+device = "cuda"
+
+
+def set_seed(seed: int):
+    """Set random seed for reproducibility."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def load_model():
+    """Load model and tokenizer (called once at startup)."""
+    global model, tokenizer
+    
+    from rnd.configuration_rnd import RND1Config
+    from rnd.modeling_rnd import RND1LM
+    
+    model_path = "radicalnumerics/RND1-Base-0910"
+    
+    print("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    
+    print("Loading model...")
+    cfg = RND1Config.from_pretrained(model_path)
+    cfg.model_type = "rnd1"
+    cfg.attn_implementation = "sdpa"
+    cfg.moe_backend = "hf"
+    
+    model = RND1LM.from_pretrained(
+        model_path,
+        config=cfg,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        trust_remote_code=True,
+        use_safetensors=True,
+        low_cpu_mem_usage=True,
+    )
+    model.eval()
+    print("Model loaded successfully!")
+
+
+@spaces.GPU(duration=120)  # Request GPU for up to 120 seconds
+def generate_text(
+    prompt: str,
+    mode: str,
+    num_steps: int,
+    max_new_tokens: int,
+    temperature: float,
+    top_k: int,
+    top_p: float,
+    seed: int,
+    progress=gr.Progress()
+):
+    """
+    Generate text using RND1 diffusion model.
+    
+    Args:
+        prompt: Input text prompt
+        mode: Generation mode ('task' or 'completion')
+        num_steps: Number of diffusion steps
+        max_new_tokens: Maximum tokens to generate
+        temperature: Sampling temperature
+        top_k: Top-k filtering (0 to disable)
+        top_p: Top-p nucleus filtering (0 to disable)
+        seed: Random seed
+        progress: Gradio progress tracker
+    """
+    if not prompt.strip():
+        return "⚠️ Please enter a prompt."
+    
+    progress(0, desc="Setting seed...")
+    set_seed(seed)
+    
+    progress(0.1, desc="Preparing prompt...")
+    
+    # Format prompt based on mode
+    if mode == "task":
+        if not prompt.strip().startswith("Question:"):
+            formatted_prompt = f"Question: {prompt}\n"
+        else:
+            formatted_prompt = prompt
+    else:
+        formatted_prompt = prompt
+    
+    # Tokenize
+    progress(0.2, desc="Tokenizing...")
+    inputs = tokenizer(formatted_prompt, return_tensors="pt")
+    input_ids = inputs.input_ids.to(device)
+    attention_mask = inputs.attention_mask.to(device) if 'attention_mask' in inputs else None
+    
+    # Prepare generation config
+    from rnd.generation_config import RND1GenerationConfig
+    
+    greedy = (temperature == 1.0)
+    gen_config = RND1GenerationConfig(
+        max_new_tokens=max_new_tokens,
+        num_diffusion_steps=num_steps,
+        mask_token_id=151669,
+        temperature=temperature if not greedy else 1.0,
+        top_k=top_k if top_k > 0 else None,
+        top_p=top_p if top_p > 0 else None,
+        greedy=greedy,
+        eos_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id else 151645,
+        pad_token_id=tokenizer.pad_token_id,
+        bos_token_id=tokenizer.bos_token_id,
+    )
+    
+    # Generate
+    progress(0.3, desc=f"Generating ({num_steps} diffusion steps)...")
+    
+    generator = torch.Generator(device=device)
+    generator.manual_seed(seed)
+    
+    with torch.no_grad():
+        output = model.generate(
+            inputs=input_ids,
+            generation_config=gen_config,
+            generator=generator,
+        )
+    
+    progress(0.9, desc="Decoding...")
+    
+    # Decode generated tokens
+    generated_tokens = output[0][len(input_ids[0]):]
+    generation = tokenizer.decode(
+        generated_tokens.tolist(),
+        skip_special_tokens=True
+    )
+    
+    progress(1.0, desc="Complete!")
+    
+    return generation
+
+
+# Create Gradio interface
+def create_interface():
+    with gr.Blocks(title="RND1 Diffusion Language Model", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("""
+        # 🌊 RND1 Diffusion Language Model
+        
+        Generate text using a diffusion-based language model. The model uses iterative denoising 
+        to progressively refine masked tokens into coherent text.
+        
+        **Note:** First generation may take longer as the model loads.
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Enter your prompt here...",
+                    lines=4,
+                    value="Write a Python function that finds the longest common subsequence of two strings."
+                )
+                
+                mode = gr.Radio(
+                    choices=["task", "completion"],
+                    value="task",
+                    label="Generation Mode",
+                    info="Task: Q&A format for instructions | Completion: Continue the text"
+                )
+                
+                with gr.Accordion("Generation Settings", open=True):
+                    num_steps = gr.Slider(
+                        minimum=16,
+                        maximum=512,
+                        value=256,
+                        step=16,
+                        label="Diffusion Steps",
+                        info="More steps = better quality but slower"
+                    )
+                    
+                    max_new_tokens = gr.Slider(
+                        minimum=32,
+                        maximum=512,
+                        value=256,
+                        step=32,
+                        label="Max New Tokens"
+                    )
+                
+                with gr.Accordion("Sampling Parameters", open=False):
+                    temperature = gr.Slider(
+                        minimum=0.1,
+                        maximum=2.0,
+                        value=1.0,
+                        step=0.1,
+                        label="Temperature",
+                        info="1.0 = greedy/deterministic"
+                    )
+                    
+                    top_k = gr.Slider(
+                        minimum=0,
+                        maximum=100,
+                        value=0,
+                        step=1,
+                        label="Top-K",
+                        info="0 to disable"
+                    )
+                    
+                    top_p = gr.Slider(
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=0.0,
+                        step=0.05,
+                        label="Top-P (Nucleus)",
+                        info="0 to disable"
+                    )
+                    
+                    seed = gr.Slider(
+                        minimum=0,
+                        maximum=100000,
+                        value=12345,
+                        step=1,
+                        label="Random Seed"
+                    )
+                
+                generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")
+            
+            with gr.Column(scale=1):
+                output = gr.Textbox(
+                    label="Generated Text",
+                    lines=20,
+                    show_copy_button=True
+                )
+        
+        gr.Markdown("""
+        ### Examples
+        Try these prompts to see what the model can do!
+        """)
+        
+        gr.Examples(
+            examples=[
+                ["Write a Python function that finds the longest common subsequence of two strings.", "task", 256, 256, 1.0, 0, 0.0, 12345],
+                ["Explain the concept of recursion with a simple example.", "task", 256, 200, 1.0, 0, 0.0, 42],
+                ["The key to understanding quantum computing lies in", "completion", 256, 256, 1.0, 0, 0.0, 9876],
+                ["Once upon a time in a distant galaxy,", "completion", 256, 300, 1.0, 0, 0.0, 7777],
+            ],
+            inputs=[prompt, mode, num_steps, max_new_tokens, temperature, top_k, top_p, seed],
+            outputs=output,
+            fn=generate_text,
+            cache_examples=False,
+        )
+        
+        generate_btn.click(
+            fn=generate_text,
+            inputs=[prompt, mode, num_steps, max_new_tokens, temperature, top_k, top_p, seed],
+            outputs=output,
+        )
+    
+    return demo
+
+
+if __name__ == "__main__":
+    # Load model at startup
+    load_model()
+    
+    # Launch Gradio interface
+    demo = create_interface()
+    demo.queue(max_size=10)  # Enable queue for ZeroGPU
+    demo.launch()