Update app.py

app.py

@@ -23,17 +23,23 @@ except Exception as e:
     raise RuntimeError(f"Failed to load InsightFace model: {e}. Ensure network access for initial download.")
 
 # Define paths for preloaded weights
-model_path = "./unet/"  # Preloaded Kolors weights should be at ./unet/
+model_path = "./"  # Start with root, adjust if needed
 ip_adapter_path = "./"
 
-# Debug: List files in the root directory to confirm preloading
+# Debug: List files to confirm preloading
 print("Files in root directory:", os.listdir("."))
 print("Files in ./unet/ directory:", os.listdir("./unet") if os.path.exists("./unet") else "No ./unet/ directory")
 
 # Check if weights exist
 kolors_weights = model_path + "diffusion_pytorch_model.safetensors"
 if not os.path.exists(kolors_weights):
-    raise FileNotFoundError(f"Kolors model weights not found at {kolors_weights}. Preloading failed. Please check .env configuration.")
+    # Check if it's in ./unet/ as a fallback
+    kolors_weights_unet = "./unet/diffusion_pytorch_model.safetensors"
+    if os.path.exists(kolors_weights_unet):
+        model_path = "./unet/"
+        kolors_weights = kolors_weights_unet
+    else:
+        raise FileNotFoundError(f"Kolors model weights not found at {kolors_weights} or {kolors_weights_unet}. Preloading failed. Please check .env configuration.")
 if not os.path.exists(ip_adapter_path + "ip-adapter.bin"):
     raise FileNotFoundError(f"IP-Adapter weights not found at {ip_adapter_path}")