Update app.py

app.py

@@ -364,8 +364,12 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
     ).images[0]
     return final_image
 
-def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, randomize_seed, seed, width, height, loras_state, progress=gr.Progress(track_tqdm=True)):
+def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_indices, 
+             lora_scale_1, lora_scale_2, lora_scale_3, randomize_seed, seed, 
+             width, height, loras_state, progress=gr.Progress(track_tqdm=True)):
     try:
+        
+
         # 한글 감지 및 번역 (이 부분은 그대로 유지)
         if any('\u3131' <= char <= '\u318E' or '\uAC00' <= char <= '\uD7A3' for char in prompt):
             translated = translator(prompt, max_length=512)[0]['translation_text']
@@ -402,26 +406,35 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
         lora_names = []
         lora_weights = []
         with calculateDuration("Loading LoRA weights"):
+
+        # LoRA 로딩 부분 수정
             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}")
+                    
                     if image_input is not None:
                         if weight_name:
-                            pipe_i2i.load_lora_weights(lora_path, weight_name=weight_name, adapter_name=lora_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)
+                            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)
+                            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)
-                    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)
-                except Exception as e:
-                    print(f"Failed to load LoRA {lora_name}: {str(e)}")
+                            pipe.load_lora_weights(lora_path, adapter_name=lora_name, 
+                                                token=huggingface_token)
+                        
 
             print("Loaded LoRAs:", lora_names)
             print("Adapter weights:", lora_weights)
@@ -525,44 +538,107 @@ def update_history(new_image, history):
 
 def refresh_models(huggingface_token):
     try:
-        # HuggingFace API를 통해 사용자의 모델 검색
-        headers = {"Authorization": f"Bearer {huggingface_token}"}
-        api_url = "https://huggingface.co/api/models"
-        params = {
-            "author": huggingface_token,  # 사용자의 모델만 검색
-            "filter": "base_model:black-forest-labs/FLUX.1-dev"  # FLUX 모델만 필터링
+        # HuggingFace API를 통해 사용자의 모델 검색 (private 포함)
+        headers = {
+            "Authorization": f"Bearer {huggingface_token}",
+            "Accept": "application/json"
         }
         
-        response = requests.get(api_url, headers=headers, params=params)
+        # 사용자 이름 가져오기
+        user_info_url = "https://huggingface.co/api/whoami"
+        user_response = requests.get(user_info_url, headers=headers)
+        if user_response.status_code != 200:
+            raise Exception("Failed to get user information")
+        
+        username = user_response.json().get('name')
+        
+        # 사용자의 모든 모델 검색 (private 포함)
+        api_url = f"https://huggingface.co/api/models?author={username}"
+        response = requests.get(api_url, headers=headers)
         if response.status_code != 200:
             raise Exception("Failed to fetch models from HuggingFace")
             
-        user_models = response.json()
+        all_models = response.json()
+        
+        # FLUX 기반 모델 필터링
+        user_models = [
+            model for model in all_models 
+            if model.get('tags') and 'flux' in [tag.lower() for tag in model.get('tags', [])]
+        ]
         
         # 새로운 모델 정보 생성
         new_models = []
         for model in user_models:
             try:
+                # 모델 카드 정보 가져오기
+                model_id = model['id']
+                model_card_url = f"https://huggingface.co/api/models/{model_id}"
+                model_info_response = requests.get(model_card_url, headers=headers)
+                model_info = model_info_response.json()
+                
+                # 프리뷰 이미지 URL 구성
+                preview_images = [
+                    f"https://huggingface.co/{model_id}/resolve/main/preview.png",
+                    f"https://huggingface.co/{model_id}/resolve/main/sample.png",
+                    f"https://huggingface.co/{model_id}/resolve/main/example.png"
+                ]
+                
+                # 이미지 존재 확인
+                image_url = None
+                for preview_url in preview_images:
+                    img_response = requests.head(preview_url, headers=headers)
+                    if img_response.status_code == 200:
+                        image_url = preview_url
+                        break
+                
+                if not image_url:
+                    image_url = "path/to/default/image.png"  # 기본 이미지
+                
+                # 트리거 워드 추출 시도
+                trigger_word = ""
+                if 'instance_prompt' in model_info:
+                    trigger_word = model_info['instance_prompt']
+                
                 model_info = {
-                    "image": f"https://huggingface.co/{model['id']}/resolve/main/preview.png",
-                    "title": model['id'].split('/')[-1],
-                    "repo": model['id'],
+                    "image": image_url,
+                    "title": f"[Private] {model_id.split('/')[-1]}" if model.get('private') else model_id.split('/')[-1],
+                    "repo": model_id,
                     "weights": "pytorch_lora_weights.safetensors",
-                    "trigger_word": ""  # 필요한 경우 모델 카드에서 추출
+                    "trigger_word": trigger_word,
+                    "private": model.get('private', False)
                 }
                 new_models.append(model_info)
+                
             except Exception as e:
                 print(f"Error processing model {model['id']}: {str(e)}")
                 continue
-                
-        # 기존 LoRA 목록과 병합
-        updated_loras = new_models + loras
+        
+        # 사용자의 모델을 최상단에 배치
+        updated_loras = new_models + [lora for lora in loras if lora['repo'] not in [m['repo'] for m in new_models]]
         
         return updated_loras
     except Exception as e:
         print(f"Error refreshing models: {str(e)}")
         return loras
 
+def load_private_model(model_id, huggingface_token):
+    """Private 모델을 로드하는 함수"""
+    try:
+        headers = {"Authorization": f"Bearer {huggingface_token}"}
+        
+        # 모델 다운로드
+        local_dir = snapshot_download(
+            repo_id=model_id,
+            token=huggingface_token,
+            local_dir=f"models/{model_id.replace('/', '_')}",
+            local_dir_use_symlinks=False
+        )
+        
+        return local_dir
+    except Exception as e:
+        print(f"Error loading private model {model_id}: {str(e)}")
+        raise e
+
 custom_theme = gr.themes.Base(
     primary_hue="blue",
     secondary_hue="purple",