Update app.py

app.py

@@ -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)
+        emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
         self.register_buffer('embeddings', emb)
 
     def forward(self, time):
         device = time.device
         embeddings = self.embeddings.to(device)
-        embeddings = time[:, None] * embeddings[None, :]
+        embeddings = time.float()[:, None] * embeddings[None, :]
         return torch.cat([embeddings.sin(), embeddings.cos()], dim=-1)
 
 class UNet(nn.Module):
@@ -133,15 +133,15 @@ class DiffusionModel(nn.Module):
         scale = 1000 / timesteps
         beta_start = scale * 0.0001
         beta_end = scale * 0.02
-        self.betas = torch.linspace(beta_start, beta_end, timesteps, dtype=torch.float64)
+        self.betas = torch.linspace(beta_start, beta_end, timesteps, dtype=torch.float32)
         self.alphas = 1. - self.betas
-        self.register_buffer('alpha_bars', torch.cumprod(self.alphas, dim=0).float())
+        self.register_buffer('alpha_bars', torch.cumprod(self.alphas, dim=0))
 
     @torch.no_grad()
     def p_sample(self, x, t, labels):
-        betas_t = self.betas[t].view(-1, 1, 1, 1)
-        sqrt_one_minus_alphas_cumprod_t = torch.sqrt(1. - self.alpha_bars[t]).view(-1, 1, 1, 1)
-        sqrt_recip_alphas_t = torch.sqrt(1.0 / (1. - self.betas[t])).view(-1, 1, 1, 1)
+        betas_t = self.betas[t].view(-1, 1, 1, 1).to(x.dtype).to(x.device)
+        sqrt_one_minus_alphas_cumprod_t = torch.sqrt(1. - self.alpha_bars[t]).view(-1, 1, 1, 1).to(x.dtype).to(x.device)
+        sqrt_recip_alphas_t = torch.sqrt(1.0 / (1. - self.betas[t])).view(-1, 1, 1, 1).to(x.dtype).to(x.device)
         
         # Model prediction
         pred_noise = self.model(x, labels, t.float())
@@ -154,11 +154,11 @@ class DiffusionModel(nn.Module):
         else:
             posterior_variance_t = self.betas[t] * (1. - self.alpha_bars[t-1]) / (1. - self.alpha_bars[t])
             noise = torch.randn_like(x)
-            return model_mean + torch.sqrt(posterior_variance_t) * noise
+            return model_mean + torch.sqrt(posterior_variance_t).to(x.device) * noise
 
     @torch.no_grad()
     def sample(self, num_images, img_size, num_classes, labels, device, progress_callback=None):
-        x = torch.randn((num_images, 3, img_size, img_size), device=device)
+        x = torch.randn((num_images, 3, img_size, img_size), device=device, dtype=torch.float32)
         
         for i in reversed(range(0, self.timesteps)):
             t = torch.full((num_images,), i, device=device, dtype=torch.long)
@@ -170,8 +170,8 @@ class DiffusionModel(nn.Module):
                 return None
         
         x = torch.clamp(x, -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)
+        mean = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32).view(1, 3, 1, 1).to(device)
+        std = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32).view(1, 3, 1, 1).to(device)
         x = std * x + mean
         x = torch.clamp(x, 0., 1.)
         
@@ -205,13 +205,14 @@ def load_model(model_path, device):
             print("Model loaded successfully")
             
         except Exception as e:
+            traceback.print_exc()
             raise ValueError(f"Error loading model: {str(e)}")
         
         diffusion_model = DiffusionModel(unet).to(device)
         try:
             diffusion_model = torch.compile(diffusion_model)
-        except:
-            print("Could not compile model - running uncompiled")
+        except Exception as e:
+            print(f"Could not compile model - running uncompiled: {str(e)}")
     else:
         raise FileNotFoundError(f"Model weights not found at {model_path}")
     
@@ -234,7 +235,7 @@ def generate_image(label_str, num_images, progress=gr.Progress()):
     if label_str not in label_map:
         raise gr.Error("Invalid condition selected")
 
-    labels = torch.zeros(num_images, NUM_CLASSES, device=device)
+    labels = torch.zeros(num_images, NUM_CLASSES, device=device, dtype=torch.float32)
     labels[:, label_map[label_str]] = 1
 
     try:
@@ -271,7 +272,7 @@ def generate_image(label_str, num_images, progress=gr.Progress()):
         torch.cuda.empty_cache()
         raise gr.Error("Out of GPU memory - try generating fewer images")
     except Exception as e:
-        print(f"Full error: {traceback.format_exc()}")
+        traceback.print_exc()
         if str(e) != "Generation was cancelled by user":
             raise gr.Error(f"Generation failed: {str(e)}")
         return None
@@ -280,7 +281,11 @@ def generate_image(label_str, num_images, progress=gr.Progress()):
 
 # --- Load Model ---
 model_path = "model_weights.pth"
-loaded_model = load_model(model_path, device)
+try:
+    loaded_model = load_model(model_path, device)
+except Exception as e:
+    print(f"Failed to load model: {str(e)}")
+    raise
 
 # --- Gradio UI ---
 with gr.Blocks(theme=gr.themes.Soft(
@@ -291,7 +296,7 @@ with gr.Blocks(theme=gr.themes.Soft(
 )) as demo:
     gr.Markdown("""
     <center>
-    <h1>CheXpert X-ray Image Generator</h1>
+    <h1>Synthetic X-ray Generator</h1>
     <p><em>Generate synthetic chest X-rays conditioned on pathology</em></p>
     </center>
     """)
@@ -352,8 +357,11 @@ with gr.Blocks(theme=gr.themes.Soft(
     """
 
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False
-    )
+    try:
+        demo.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=False
+        )
+    except Exception as e:
+        print(f"Failed to launch app: {str(e)}")