mnist diff vs flow matching

.gitignore

@@ -0,0 +1,182 @@
+models/
+*.jpg
+*.png
+*.pth 
+data/MNIST/
+flagged/
+demo/flagged/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc

app.py

@@ -0,0 +1,172 @@
+import os
+import cv2
+import sys
+import torch
+import numpy as np
+import gradio as gr
+import matplotlib.pyplot as plt
+from src.model import ConditionalUNet
+from huggingface_hub import hf_hub_download
+
+device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+img_shape = (1, 28, 28)
+
+
+def resize(image,size=(200,200)):
+    stretch_near = cv2.resize(image, size, interpolation = cv2.INTER_LINEAR)
+    return stretch_near
+        
+
+model_diff = ConditionalUNet().to(device)
+model_path = hf_hub_download(repo_id="CristianLazoQuispe/MNIST_Diff_Flow_matching", filename="outputs/diffusion/diffusion_model.pth",
+                        cache_dir="models")
+print("Diff Downloaded!")
+model_diff.load_state_dict(torch.load(model_path, map_location=device))
+model_diff.eval()
+
+
+model_flow = ConditionalUNet().to(device)
+model_path = hf_hub_download(repo_id="CristianLazoQuispe/MNIST_Diff_Flow_matching", filename="outputs/flow_matching/flow_model.pth",
+                        cache_dir="models")
+print("Flow Downloaded!")
+model_flow.load_state_dict(torch.load(model_path, map_location=device))
+model_flow.eval()
+
+@torch.no_grad()
+def generate_diffusion_intermediates(label):
+    timesteps = 500
+    img_shape = (1, 28, 28)
+    betas = torch.linspace(1e-4, 0.02, timesteps)
+    alphas = 1.0 - betas
+    alphas_cumprod = torch.cumprod(alphas, dim=0).to(device)
+
+    x = torch.randn(1, *img_shape).to(device)
+    y = torch.tensor([label], dtype=torch.long, device=device)
+    noise_magnitudes = []
+    intermediates = [resize(((x + 1) / 2.0)[0][0].clamp(0, 1).cpu().numpy())]
+
+    for t in reversed(range(timesteps)):
+        t_tensor = torch.full((x.size(0),), t, device=device, dtype=torch.float)
+        noise_pred = model_diff(x, t_tensor, y)
+        x = (1 / alphas[t].sqrt()) * (x - noise_pred * betas[t] / (1 - alphas_cumprod[t]).sqrt() )
+        if t > 0:
+            noise = torch.randn(1, *img_shape).to(device)
+            v = (1 - alphas_cumprod[t - 1]) / (1 - alphas_cumprod[t]) * betas[t]
+            x += v.sqrt() * noise
+            
+        x = x.clamp(-1, 1)
+        if t in [400, 300, 200, 100,0]:
+            #print("t:",t)
+            img_np = ((x + 1) / 2)[0, 0].cpu().numpy()
+            intermediates.append(resize(img_np))
+
+        if t in [499, 399, 299, 199,99,0]:
+            # Compute velocity magnitude and convert to numpy for visualization
+            v_mag = noise_pred[0, 0].abs().clamp(0, 3).cpu().numpy()  # Clamp to max value for better contrast
+            v_mag = (v_mag - v_mag.min()) / (v_mag.max() - v_mag.min() + 1e-5)
+            vel_colored = plt.get_cmap("coolwarm")(v_mag)[:, :, :3]  # (H,W,3)
+            vel_colored = (vel_colored * 255).astype(np.uint8)
+            noise_magnitudes.append(resize(vel_colored, (100, 100)))
+
+    return intermediates+noise_magnitudes
+
+
+def generate_localized_noise(shape, radius=5):
+    """Genera una imagen con ruido solo en un círculo en el centro."""
+    B, C, H, W = shape
+    assert C == 1, "Solo imágenes en escala de grises."
+
+    # Crear máscara circular
+    yy, xx = torch.meshgrid(torch.arange(H), torch.arange(W), indexing='ij')
+    center_y, center_x = H // 2, W // 2
+    mask = ((yy - center_y)**2 + (xx - center_x)**2) >= radius**2
+    mask = mask.float().unsqueeze(0).unsqueeze(0)  # (1, 1, H, W)
+
+    # Aplicar máscara a ruido
+    noise = torch.randn(B, C, H, W)
+    localized_noise = noise * mask + -1*(1-mask)  # solo hay ruido dentro del círculo
+    return localized_noise
+
+
+@torch.no_grad()
+def generate_flow_intermediates(label):
+    x = torch.randn(1, *img_shape).to(device)
+    #x = generate_localized_noise((1, 1, 28, 28), radius=12).to(device)
+    y = torch.full((1,), label, dtype=torch.long, device=device)
+    steps = 500
+    dt = 1.0 / steps
+    
+    images = [(x + 1) / 2.0]  # initial noise
+    vel_magnitudes = []
+    for i in range(steps):
+            
+        t = torch.full((1,), i * dt, device=device)
+        v = model_flow(x, t, y)
+        x = x + v * dt
+
+        if i in [100,200,300,400,499]:
+            images.append((x + 1) / 2.0)
+            # Compute velocity magnitude and convert to numpy for visualization
+        if i in [0,100,200,300,400,499]:
+            v_mag = dt*v[0, 0].abs().clamp(0, 3).cpu().numpy()  # Clamp to max value for better contrast
+            v_mag = (v_mag - v_mag.min()) / (v_mag.max() - v_mag.min() + 1e-5)
+            vel_colored = plt.get_cmap("coolwarm")(v_mag)[:, :, :3]  # (H,W,3)
+            vel_colored = (vel_colored * 255).astype(np.uint8)
+            vel_magnitudes.append(resize(vel_colored, (100, 100)))
+
+    return [resize(images[0][0][0].clamp(0, 1).cpu().numpy())]+[resize(img[0][0].clamp(0, 1).cpu().numpy()) for img in images[-5:]]+vel_magnitudes
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Conditional MNIST Generation: Diffusion vs Flow Matching")
+
+    with gr.Tab("Diffusion"):
+        label_d = gr.Slider(0, 9, step=1, label="Digit Label")
+        btn_d = gr.Button("Generate")
+        with gr.Row():
+            outs_d = [
+                gr.Image(label="Noise"),
+                gr.Image(label="Diffusion t=400"),
+                gr.Image(label="Diffusion t=300"),
+                gr.Image(label="Diffusion t=200"),
+                gr.Image(label="Diffusion t=100"),
+                gr.Image(label="Diffusion t=0"),
+            ]
+        with gr.Row():
+            #400, 300, 200, 100,0
+            flow_noise_imgs = [
+                gr.Image(label="Noise pred t=500"),
+                gr.Image(label="Noise pred t=400"),
+                gr.Image(label="Noise pred t=300"),
+                gr.Image(label="Noise pred t=200"),
+                gr.Image(label="Noise pred t=100"),
+                gr.Image(label="Noise pred t=0")
+            ]
+        btn_d.click(fn=generate_diffusion_intermediates, inputs=label_d, outputs=outs_d+flow_noise_imgs)
+
+    with gr.Tab("Flow Matching"):
+        label_f = gr.Slider(0, 9, step=1, label="Digit Label")
+        btn_f = gr.Button("Generate")
+        with gr.Row():
+            outs_f = [
+                gr.Image(label="Noise"),
+                gr.Image(label="Flow step=100"),
+                gr.Image(label="Flow step=200"),
+                gr.Image(label="Flow step=300"),
+                gr.Image(label="Flow step=400"),
+                gr.Image(label="Flow step=499"),
+            ]
+        with gr.Row():
+            #100,200,300,400,499
+            flow_vel_imgs = [
+                gr.Image(label="Velocity step=0"),
+                gr.Image(label="Velocity step=100"),
+                gr.Image(label="Velocity step=200"),
+                gr.Image(label="Velocity step=300"),
+                gr.Image(label="Velocity step=400"),
+                gr.Image(label="Velocity step=499")
+            ]
+
+        btn_f.click(fn=generate_flow_intermediates, inputs=label_f, outputs=outs_f+flow_vel_imgs)
+
+#demo.launch()
+demo.launch(share=False, server_port=9070)

