Update app.py

app.py

@@ -7,11 +7,11 @@ import math
 import os
 from threading import Event
 import traceback
-import cv2  # Added for bilateral filtering
+import cv2
 
 # Constants
 IMG_SIZE = 128
-TIMESTEPS = 300  # From second code
+TIMESTEPS = 300
 NUM_CLASSES = 2
 
 # Global Cancellation Flag
@@ -27,13 +27,13 @@ class SinusoidalPositionEmbeddings(nn.Module):
         self.dim = dim
         half_dim = dim // 2
         emb = math.log(10000) / (half_dim - 1)
-        emb = torch.exp(torch.arange(half_dim) * -emb)  # From second code (no dtype specified)
+        emb = torch.exp(torch.arange(half_dim) * -emb)
         self.register_buffer('embeddings', emb)
 
     def forward(self, time):
-        device = time.device  # From second code
+        device = time.device
         embeddings = self.embeddings.to(device)
-        embeddings = time[:, None] * embeddings[None, :]  # From second code
+        embeddings = time[:, None] * embeddings[None, :]
         return torch.cat([embeddings.sin(), embeddings.cos()], dim=-1)
 
 class UNet(nn.Module):
@@ -130,7 +130,7 @@ class DiffusionModel(nn.Module):
         self.timesteps = timesteps
         self.time_dim = time_dim
 
-        # Linear beta schedule with scaling from second code
+        # Linear beta schedule with scaling
         scale = 1000 / timesteps
         beta_start = scale * 0.0001
         beta_end = scale * 0.02
@@ -165,7 +165,7 @@ class DiffusionModel(nn.Module):
         else:
             labels = labels.to(device)
 
-        # REVERTED SAMPLING LOOP WITH NOISE REDUCTION
+        # ---- REVERTED SAMPLING LOOP WITH NOISE REDUCTION ----
         for t in reversed(range(self.timesteps)):
             if cancel_event.is_set():
                 return None
@@ -199,7 +199,7 @@ class DiffusionModel(nn.Module):
         x_0 = std * x_0 + mean
         x_0 = torch.clamp(x_0, 0., 1.)
 
-        # ENHANCED SHARPENING
+        # ---- ENHANCED SHARPENING ----
         # First apply mild bilateral filtering to reduce noise while preserving edges
         x_np = x_0.cpu().permute(0, 2, 3, 1).numpy()
         filtered = []
@@ -208,7 +208,7 @@ class DiffusionModel(nn.Module):
             filtered_img = cv2.bilateralFilter(img, d=5, sigmaColor=15, sigmaSpace=15)
             filtered.append(filtered_img / 255.0)
         x_0 = torch.tensor(np.array(filtered), device=device).permute(0, 3, 1, 2)
-        
+
         # Then apply stronger unsharp masking
         kernel = torch.ones(3, 1, 5, 5, device=device) / 75
         kernel = kernel.to(x_0.dtype)
@@ -328,7 +328,7 @@ except Exception as e:
     print("Creating dummy model for demonstration")
     loaded_model = DiffusionModel(UNet(num_classes=NUM_CLASSES), timesteps=TIMESTEPS).to(device)
 
-# Gradio UI (from first code)
+# Gradio UI
 with gr.Blocks(theme=gr.themes.Soft(
     primary_hue="violet",
     neutral_hue="slate",
@@ -421,4 +421,4 @@ with gr.Blocks(theme=gr.themes.Soft(
     """
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860,share= True)