Update app-backup1.py

app-backup1.py

@@ -19,6 +19,8 @@ from gradio_imageslider import ImageSlider
 import numpy as np
 import warnings
 
+# 상단에 허깅페이스 USERNAME (해당 계정) 반드시 개별 지정할것
+USERNAME = "openfree"
 
 huggingface_token = os.getenv("HF_TOKEN")
 
@@ -325,7 +327,7 @@ def remove_custom_lora(selected_indices, current_loras):
         lora_image_3
     )
 
-@spaces.GPU(duration=75)
+
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress):
     print("Generating image...")
     pipe.to("cuda")
@@ -345,7 +347,7 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
         ):
             yield img
 
-@spaces.GPU(duration=75)
+
 def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, seed):
     pipe_i2i.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(seed)
@@ -364,9 +366,11 @@ 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']
             print(f"Original prompt: {prompt}")
@@ -378,7 +382,7 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
 
         selected_loras = [loras_state[idx] for idx in selected_indices]
 
-        # Build the prompt with trigger words (이 부분은 그대로 유지)
+        # Build the prompt with trigger words
         prepends = []
         appends = []
         for lora in selected_loras:
@@ -401,27 +405,40 @@ 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']
-                    weight_name = lora.get("weights")
-                    print(f"Loading LoRA {lora_name} from {lora_path}")
+            
+                    # Private 모델인 경우 특별 처리
+                    if lora.get('private', False):
+                        lora_path = load_private_model(lora_path, huggingface_token)
+                        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)
-                        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)
-                        else:
-                            pipe.load_lora_weights(lora_path, adapter_name=lora_name)
+                        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)
@@ -437,11 +454,12 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
 
         print(f"Active adapters after loading: {pipe.get_active_adapters()}")
 
-        # 여기서부터 이미지 생성 로직 (이 부분은 그대로 유지)
+        # Randomize seed if needed
         with calculateDuration("Randomizing seed"):
             if randomize_seed:
                 seed = random.randint(0, MAX_SEED)
 
+        # Generate image
         if image_input is not None:
             final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, seed)
         else:
@@ -522,6 +540,138 @@ def update_history(new_image, history):
         history.insert(0, new_image)
     return history
 
+
+
+def refresh_models(huggingface_token):
+    try:
+        headers = {
+            "Authorization": f"Bearer {huggingface_token}",
+            "Accept": "application/json"
+        }
+        
+        username = USERNAME
+        api_url = f"https://huggingface.co/api/models?author={username}"
+        response = requests.get(api_url, headers=headers)
+        if response.status_code != 200:
+            raise Exception(f"Failed to fetch models from HuggingFace. Status code: {response.status_code}")
+            
+        all_models = response.json()
+        print(f"Found {len(all_models)} models for user {username}")
+        
+        user_models = [
+            model for model in all_models 
+            if model.get('tags') and ('flux' in [tag.lower() for tag in model.get('tags', [])] or
+                                    'flux-lora' in [tag.lower() for tag in model.get('tags', [])])
+        ]
+        
+        print(f"Found {len(user_models)} FLUX models")
+        
+        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에 토큰을 포함시키는 방식으로 변경
+                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 
+                        'name' in f and 
+                        f['name'].lower().endswith('.jpg') and 
+                        any(char.isdigit() for char in os.path.basename(f['name']))
+                    ]
+                    
+                    if jpg_files:
+                        base_image_name = os.path.basename(jpg_files[0])
+                except Exception as e:
+                    print(f"Error accessing samples folder for {model_id}: {str(e)}")
+                
+                # 이미지 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 is_private else model_id.split('/')[-1],
+                    "repo": model_id,
+                    "weights": "pytorch_lora_weights.safetensors",
+                    "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}")
+                
+            except Exception as e:
+                print(f"Error processing model {model['id']}: {str(e)}")
+                continue
+        
+        updated_loras = new_models + [lora for lora in loras if lora['repo'] not in [m['repo'] for m in new_models]]
+        
+        print(f"Total models after refresh: {len(updated_loras)}")
+        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
+        )
+        
+        # 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
+
 custom_theme = gr.themes.Base(
     primary_hue="blue",
     secondary_hue="purple",
@@ -825,6 +975,25 @@ input:focus, textarea:focus {
     max-width: 90% !important;
     margin: 0 !important;  /* auto에서 0으로 변경 */
     margin-left: 20px !important;  /* 왼쪽 여백 추가 */
+
+/* 새로고침 버튼 스타일 */
+#refresh-button {
+    margin: 10px;
+    padding: 8px 16px;
+    background-color: #4a5568;
+    color: white;
+    border-radius: 8px;
+    transition: all 0.3s ease;
+}
+
+#refresh-button:hover {
+    background-color: #2d3748;
+    transform: scale(1.05);
+}
+
+#refresh-button:active {
+    transform: scale(0.95);
+}    
 }
 '''
 
@@ -839,6 +1008,9 @@ with gr.Blocks(theme=custom_theme, css=css, delete_cache=(60, 3600)) as app:
         갤러리에서 원하는 모델을 선택(최대 3개까지) < 프롬프트에 한글 또는 영문으로 원하는 내용을 입력 < Generate 버튼 실행  
         """
     )
+    # 새로고침 버튼 추가
+    with gr.Row():
+        refresh_button = gr.Button("🔄 모델 새로고침(나만의 맞춤 학습된 Private 모델 불러오기)", variant="secondary")
 
     with gr.Row(elem_id="lora_gallery", equal_height=True):
         gallery = gr.Gallery(
@@ -855,6 +1027,7 @@ with gr.Blocks(theme=custom_theme, css=css, delete_cache=(60, 3600)) as app:
             preview=False
         )
 
+
     with gr.Tab(label="Generate"):
         # Prompt and Generate Button
         with gr.Row():
@@ -1023,6 +1196,19 @@ with gr.Blocks(theme=custom_theme, css=css, delete_cache=(60, 3600)) as app:
             outputs=history_gallery
         )
 
+# 새로고침 버튼 이벤트 핸들러
+    def refresh_gallery():
+        updated_loras = refresh_models(huggingface_token)
+        return (
+            gr.update(value=[(item["image"], item["title"]) for item in updated_loras]),
+            updated_loras
+        )
+
+    refresh_button.click(
+        refresh_gallery,
+        outputs=[gallery, loras_state]
+    )
+
 if __name__ == "__main__":
     app.queue(max_size=20)
     app.launch(debug=True)