requirements.txt

@@ -0,0 +1,4 @@
+gradio==4.44.1
+torch==2.0.0
+opencv-python==4.6.0.66
+numpy==1.26.4

src/__init__.py

@@ -0,0 +1 @@
+from . import *

src/dataset.py

@@ -0,0 +1,23 @@
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import datasets, transforms, utils
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+
+
+def get_data(batch_size):
+    # --- Dataset ---
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Lambda(lambda x: x * 2 - 1)
+    ])
+    full_dataset = datasets.MNIST(root='data', train=True, download=True, transform=transform)
+    train_size = int(0.9 * len(full_dataset))
+    val_size = len(full_dataset) - train_size
+    train_dataset, val_dataset = torch.utils.data.random_split(full_dataset, [train_size, val_size])
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
+    return train_loader, val_loader

src/model.py

@@ -0,0 +1,86 @@
+# conditional_mnist_diffusion_flow.py
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import datasets, transforms, utils
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+import os
+import sys
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+# --- Utilidades ---
+def timestep_embedding(timesteps, dim):
+    half = dim // 2
+    freqs = torch.exp(-torch.arange(half, dtype=torch.float32) * torch.log(torch.tensor(10000.0)) / half)
+    args = timesteps[:, None].float() * freqs[None]
+    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
+    return emb
+
+
+# --- Bloque residual con condición ---
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, emb_channels):
+        super().__init__()
+        self.norm1 = nn.GroupNorm(1, in_channels)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
+        self.emb_proj = nn.Linear(emb_channels, out_channels)
+        self.norm2 = nn.GroupNorm(1, out_channels)
+        self.dropout = nn.Dropout(0.1)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
+        self.skip = nn.Conv2d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()
+
+    def forward(self, x, emb):
+        h = F.silu(self.norm1(x))
+        h = self.conv1(h)
+        h += self.emb_proj(F.silu(emb))[:, :, None, None]
+        h = F.silu(self.norm2(h))
+        h = self.dropout(h)
+        h = self.conv2(h)
+        return h + self.skip(x)
+
+
+# --- UNet condicional ---
+class ConditionalUNet(nn.Module):
+    def __init__(self, num_classes=10, base_channels=64):
+        super().__init__()
+        self.time_mlp = nn.Sequential(
+            nn.Linear(1, base_channels),
+            nn.SiLU(),
+            nn.Linear(base_channels, base_channels),
+        )
+        self.label_emb = nn.Embedding(num_classes, base_channels)
+
+        self.enc1 = ResidualBlock(1, base_channels, base_channels)
+        self.enc2 = ResidualBlock(base_channels, base_channels * 2, base_channels)
+        self.down = nn.Conv2d(base_channels * 2, base_channels * 2, 4, stride=2, padding=1)
+
+        self.mid = ResidualBlock(base_channels * 2, base_channels * 2, base_channels)
+
+        self.up = nn.ConvTranspose2d(base_channels * 2, base_channels * 2, 4, stride=2, padding=1)
+        self.dec2 = ResidualBlock(base_channels * 4, base_channels, base_channels)
+        self.dec1 = ResidualBlock(base_channels * 2, base_channels, base_channels)
+
+        self.out_norm = nn.GroupNorm(8, base_channels)
+        self.out_conv = nn.Conv2d(base_channels, 1, 3, padding=1)
+
+    def forward(self, x, t, y):
+        emb_t = self.time_mlp(t.view(-1, 1))
+        emb_y = self.label_emb(y)
+        emb = emb_t + emb_y
+
+        x1 = self.enc1(x, emb)
+        x2 = self.enc2(x1, emb)
+        x3 = self.down(x2)
+        m = self.mid(x3, emb)
+        u = self.up(m)
+
+        d2 = self.dec2(torch.cat([u, x2], dim=1), emb)
+        d1 = self.dec1(torch.cat([d2, x1], dim=1), emb)
+
+        out = self.out_conv(F.silu(self.out_norm(d1)))
+        return out

src/utils.py

@@ -0,0 +1,14 @@
+import os
+import torch
+import random
+import numpy as np
+
+
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True