fix: add fallback model loading with multiple Stable Diffusion model paths

configs/hair_transfer.yaml

@@ -1,4 +1,4 @@
-pretrained_model_path: "stable-diffusion-v1-5/stable-diffusion-v1-5"  # Updated model path
+pretrained_model_path: "runwayml/stable-diffusion-v1-5"  # Original model path
 
 pretrained_folder: "./models/stage2"
 encoder_path: "pytorch_model.bin"

infer_full.py

@@ -35,7 +35,28 @@ class StableHair:
 
         try:
             ### Load controlnet
-            unet = UNet2DConditionModel.from_pretrained(self.config.pretrained_model_path, subfolder="unet").to(device)
+            # Try multiple model paths in case of access issues
+            model_paths = [
+                "runwayml/stable-diffusion-v1-5",
+                "stabilityai/stable-diffusion-2-1", 
+                "stabilityai/stable-diffusion-2-1-base"
+            ]
+            
+            unet = None
+            for model_path in model_paths:
+                try:
+                    print(f"Trying to load model from: {model_path}")
+                    unet = UNet2DConditionModel.from_pretrained(model_path, subfolder="unet").to(device)
+                    self.config.pretrained_model_path = model_path  # Update config with working path
+                    print(f"Successfully loaded model from: {model_path}")
+                    break
+                except Exception as e:
+                    print(f"Failed to load {model_path}: {str(e)}")
+                    continue
+            
+            if unet is None:
+                raise Exception("Could not load any Stable Diffusion model")
+            
             controlnet = ControlNetModel.from_unet(unet).to(device)
             
             # Try to load custom controlnet weights, fallback to default if not available