Update app.py

app.py

@@ -24,8 +24,8 @@ global_vocoder = None
 global_diffusion = None
 
 # Set the models directory
-MODELS_DIR = "/content/models"
-GENERATIONS_DIR = "/content/generations"
+MODELS_DIR = os.path.join(os.path.dirname(__file__), "models")
+GENERATIONS_DIR = os.path.join(os.path.dirname(__file__), "generations")
 
 def prepare(t5, clip, img, prompt):
     bs, c, h, w = img.shape
@@ -211,7 +211,9 @@ model_choices = [os.path.basename(f) for f in model_files]
 
 # Ensure we have at least one model
 if not model_choices:
-    print("No models found in the models directory. Please add at least one .pt file.")
+    print(f"No models found in the models directory: {MODELS_DIR}")
+    print("Available files in the directory:")
+    print(os.listdir(MODELS_DIR))
     model_choices = ["No models available"]
 
 # Set default model
@@ -253,7 +255,11 @@ with gr.Blocks(theme=theme) as iface:
 
     def on_model_change(model_name):
         if model_name != "No models available":
-            load_model(model_name)
+            try:
+                load_model(model_name)
+                print(f"Successfully loaded model: {model_name}")
+            except Exception as e:
+                print(f"Error loading model {model_name}: {str(e)}")
         else:
             print("No valid model selected.")