Update app.py

app.py

@@ -1,3 +1,4 @@
+
 import random
 import os
 import uuid
@@ -8,6 +9,14 @@ import spaces
 import torch
 from diffusers import DiffusionPipeline
 from PIL import Image
+import huggingface_hub
+import requests
+from tqdm.auto import tqdm
+import time
+
+# 타임아웃 값 증가 설정
+huggingface_hub.constants.DEFAULT_ETAG_TIMEOUT = 30
+huggingface_hub.constants.DEFAULT_DOWNLOAD_TIMEOUT = 120
 
 # Temporary fix to patch the gradio_client.utils module
 import gradio_client.utils
@@ -25,13 +34,37 @@ SAVE_DIR = "saved_images"  # Gradio will handle the persistence
 if not os.path.exists(SAVE_DIR):
     os.makedirs(SAVE_DIR, exist_ok=True)
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-repo_id = "black-forest-labs/FLUX.1-dev"
-adapter_id = "openfree/flux-chatgpt-ghibli-lora"
+# 모델 로딩 함수 - 재시도 로직 추가
+def load_model_with_retry(max_retries=5):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    repo_id = "black-forest-labs/FLUX.1-dev"
+    adapter_id = "openfree/flux-chatgpt-ghibli-lora"
+    
+    for attempt in range(max_retries):
+        try:
+            print(f"Loading model attempt {attempt+1}/{max_retries}...")
+            pipeline = DiffusionPipeline.from_pretrained(
+                repo_id, 
+                torch_dtype=torch.bfloat16,
+                use_safetensors=True,
+                resume_download=True
+            )
+            print("Model loaded successfully, loading LoRA weights...")
+            pipeline.load_lora_weights(adapter_id)
+            pipeline = pipeline.to(device)
+            print("Pipeline ready!")
+            return pipeline, device
+        except (requests.exceptions.ReadTimeout, requests.exceptions.ConnectionError) as e:
+            if attempt < max_retries - 1:
+                wait_time = 10 * (attempt + 1)  # 점진적으로 대기 시간 증가
+                print(f"Download timed out or connection error: {e}. Retrying in {wait_time} seconds...")
+                time.sleep(wait_time)
+            else:
+                raise Exception(f"Failed to download model after {max_retries} attempts: {e}")
 
-pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
-pipeline.load_lora_weights(adapter_id)
-pipeline = pipeline.to(device)
+# 모델 로드 시작
+print("Starting model loading process...")
+pipeline, device = load_model_with_retry()
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -233,5 +266,12 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css, analytics_enabled=Fa
         outputs=[result, seed, generated_gallery],
     )
 
+# Launch with explicit host and port
 demo.queue()
-demo.launch()
+try:
+    demo.launch(share=False)
+except Exception as e:
+    print(f"Error during launch: {e}")
+    # 에러 발생 시, 간소화된 버전으로 재시도
+    print("Retrying with simplified launch...")
+    demo.launch(share=False, ssl_verify=False)