Update app.py

app.py

@@ -14,8 +14,183 @@ NUM_CLASSES = 2
 # Define the device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# SinusoidalPositionEmbeddings and UNet classes remain the same as your original code
-# DiffusionModel class remains the same as your original code
+# Define the SinusoidalPositionEmbeddings class
+class SinusoidalPositionEmbeddings(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+        self.register_buffer('embeddings', self._precompute_embeddings(dim))
+
+    def _precompute_embeddings(self, dim):
+        half_dim = dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim) * -emb)
+        return emb
+
+    def forward(self, time):
+        device = time.device
+        embeddings = self.embeddings.to(device)
+        embeddings = time[:, None] * embeddings[None, :]
+        output = torch.cat([embeddings.sin(), embeddings.cos()], dim=-1)
+        return output
+
+# Define the UNet class
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=3, num_classes=2, time_dim=256):
+        super().__init__()
+
+        self.num_classes = num_classes
+        self.label_embedding = nn.Embedding(num_classes, time_dim)
+
+        self.time_mlp = nn.Sequential(
+            SinusoidalPositionEmbeddings(time_dim),
+            nn.Linear(time_dim, time_dim),
+            nn.ReLU(),
+            nn.Linear(time_dim, time_dim)
+        )
+
+        self.inc = self.double_conv(in_channels, 64)
+        self.down1 = self.down(64 + time_dim * 2, 128)
+        self.down2 = self.down(128 + time_dim * 2, 256)
+        self.down3 = self.down(256 + time_dim * 2, 512)
+
+        self.bottleneck = self.double_conv(512 + time_dim * 2, 1024)
+
+        self.up1 = nn.ConvTranspose2d(1024, 256, kernel_size=2, stride=2)
+        self.upconv1 = self.double_conv(256 + 256 + time_dim * 2, 256)
+
+        self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
+        self.upconv2 = self.double_conv(128 + 128 + time_dim * 2, 128)
+
+        self.up3 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
+        self.upconv3 = self.double_conv(64 + 64 + time_dim * 2, 64)
+
+        self.outc = nn.Conv2d(64, out_channels, kernel_size=1)
+
+    def double_conv(self, in_channels, out_channels):
+        return nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True)
+        )
+
+    def down(self, in_channels, out_channels):
+        return nn.Sequential(
+            nn.MaxPool2d(2),
+            self.double_conv(in_channels, out_channels)
+        )
+
+    def forward(self, x, labels, time):
+        label_indices = torch.argmax(labels, dim=1)
+        label_emb = self.label_embedding(label_indices)
+        t_emb = self.time_mlp(time)
+
+        combined_emb = torch.cat([t_emb, label_emb], dim=1)
+        combined_emb = combined_emb.unsqueeze(-1).unsqueeze(-1)
+
+        x1 = self.inc(x)
+        x1_cat = torch.cat([x1, combined_emb.repeat(1, 1, x1.shape[-2], x1.shape[-1])], dim=1)
+
+        x2 = self.down1(x1_cat)
+        x2_cat = torch.cat([x2, combined_emb.repeat(1, 1, x2.shape[-2], x2.shape[-1])], dim=1)
+
+        x3 = self.down2(x2_cat)
+        x3_cat = torch.cat([x3, combined_emb.repeat(1, 1, x3.shape[-2], x3.shape[-1])], dim=1)
+
+        x4 = self.down3(x3_cat)
+        x4_cat = torch.cat([x4, combined_emb.repeat(1, 1, x4.shape[-2], x4.shape[-1])], dim=1)
+
+        x5 = self.bottleneck(x4_cat)
+
+        x = self.up1(x5)
+        x = torch.cat([x, x3], dim=1)
+        x = torch.cat([x, combined_emb.repeat(1, 1, x.shape[-2], x.shape[-1])], dim=1)
+        x = self.upconv1(x)
+
+        x = self.up2(x)
+        x = torch.cat([x, x2], dim=1)
+        x = torch.cat([x, combined_emb.repeat(1, 1, x.shape[-2], x.shape[-1])], dim=1)
+        x = self.upconv2(x)
+
+        x = self.up3(x)
+        x = torch.cat([x, x1], dim=1)
+        x = torch.cat([x, combined_emb.repeat(1, 1, x.shape[-2], x.shape[-1])], dim=1)
+        x = self.upconv3(x)
+
+        output = self.outc(x)
+        return output
+
+# Define the DiffusionModel class
+class DiffusionModel(nn.Module):
+    def __init__(self, model, timesteps=500, time_dim=256):
+        super().__init__()
+        self.model = model
+        self.timesteps = timesteps
+        self.time_dim = time_dim
+
+        self.betas = self.linear_schedule(timesteps)
+        self.alphas = 1. - self.betas
+        self.register_buffer('alpha_bars', torch.cumprod(self.alphas, dim=0).float())
+
+    def linear_schedule(self, timesteps):
+        scale = 1000 / timesteps
+        beta_start = scale * 0.0001
+        beta_end = scale * 0.02
+        return torch.linspace(beta_start, beta_end, timesteps, dtype=torch.float64)
+
+    def forward_diffusion(self, x_0, t, noise):
+        x_0 = x_0.float()
+        noise = noise.float()
+        alpha_bar_t = self.alpha_bars[t].view(-1, 1, 1, 1)
+        x_t = torch.sqrt(alpha_bar_t) * x_0 + torch.sqrt(1. - alpha_bar_t) * noise
+        return x_t
+
+    def forward(self, x_0, labels):
+        t = torch.randint(0, self.timesteps, (x_0.shape[0],), device=x_0.device).long()
+        noise = torch.randn_like(x_0)
+        x_t = self.forward_diffusion(x_0, t, noise)
+        predicted_noise = self.model(x_t, labels, t.float())
+        return predicted_noise, noise, t
+
+@torch.no_grad()
+def sample(model, num_images, timesteps, img_size, num_classes, labels, device):
+    x_t = torch.randn(num_images, 3, img_size, img_size).to(device)
+
+    if labels.ndim == 1:
+        labels_one_hot = torch.zeros(num_images, num_classes).to(device)
+        labels_one_hot[torch.arange(num_images), labels] = 1
+        labels = labels_one_hot
+    else:
+        labels = labels.to(device)
+
+    for t in reversed(range(timesteps)):
+        t_tensor = torch.full((num_images,), t, device=device, dtype=torch.float)
+
+        predicted_noise = model.model(x_t, labels, t_tensor)
+
+        beta_t = model.betas[t].to(device)
+        alpha_t = model.alphas[t].to(device)
+        alpha_bar_t = model.alpha_bars[t].to(device)
+
+        mean = (1 / torch.sqrt(alpha_t)) * (x_t - (beta_t / torch.sqrt(1 - alpha_bar_t)) * predicted_noise)
+        variance = beta_t
+
+        if t > 0:
+            noise = torch.randn_like(x_t)
+        else:
+            noise = torch.zeros_like(x_t)
+
+        x_t = mean + torch.sqrt(variance) * noise
+
+    x_0 = torch.clamp(x_t, -1., 1.)
+
+    mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device)
+    std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device)
+    x_0 = std * x_0 + mean
+    x_0 = torch.clamp(x_0, 0., 1.)
+
+    return x_0
 
 # Load the trained model with improved error handling
 def load_model(model_path, device):