Update app.py

app.py

@@ -405,41 +405,41 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
         # Load LoRA weights with respective scales
         lora_names = []
         lora_weights = []
-        
+
         with calculateDuration("Loading LoRA weights"):
             for idx, lora in enumerate(selected_loras):
                 try:
                     lora_name = f"lora_{idx}"
                     lora_path = lora['repo']
-                
+            
                     # Private 모델인 경우 특별 처리
                     if lora.get('private', False):
                         lora_path = load_private_model(lora_path, huggingface_token)
-                
-                    weight_name = lora.get("weights")
-                    print(f"Loading LoRA {lora_name} from {lora_path}")
-                    
+                        print(f"Using private model path: {lora_path}")
+            
                     if image_input is not None:
-                        if weight_name:
-                            pipe_i2i.load_lora_weights(lora_path, weight_name=weight_name, 
-                                                     adapter_name=lora_name, token=huggingface_token)
-                        else:
-                            pipe_i2i.load_lora_weights(lora_path, adapter_name=lora_name, 
-                                                     token=huggingface_token)
+                        pipe_i2i.load_lora_weights(
+                            lora_path,
+                            adapter_name=lora_name,
+                            token=huggingface_token
+                        )
                     else:
-                        if weight_name:
-                            pipe.load_lora_weights(lora_path, weight_name=weight_name, 
-                                                adapter_name=lora_name, token=huggingface_token)
-                        else:
-                            pipe.load_lora_weights(lora_path, adapter_name=lora_name, 
-                                                token=huggingface_token)
-                    
+                        pipe.load_lora_weights(
+                            lora_path,
+                            adapter_name=lora_name,
+                            token=huggingface_token
+                        )
+            
                     lora_names.append(lora_name)
                     lora_weights.append(lora_scale_1 if idx == 0 else lora_scale_2 if idx == 1 else lora_scale_3)
+                    print(f"Successfully loaded LoRA {lora_name} from {lora_path}")
+            
                 except Exception as e:
                     print(f"Failed to load LoRA {lora_name}: {str(e)}")
                     continue
 
+
+
             print("Loaded LoRAs:", lora_names)
             print("Adapter weights:", lora_weights)
             
@@ -652,7 +652,22 @@ def load_private_model(model_id, huggingface_token):
             local_dir_use_symlinks=False
         )
         
-        return local_dir
+        # safetensors 파일 찾기
+        safetensors_file = None
+        for root, dirs, files in os.walk(local_dir):
+            for file in files:
+                if file.endswith('.safetensors'):
+                    safetensors_file = os.path.join(root, file)
+                    break
+            if safetensors_file:
+                break
+        
+        if not safetensors_file:
+            raise Exception(f"No .safetensors file found in {local_dir}")
+            
+        print(f"Found safetensors file: {safetensors_file}")
+        return safetensors_file  # 전체 경로를 반환
+        
     except Exception as e:
         print(f"Error loading private model {model_id}: {str(e)}")
         raise e