Update app.py

app.py

@@ -97,7 +97,7 @@ current_pipe = None
 model_lock = threading.Lock()
 
 def load_model(model_name="dreamshaper-8"):
-    """Thread-safe model loading with HIGH-QUALITY settings"""
+    """Thread-safe model loading with HIGH-QUALITY settings and better error handling"""
     global model_cache, current_model_name, current_pipe
     
     with model_lock:
@@ -110,11 +110,15 @@ def load_model(model_name="dreamshaper-8"):
         try:
             model_id = MODEL_CHOICES.get(model_name, "lykon/dreamshaper-8")
             
+            print(f"🔧 Attempting to load: {model_id}")
+            
             pipe = StableDiffusionPipeline.from_pretrained(
                 model_id, 
                 torch_dtype=torch.float32,
                 safety_checker=None,
-                requires_safety_checker=False
+                requires_safety_checker=False,
+                local_files_only=False,  # Allow downloading if not cached
+                cache_dir="./model_cache"  # Specific cache directory
             )
             
             pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
@@ -128,13 +132,28 @@ def load_model(model_name="dreamshaper-8"):
             return pipe
             
         except Exception as e:
-            print(f"❌ Model loading failed: {e}")
-            return StableDiffusionPipeline.from_pretrained(
-                "runwayml/stable-diffusion-v1-5", 
-                torch_dtype=torch.float32,
-                safety_checker=None,
-                requires_safety_checker=False
-            ).to("cpu")
+            print(f"❌ Model loading failed for {model_name}: {e}")
+            print(f"🔄 Falling back to stable-diffusion-v1-5")
+            
+            # Fallback to base model
+            try:
+                pipe = StableDiffusionPipeline.from_pretrained(
+                    "runwayml/stable-diffusion-v1-5", 
+                    torch_dtype=torch.float32,
+                    safety_checker=None,
+                    requires_safety_checker=False
+                ).to("cpu")
+                
+                model_cache[model_name] = pipe
+                current_pipe = pipe
+                current_model_name = "sd-1.5"
+                
+                print(f"✅ Fallback model loaded: stable-diffusion-v1-5")
+                return pipe
+                
+            except Exception as fallback_error:
+                print(f"❌ Critical: Fallback model also failed: {fallback_error}")
+                raise
 
 # Initialize default model
 print("🚀 Initializing Storybook Generator API...")