Update app.py

app.py

@@ -288,13 +288,19 @@ def generate_images(label_str, num_images, progress=gr.Progress()):
         torch.cuda.empty_cache()
 
 # Load model
-MODEL_DIR = "models"
-MODEL_NAME = "model_weights.pth"
-model_path = os.path.join(MODEL_DIR, MODEL_NAME)
+MODEL_NAME = "model_weights.pth"  # Updated to look in root folder
+model_path = MODEL_NAME
 print("Loading model...")
-loaded_model = load_model(model_path, device)
-print("Model loaded successfully!")
+try:
+    loaded_model = load_model(model_path, device)
+    print("Model loaded successfully!")
+except Exception as e:
+    print(f"Failed to load model: {e}")
+    # Create a dummy model for demo purposes
+    print("Creating dummy model for demonstration")
+    loaded_model = DiffusionModel(UNet(num_classes=NUM_CLASSES)).to(device)
 
+    
 # Gradio UI
 with gr.Blocks(theme=gr.themes.Soft(
     primary_hue="violet",