intital commit

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+text_to_digit_diffusion_mnist.pth filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,143 @@
+# -*- coding: utf-8 -*-
+"""app.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/104LGaVaxm2qBgfuuNORSOHuzKCmgOVKm
+"""
+
+### 1. Imports and setup ###
+
+import gradio as gr
+import torch
+import numpy as np
+
+from model import create_diffusion_model
+from timeit import default_timer as timer
+from typing import Tuple
+
+# ------------------------------
+# 2. Model preparation
+# ------------------------------
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Digit vocabulary
+digit_words = ["zero","one","two","three","four","five","six","seven","eight","nine"]
+word_to_idx = {w:i for i,w in enumerate(digit_words)}
+
+# Create diffusion model
+model = create_diffusion_model(
+    emb_dim=64,
+    timesteps=100,
+    seed=42,
+    device=device
+)
+
+# Load trained weights
+model.load_state_dict(
+    torch.load(
+        "text_to_digit_diffusion_mnist.pth",
+        map_location=device
+    )
+)
+
+model.eval()
+
+# ------------------------------
+# 3. Diffusion scheduler utils
+# ------------------------------
+
+timesteps = 100
+betas = torch.linspace(1e-4, 0.02, timesteps).to(device)
+alphas = 1.0 - betas
+alphas_cumprod = torch.cumprod(alphas, dim=0)
+
+def extract(a, t, x_shape):
+    return a.gather(-1, t).reshape(-1,1,1,1).expand(x_shape)
+
+# ------------------------------
+# 4. Predict / Generate function
+# ------------------------------
+
+@torch.inference_mode()
+def predict(text: str) -> Tuple[np.ndarray, float]:
+    start_time = timer()
+
+    text = text.strip().lower()
+
+    # Normalize input
+    if text.isdigit():
+        label = int(text)
+    else:
+        label = word_to_idx.get(text, None)
+
+    if label is None or not (0 <= label <= 9):
+        raise ValueError("Please enter a digit (0–9) or its word form.")
+
+    labels = torch.tensor([label], device=device)
+    uncond_labels = torch.tensor([0], device=device)
+
+    # Start from noise
+    x = torch.randn(1, 1, 28, 28, device=device)
+
+    guidance_scale = 3.0
+
+    for i in reversed(range(1, timesteps)):
+        t = torch.full((1,), i, device=device)
+
+        pred_cond = model(x, t, labels)
+        pred_uncond = model(x, t, uncond_labels)
+
+        pred_noise = pred_uncond + guidance_scale * (pred_cond - pred_uncond)
+
+        beta_t = extract(betas, t, x.shape)
+        alpha_t = extract(alphas, t, x.shape)
+        alpha_bar_t = extract(alphas_cumprod, t, x.shape)
+
+        pred_x0 = (x - torch.sqrt(1 - alpha_bar_t) * pred_noise) / torch.sqrt(alpha_bar_t)
+        x = torch.sqrt(alpha_t) * pred_x0 + torch.sqrt(beta_t) * pred_noise
+
+    img = (x.clamp(-1,1) + 1) / 2
+    img = img[0,0].cpu().numpy()
+
+    end_time = timer()
+    gen_time = round(end_time - start_time, 4)
+
+    return img, gen_time
+
+# ------------------------------
+# 5. Gradio app
+# ------------------------------
+
+title = "Text-to-Digit Diffusion (MNIST)"
+description = (
+    "A **conditional diffusion model** trained on MNIST. "
+    "Type a digit (e.g. `seven` or `7`) to generate a handwritten number."
+)
+article = "Created by [Programming Ocean Academy](https://www.programming-ocean.com/)"
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Textbox(
+        label="Enter digit (0–9 or word)",
+        placeholder="seven or 7"
+    ),
+    outputs=[
+        gr.Image(
+            label="Generated Digit",
+            type="numpy",
+            width=256,
+            height=256
+        ),
+        gr.Number(label="generation time (s)")
+    ],
+    title=title,
+    description=description,
+    article=article
+)
+
+# Launch demo
+demo.launch(debug=False)
+

model.py

@@ -0,0 +1,104 @@
+# -*- coding: utf-8 -*-
+"""model.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/1_9Ac8If0MzzNMrKBdneReuZPjByJNzHE
+"""
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+# --------------------------------------------------
+# UNet for Text-to-Digit Diffusion (MNIST)
+# --------------------------------------------------
+class UNet(nn.Module):
+    def __init__(self, emb_dim: int = 64, num_classes: int = 10, timesteps: int = 100):
+        super().__init__()
+
+        # Time embedding
+        self.time_mlp = nn.Sequential(
+            nn.Linear(1, emb_dim),
+            nn.ReLU(),
+            nn.Linear(emb_dim, emb_dim)
+        )
+
+        # Label embedding
+        self.label_embed = nn.Embedding(num_classes, emb_dim)
+
+        # Encoder
+        self.enc1 = nn.Conv2d(1, 32, 3, padding=1)
+        self.enc2 = nn.Conv2d(32, 64, 3, stride=2, padding=1)
+        self.enc3 = nn.Conv2d(64, 128, 3, stride=2, padding=1)
+
+        # Bottleneck
+        self.bot = nn.Conv2d(128, 128, 3, padding=1)
+
+        # Conditioning projection
+        self.cond_proj = nn.Linear(emb_dim, 128)
+
+        # Decoder
+        self.dec3 = nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1)
+        self.dec2 = nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1)
+        self.dec1 = nn.Conv2d(32, 1, 3, padding=1)
+
+        self.timesteps = timesteps
+
+    def forward(self, x, t, labels):
+        # Time embedding
+        t = t.unsqueeze(-1).float() / self.timesteps
+        t_emb = self.time_mlp(t)
+
+        # Label embedding
+        l_emb = self.label_embed(labels)
+
+        # Conditioning
+        cond = t_emb + l_emb
+        cond = self.cond_proj(cond).unsqueeze(-1).unsqueeze(-1)
+
+        # Encoder
+        x1 = F.relu(self.enc1(x))
+        x2 = F.relu(self.enc2(x1))
+        x3 = F.relu(self.enc3(x2))
+
+        # Bottleneck + conditioning
+        h = F.relu(self.bot(x3 + cond))
+
+        # Decoder with skip connections
+        h = F.relu(self.dec3(h)) + x2
+        h = F.relu(self.dec2(h)) + x1
+
+        return self.dec1(h)
+
+
+# --------------------------------------------------
+# Factory function (EffNet-style)
+# --------------------------------------------------
+def create_diffusion_model(
+    emb_dim: int = 64,
+    num_classes: int = 10,
+    timesteps: int = 100,
+    seed: int = 42,
+    device: str = "cpu"
+):
+    """
+    Creates a conditional diffusion UNet model.
+
+    Returns:
+        model (nn.Module): diffusion UNet
+    """
+
+    # Reproducibility
+    torch.manual_seed(seed)
+
+    model = UNet(
+        emb_dim=emb_dim,
+        num_classes=num_classes,
+        timesteps=timesteps
+    ).to(device)
+
+    return model
+

requirements.txt

@@ -0,0 +1,3 @@
+torch==2.9.0
+torchvision==0.24.0
+gradio==5.50.0

text_to_digit_diffusion_mnist.pt

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

text_to_digit_diffusion_mnist.pth

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