Update app.py

app.py

@@ -197,27 +197,30 @@ def load_model(model_path, device):
         checkpoint = torch.load(model_path, map_location=device)
         
         if 'model_state_dict' in checkpoint:
-            # Filter out DiffusionModel-specific keys
+            # Handle training checkpoint format
             state_dict = {
-                k[6:]: v for k, v in checkpoint['model_state_dict'].items() 
-                if k.startswith('model.') and not k.startswith('model.alpha_bars')
+                k[6:]: v for k, v in checkpoint['model_state_dict'].items()
+                if k.startswith('model.')
             }
             
-            # Load into UNet only
-            missing, unexpected = unet_model.load_state_dict(state_dict, strict=False)
+            # Load UNet weights
+            unet_model.load_state_dict(state_dict, strict=False)
             
-            print(f"Loaded UNet weights. Missing keys: {missing}. Unexpected keys: {unexpected}")
-            
-            # Reinitialize diffusion model with loaded UNet
+            # Initialize diffusion model with loaded UNet
             diffusion_model = DiffusionModel(unet_model, timesteps=TIMESTEPS).to(device)
-        else:
-            # Handle case where it's not a training checkpoint
-            diffusion_model.load_state_dict({
-                k: v for k, v in checkpoint.items()
-                if not k.startswith('alpha_bars')
-            })
             
-        print(f"Model successfully loaded from {model_path}")
+            print(f"Loaded UNet weights from {model_path}")
+        else:
+            # Handle direct model weights format
+            try:
+                # First try loading full DiffusionModel
+                diffusion_model.load_state_dict(checkpoint)
+                print(f"Loaded full DiffusionModel from {model_path}")
+            except RuntimeError:
+                # If that fails, load just the UNet weights
+                unet_model.load_state_dict(checkpoint, strict=False)
+                diffusion_model = DiffusionModel(unet_model, timesteps=TIMESTEPS).to(device)
+                print(f"Loaded UNet weights only from {model_path}")
     else:
         print(f"Weights file not found at {model_path}")
         print("Using randomly initialized weights")