Create app.py

app.py

@@ -0,0 +1,105 @@
+import gradio as gr
+import json
+import os
+import urllib.request
+from pathlib import Path
+
+import torch
+from PIL import Image
+import numpy as np
+
+# Global variables
+model = None
+checkpoint = None
+device = None
+
+# Download and load the model
+def initialize_model():
+    global model, checkpoint, device
+    
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    model_url = "https://huggingface.co/lazerkat/randomdiffusion/resolve/main/newest.pth"
+    model_path = "newest.pth"
+    
+    # Download if not already present
+    if not os.path.exists(model_path):
+        gr.Info("Downloading model...")
+        urllib.request.urlretrieve(model_url, model_path)
+    
+    # Load checkpoint
+    checkpoint = torch.load(model_path, map_location=device)
+    
+    # Recreate the model architecture
+    from train import DiffusionUNet  # Import directly from training script
+    model = DiffusionUNet(vocab_size=checkpoint['vocab_size']).to(device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+    
+    return "Model loaded successfully!"
+
+# Generate image from prompt
+def generate_image(prompt):
+    global model, checkpoint, device
+    
+    if model is None:
+        return None, "Model not loaded yet. Please wait for initialization."
+    
+    # Tokenize prompt using the saved vocab
+    vocab_data = checkpoint['word_to_idx']
+    max_len = 20
+    words = [w.strip('.,!?"\'') for w in prompt.lower().split()][:max_len]
+    indices = [vocab_data.get(w, 1) for w in words]
+    indices += [0] * (max_len - len(indices))
+    text_tokens = torch.tensor(indices).unsqueeze(0).to(device)
+    
+    # Diffusion sampling
+    from train import Diffusion
+    diffusion = Diffusion(timesteps=500, device=device)
+    
+    with torch.no_grad():
+        generated = diffusion.sample(model, text_tokens, image_size=64, batch_size=1)
+    
+    # Convert to PIL image
+    image = generated.cpu().squeeze(0)
+    image = (image + 1) / 2
+    image = image.clamp(0, 1)
+    image = image.permute(1, 2, 0).numpy()
+    image = (image * 255).astype(np.uint8)
+    img = Image.fromarray(image)
+    
+    return img, f"Generated image for: '{prompt}'"
+
+# Create the interface
+with gr.Blocks(title="RandomDiffusion", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# RandomDiffusion")
+    gr.Markdown("Text-to-Image Diffusion Model")
+    
+    # Model status
+    status = gr.Textbox(label="Model Status", value="Initializing...", interactive=False)
+    
+    # Image generation
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(label="Enter Prompt", placeholder="a beautiful landscape")
+            generate_btn = gr.Button("Generate")
+        with gr.Column():
+            output_image = gr.Image(label="Generated Image", type="pil")
+            result_text = gr.Textbox(label="Result")
+    
+    # Load model on startup
+    demo.load(
+        lambda: initialize_model(),
+        inputs=[],
+        outputs=[status]
+    )
+    
+    # Generate on button click
+    generate_btn.click(
+        generate_image,
+        inputs=[prompt],
+        outputs=[output_image, result_text]
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True)