Update app.py

app.py

@@ -16,70 +16,61 @@ pipe = StableDiffusionPipeline.from_pretrained(
     torch_dtype=dtype
 )
 
-# Debug LoRA files
+# Try to load LoRA weights with multiple file name options
 lora_loaded = False
 lora_path = "./lora"
 
-print("🔍 Debugging LoRA files...")
+print("🔍 Looking for LoRA files...")
 
 if os.path.exists(lora_path):
-    print(f"✅ LoRA folder exists: {lora_path}")
-    
-    # List all files in lora folder
     files = os.listdir(lora_path)
     print(f"📁 Files in lora folder: {files}")
     
-    # Check adapter_config.json
-    config_path = os.path.join(lora_path, "adapter_config.json")
-    if os.path.exists(config_path):
-        print("✅ adapter_config.json exists")
-        with open(config_path, 'r') as f:
-            config_content = f.read()
-            print(f"📄 Config content: {config_content[:200]}...")
-    else:
-        print("❌ adapter_config.json missing")
+    # List of possible LoRA filenames to try
+    possible_names = [
+        "adapter_model.safetensors",
+        "pytorch_lora_weights.safetensors", 
+        "lora_weights.safetensors"
+    ]
+    
+    lora_file_found = None
+    for filename in possible_names:
+        if filename in files:
+            lora_file_found = filename
+            print(f"✅ Found LoRA file: {filename}")
+            break
     
-    # Check adapter_model.safetensors
-    model_path = os.path.join(lora_path, "adapter_model.safetensors")
-    if os.path.exists(model_path):
+    if lora_file_found:
+        model_path = os.path.join(lora_path, lora_file_found)
         file_size = os.path.getsize(model_path)
-        print(f"✅ adapter_model.safetensors exists")
         print(f"📊 File size: {file_size:,} bytes ({file_size/1024/1024:.2f} MB)")
         
-        # Try to read the safetensors file directly
+        # Test if the file is readable
         try:
-            print("🔧 Attempting to read safetensors file directly...")
+            print("🔧 Testing file integrity...")
             with safe_open(model_path, framework="pt") as f:
                 keys = list(f.keys())
-                print(f"✅ File is readable! Found {len(keys)} tensors")
-                print(f"📝 First few keys: {keys[:5]}")
+                print(f"✅ File is valid! Found {len(keys)} tensors")
+                
+                # Now try to load it with diffusers
+                print("🎯 Loading with diffusers...")
+                
+                if lora_file_found == "pytorch_lora_weights.safetensors":
+                    # For pytorch_lora_weights.safetensors, try loading it directly
+                    pipe.load_lora_weights(lora_path, weight_name=lora_file_found)
+                else:
+                    # For standard naming, load normally  
+                    pipe.load_lora_weights(lora_path)
+                
+                lora_loaded = True
+                print("✅ LoRA weights loaded successfully!")
                 
-                # Check tensor shapes
-                for key in keys[:3]:
-                    tensor = f.get_tensor(key)
-                    print(f"   {key}: shape {tensor.shape}, dtype {tensor.dtype}")
-                    
-            print("✅ Safetensors file is valid!")
-            
         except Exception as e:
-            print(f"❌ Cannot read safetensors file: {e}")
-            print("💡 This means the file is corrupted or not a valid safetensors file")
-            
+            print(f"❌ Error with {lora_file_found}: {e}")
+            print("💡 File appears to be corrupted or incompatible")
     else:
-        print("❌ adapter_model.safetensors missing")
+        print("❌ No LoRA files found with expected names")
         
-    # Now try to load with diffusers
-    if os.path.exists(config_path) and os.path.exists(model_path):
-        try:
-            print("🎯 Attempting to load LoRA weights with diffusers...")
-            pipe.load_lora_weights(lora_path)
-            lora_loaded = True
-            print("✅ LoRA weights loaded successfully!")
-        except Exception as e:
-            print(f"❌ Failed to load LoRA weights: {e}")
-            print("🔄 Running with base model only")
-            lora_loaded = False
-    
 else:
     print("❌ LoRA folder not found")
 
@@ -109,11 +100,11 @@ def generate(prompt, quality):
 # Build Gradio UI
 title = "Fine-tuned Stable Diffusion"
 if lora_loaded:
-    title += " with LoRA"
+    title += " with LoRA ✨"
     description = "✅ LoRA weights loaded! Your custom trained model is active."
 else:
     title += " (Base Model)"
-    description = "⚠️ Running with base model only - check logs for LoRA loading issues."
+    description = "⚠️ Running with base model only. Check logs for details."
 
 description += "\nChoose 'Fast' for quicker generation or 'High Quality' for better details."