Create app.py

app.py

@@ -0,0 +1,133 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import gradio as gr
+import numpy as np
+from PIL import Image
+import os
+IMAGE_SIZE = 32
+CHANNELS = 3
+BATCH_SIZE = 128
+NUM_FEATURES = 128  
+Z_DIM = 200
+LEARNING_RATE = 0.0005 
+EPOCHS = 30
+BETA = 2000
+LOAD_MODEL = False
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+def reparameterize(mu, log_var):
+    std = torch.exp(0.5 * log_var)
+    epsilon = torch.randn_like(std).to(mu.device) 
+    return mu + std * epsilon
+
+class Encoder(nn.Module):
+    def __init__(self, image_size, channels, num_features, z_dim):
+        super(Encoder, self).__init__()
+        self.output_size = image_size // (2**4)
+
+        self.main = nn.Sequential(
+            nn.Conv2d(channels, num_features, kernel_size=3, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(num_features),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(num_features, num_features, kernel_size=3, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(num_features),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(num_features, num_features, kernel_size=3, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(num_features),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(num_features, num_features, kernel_size=3, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(num_features),
+            nn.LeakyReLU(0.2, inplace=True),
+        )
+        self.flatten_size = num_features * self.output_size * self.output_size
+        self.fc_mu = nn.Linear(self.flatten_size, z_dim)
+        self.fc_log_var = nn.Linear(self.flatten_size, z_dim)
+
+    def forward(self, x):
+        x = self.main(x)
+        x = torch.flatten(x, 1)
+        mu = self.fc_mu(x)
+        log_var = self.fc_log_var(x)
+        z = reparameterize(mu, log_var)
+
+        return mu, log_var, z
+class Decoder(nn.Module):
+    def __init__(self, image_size, channels, num_features, z_dim):
+        super(Decoder, self).__init__()
+
+        self.input_size = image_size // (2**4)
+        self.num_features = num_features
+        self.fc = nn.Linear(z_dim, num_features * self.input_size * self.input_size)
+
+        self.main = nn.Sequential(
+            nn.BatchNorm2d(num_features),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.ConvTranspose2d(num_features, num_features, kernel_size=3, stride=2, padding=1, output_padding=1, bias=False),
+            nn.BatchNorm2d(num_features),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.ConvTranspose2d(num_features, num_features, kernel_size=3, stride=2, padding=1, output_padding=1, bias=False),
+            nn.BatchNorm2d(num_features),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.ConvTranspose2d(num_features, num_features, kernel_size=3, stride=2, padding=1, output_padding=1, bias=False),
+            nn.BatchNorm2d(num_features),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.ConvTranspose2d(num_features, num_features, kernel_size=3, stride=2, padding=1, output_padding=1, bias=False),
+            nn.BatchNorm2d(num_features),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.ConvTranspose2d(num_features, channels, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.Sigmoid()
+        )
+
+    def forward(self, z):
+        x = self.fc(z)
+        x = x.view(-1, self.num_features, self.input_size, self.input_size)
+        x = self.main(x)
+        return x
+class VAE(nn.Module):
+    def __init__(self, encoder, decoder, beta):
+        super(VAE, self).__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.beta = beta
+
+    def forward(self, x):
+        mu, log_var, z = self.encoder(x)
+        reconstruction = self.decoder(z)
+        return reconstruction, mu, log_var
+
+    def vae_loss(self, recon_x, x, mu, log_var):
+        recon_loss = self.beta * F.mse_loss(recon_x, x, reduction='sum') / x.size(0)
+        kl_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp()) / x.size(0)
+        total_loss = recon_loss + kl_loss
+        return total_loss, recon_loss, kl_loss
+encoder = Encoder(IMAGE_SIZE, CHANNELS, NUM_FEATURES, Z_DIM).to(device)
+decoder = Decoder(IMAGE_SIZE, CHANNELS, NUM_FEATURES, Z_DIM).to(device)
+model = VAE(encoder, decoder, BETA).to(device)
+model_weights_path = 'vae_final.pth'
+if os.path.exists(model_weights_path):
+    try:
+        model.load_state_dict(torch.load(model_weights_path, map_location=device))
+        model.eval() # Set to evaluation mode
+        print("Model weights loaded successfully.")
+    except Exception as e:
+        print(f"Error loading model weights: {e}")
+        model = None
+else:
+    print(f"Error: Model weights file not found at: {model_weights_path}")
+    model = None
+def generate_image():
+    if model is None:
+        return "Error: Model not loaded. Please ensure 'vae_final.pth' is available."
+
+    with torch.no_grad():
+        z = torch.randn(1, Z_DIM).to(device)
+        generated_image = model.decoder(z).squeeze(0)
+        generated_image = generated_image.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
+        pil_image = Image.fromarray(generated_image, 'RGB')
+    return pil_image
+if model is not None:
+    # Create Gradio interface
+    iface = gr.Interface(fn=generate_image, inputs=None, outputs="image")
+    iface.launch(debug=True)
+else:
+    print("Cannot launch Gradio interface because the model was not loaded.")