Update app.py

app.py

@@ -574,13 +574,15 @@ def refresh_models(huggingface_token):
                 model_info_response = requests.get(model_card_url, headers=headers)
                 model_info = model_info_response.json()
                 
-                # samples 폴더의 숫자로 된 jpg 파일 찾기
-                fs = HfFileSystem(token=huggingface_token)
-                samples_path = f"{model_id}/samples"
+                # 이미지 URL에 토큰을 포함시키는 방식으로 변경
+                is_private = model.get('private', False)
+                base_image_name = "1732195028106__000001000_0.jpg"  # 기본 이미지 이름
                 
                 try:
+                    # 실제 이미지 파일 확인
+                    fs = HfFileSystem(token=huggingface_token)
+                    samples_path = f"{model_id}/samples"
                     files = fs.ls(samples_path, detail=True)
-                    # 파일 이름만 추출
                     jpg_files = [
                         f['name'] for f in files 
                         if isinstance(f, dict) and 
@@ -590,29 +592,37 @@ def refresh_models(huggingface_token):
                     ]
                     
                     if jpg_files:
-                        # 숫자가 포함된 첫 번째 jpg 파일 사용
-                        image_url = f"https://huggingface.co/{model_id}/resolve/main/samples/{os.path.basename(jpg_files[0])}"
-                        print(f"Found image for {model_id}: {image_url}")
-                    else:
-                        # 기본 이미지 대신 첫 번째 발견된 jpg 사용
-                        image_url = f"https://huggingface.co/{model_id}/resolve/main/samples/1732195028106__000001000_0.jpg"
-                        print(f"Using default numbered image for {model_id}")
+                        base_image_name = os.path.basename(jpg_files[0])
                 except Exception as e:
                     print(f"Error accessing samples folder for {model_id}: {str(e)}")
-                    image_url = f"https://huggingface.co/{model_id}/resolve/main/samples/1732195028106__000001000_0.jpg"
                 
-                # 트리거 워드 추출
-                trigger_word = ""
-                if 'instance_prompt' in model_info:
-                    trigger_word = model_info['instance_prompt']
+                # 이미지 URL 구성 (토큰 포함)
+                if is_private:
+                    # Private 모델의 경우 로컬 캐시 경로 사용
+                    cache_dir = f"models/{model_id.replace('/', '_')}/samples"
+                    os.makedirs(cache_dir, exist_ok=True)
+                    
+                    # 이미지 다운로드
+                    image_url = f"https://huggingface.co/{model_id}/resolve/main/samples/{base_image_name}"
+                    local_image_path = os.path.join(cache_dir, base_image_name)
+                    
+                    if not os.path.exists(local_image_path):
+                        response = requests.get(image_url, headers=headers)
+                        if response.status_code == 200:
+                            with open(local_image_path, 'wb') as f:
+                                f.write(response.content)
+                    
+                    image_url = local_image_path
+                else:
+                    image_url = f"https://huggingface.co/{model_id}/resolve/main/samples/{base_image_name}"
                 
                 model_info = {
                     "image": image_url,
-                    "title": f"[Private] {model_id.split('/')[-1]}" if model.get('private') else model_id.split('/')[-1],
+                    "title": f"[Private] {model_id.split('/')[-1]}" if is_private else model_id.split('/')[-1],
                     "repo": model_id,
                     "weights": "pytorch_lora_weights.safetensors",
-                    "trigger_word": trigger_word,
-                    "private": model.get('private', False)
+                    "trigger_word": model_info.get('instance_prompt', ''),
+                    "private": is_private
                 }
                 new_models.append(model_info)
                 print(f"Added model: {model_id} with image: {image_url}")