Update app.py

app.py

@@ -12,10 +12,29 @@ import torch
 from diffusers import DiffusionPipeline
 from PIL import Image
 
+# ===== Fix: Import spaces for Hugging Face Spaces =====
+try:
+    import spaces
+    HF_SPACES = True
+except ImportError:
+    # If running locally, create a dummy decorator
+    def spaces_gpu_decorator(duration=60):
+        def decorator(func):
+            return func
+        return decorator
+    spaces = type('spaces', (), {'GPU': spaces_gpu_decorator})()
+    HF_SPACES = False
+    print("Warning: Running without Hugging Face Spaces GPU allocation")
+
 # ===== OpenAI 설정 =====
 from openai import OpenAI
 
-client = OpenAI(api_key=os.getenv("LLM_API"))  # 환경 변수에 API 키가 있어야 합니다.
+# Add error handling for API key
+try:
+    client = OpenAI(api_key=os.getenv("LLM_API"))
+except Exception as e:
+    print(f"Warning: OpenAI client initialization failed: {e}")
+    client = None
 
 # ===== 프롬프트 증강용 스타일 프리셋 =====
 STYLE_PRESETS = {
@@ -27,18 +46,26 @@ STYLE_PRESETS = {
 }
 
 # ===== 저장 폴더 =====
-SAVE_DIR = "saved_images"  # Gradio will handle the persistence
+SAVE_DIR = "saved_images"
 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 = "seawolf2357/kim-korea"  # 특정 정치인을 학습한 LoRA 모델
+print(f"Using device: {device}")
 
-pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
-pipeline.load_lora_weights(adapter_id)
-pipeline = pipeline.to(device)
+repo_id = "black-forest-labs/FLUX.1-dev"
+adapter_id = "seawolf2357/kim-korea"
+
+# Add error handling for model loading
+try:
+    pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
+    pipeline.load_lora_weights(adapter_id)
+    pipeline = pipeline.to(device)
+    print("Model loaded successfully")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    pipeline = None
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -55,11 +82,15 @@ def openai_translate(text: str, retries: int = 3) -> str:
     """한글을 영어로 번역 (OpenAI GPT-4.1-mini 사용). 영어 입력이면 그대로 반환."""
     if not is_korean(text):
         return text
+    
+    if client is None:
+        print("Warning: OpenAI client not available, returning original text")
+        return text
 
     for attempt in range(retries):
         try:
             res = client.chat.completions.create(
-                model="gpt-4.1-mini",
+                model="gpt-4o-mini",  # Fixed: gpt-4.1-mini doesn't exist, use gpt-4o-mini
                 messages=[
                     {
                         "role": "system",
@@ -71,7 +102,7 @@ def openai_translate(text: str, retries: int = 3) -> str:
                 max_tokens=256,
             )
             return res.choices[0].message.content.strip()
-        except (requests.exceptions.RequestException, Exception) as e:
+        except Exception as e:
             print(f"[translate] attempt {attempt + 1} failed: {e}")
             time.sleep(2)
     return text  # 번역 실패 시 원문 그대로
@@ -104,6 +135,9 @@ def save_generated_image(image: Image.Image, prompt: str) -> str:
 # ===== Diffusion 호출 =====
 
 def run_pipeline(prompt: str, seed: int, width: int, height: int, guidance_scale: float, num_steps: int, lora_scale: float):
+    if pipeline is None:
+        raise ValueError("Model pipeline not loaded")
+    
     generator = torch.Generator(device=device).manual_seed(int(seed))
     result = pipeline(
         prompt=prompt,
@@ -131,19 +165,27 @@ def generate_image(
     lora_scale: float = 1.0,
     progress=None,
 ):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+    try:
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
 
-    # 1) 번역 + 증강
-    final_prompt = prepare_prompt(user_prompt, style_key)
+        # 1) 번역 + 증강
+        final_prompt = prepare_prompt(user_prompt, style_key)
+        print(f"Final prompt: {final_prompt}")
 
-    # 2) 파이프라인 호출
-    image = run_pipeline(final_prompt, seed, width, height, guidance_scale, num_inference_steps, lora_scale)
+        # 2) 파이프라인 호출
+        image = run_pipeline(final_prompt, seed, width, height, guidance_scale, num_inference_steps, lora_scale)
 
-    # 3) 저장
-    save_generated_image(image, final_prompt)
+        # 3) 저장
+        save_generated_image(image, final_prompt)
 
-    return image, seed
+        return image, seed
+    
+    except Exception as e:
+        print(f"Error generating image: {e}")
+        # Return a placeholder or error message
+        error_image = Image.new('RGB', (width, height), color='red')
+        return error_image, seed
 
 # ===== 예시 프롬프트 (한국어/영어 혼용 허용) =====
 
@@ -154,12 +196,13 @@ examples = [
     "Mr. KIM이 붐비는 거리에서 여성 시민들과 따뜻하게 악수하는 ��습, 여성 유권자들에 대한 진정한 관심과 소통을 보여주고 있다.",
     "Mr. KIM이 선거 유세장에서 지평선을 향해 손가락으로 가리키며 영감을 주는 제스처를 취하고 있고, 여성들과 아이들이 박수를 치고 있다.",
     "Mr. KIM이 지역 행사에 참여하여 열정적으로 응원하는 여성 지지자들에게 둘러싸여 있는 모습.",
-    "Mr. KIM visiting a local market, engaging in friendly conversation with female vendors and shopkeepers. ",
-    "Mr. KIM walking through a university campus, discussing education policies with female students and professors. ",    
-    "Mr. KIM delivering a powerful speech in front of a large crowd with confident gestures and determined expression. ",
+    "Mr. KIM visiting a local market, engaging in friendly conversation with female vendors and shopkeepers.",
+    "Mr. KIM walking through a university campus, discussing education policies with female students and professors.",    
+    "Mr. KIM delivering a powerful speech in front of a large crowd with confident gestures and determined expression.",
     "Mr. KIM in a dynamic interview setting, passionately outlining his visions for the future.",
-    "Mr. KIM preparing for an important debate, surrounded by paperwork, looking focused and resolute. ",
+    "Mr. KIM preparing for an important debate, surrounded by paperwork, looking focused and resolute.",
 ]
+
 # ===== 커스텀 CSS (붉은 톤 유지) =====
 custom_css = """
 :root {
@@ -183,63 +226,68 @@ button:hover{transform:translateY(-2px); box-shadow:0 5px 15px rgba(0,0,0,.1);}
 """
 
 # ===== Gradio UI =====
-with gr.Blocks(css=custom_css, analytics_enabled=False) as demo:
-    gr.HTML('<div class="title">Mr. KIM in KOREA</div>')
-    gr.HTML('<div class="collection-link"><a href="https://huggingface.co/collections/openfree/painting-art-ai-681453484ec15ef5978bbeb1" target="_blank">Visit the LoRA Model Collection</a></div>')
-
-    with gr.Group(elem_classes="model-description"):
-        gr.HTML("""
-        <p>
-        본 모델은 연구 목적으로 특정인의 얼굴과 외모를 학습한 LoRA 모델입니다.<br>
-        목적외의 용도로 무단 사용 않도록 유의해 주세요.<br>
-        (예시 prompt 사용 시 반드시 'kim'을 포함하여야 최적의 결과를 얻을 수 있습니다.)
-        </p>
-        """)
-
-    # ===== 메인 입력 =====
-    with gr.Column():
-        with gr.Row(elem_classes="input-container"):
-            user_prompt = gr.Text(label="Prompt", max_lines=1, value=examples[0])
-            style_select = gr.Radio(label="Style Preset", choices=list(STYLE_PRESETS.keys()), value="None", interactive=True)
-            run_button = gr.Button("Generate", variant="primary")
-
-        result_image = gr.Image(label="Generated Image")
-        seed_output = gr.Number(label="Seed")
-
-        # ===== 고급 설정 =====
-        with gr.Accordion("Advanced Settings", open=False, elem_classes="advanced-settings"):
-            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-            with gr.Row():
-                width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
-                height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=768)
-            with gr.Row():
-                guidance_scale = gr.Slider(label="Guidance scale", minimum=0.0, maximum=10.0, step=0.1, value=3.5)
-                num_inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=50, step=1, value=30)
-                lora_scale = gr.Slider(label="LoRA scale", minimum=0.0, maximum=1.0, step=0.1, value=1.0)
-
-        # ===== 예시 영역 =====
-        with gr.Group(elem_classes="example-region"):
-            gr.Markdown("### Examples")
-            gr.Examples(examples=examples, inputs=user_prompt, cache_examples=False)
-
-    # ===== 이벤트 =====
-    run_button.click(
-        fn=generate_image,
-        inputs=[
-            user_prompt,
-            style_select,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-            lora_scale,
-        ],
-        outputs=[result_image, seed_output],
-    )
-
-
-demo.queue()
-demo.launch()
+def create_interface():
+    with gr.Blocks(css=custom_css, analytics_enabled=False) as demo:
+        gr.HTML('<div class="title">Mr. KIM in KOREA</div>')
+        gr.HTML('<div class="collection-link"><a href="https://huggingface.co/collections/openfree/painting-art-ai-681453484ec15ef5978bbeb1" target="_blank">Visit the LoRA Model Collection</a></div>')
+
+        with gr.Group(elem_classes="model-description"):
+            gr.HTML("""
+            <p>
+            본 모델은 연구 목적으로 특정인의 얼굴과 외모를 학습한 LoRA 모델입니다.<br>
+            목적외의 용도로 무단 사용 않도록 유의해 주세��.<br>
+            (예시 prompt 사용 시 반드시 'kim'을 포함하여야 최적의 결과를 얻을 수 있습니다.)
+            </p>
+            """)
+
+        # ===== 메인 입력 =====
+        with gr.Column():
+            with gr.Row(elem_classes="input-container"):
+                user_prompt = gr.Text(label="Prompt", max_lines=1, value=examples[0])
+                style_select = gr.Radio(label="Style Preset", choices=list(STYLE_PRESETS.keys()), value="None", interactive=True)
+                run_button = gr.Button("Generate", variant="primary")
+
+            result_image = gr.Image(label="Generated Image")
+            seed_output = gr.Number(label="Seed")
+
+            # ===== 고급 설정 =====
+            with gr.Accordion("Advanced Settings", open=False, elem_classes="advanced-settings"):
+                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                with gr.Row():
+                    width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
+                    height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=768)
+                with gr.Row():
+                    guidance_scale = gr.Slider(label="Guidance scale", minimum=0.0, maximum=10.0, step=0.1, value=3.5)
+                    num_inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=50, step=1, value=30)
+                    lora_scale = gr.Slider(label="LoRA scale", minimum=0.0, maximum=1.0, step=0.1, value=1.0)
+
+            # ===== 예시 영역 =====
+            with gr.Group(elem_classes="example-region"):
+                gr.Markdown("### Examples")
+                gr.Examples(examples=examples, inputs=user_prompt, cache_examples=False)
+
+        # ===== 이벤트 =====
+        run_button.click(
+            fn=generate_image,
+            inputs=[
+                user_prompt,
+                style_select,
+                seed,
+                randomize_seed,
+                width,
+                height,
+                guidance_scale,
+                num_inference_steps,
+                lora_scale,
+            ],
+            outputs=[result_image, seed_output],
+        )
+    
+    return demo
+
+# ===== 애플리케이션 실행 =====
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.queue()
+    demo.launch()