Update app.py

app.py

@@ -5,15 +5,16 @@ import os
 
 # 환경변수에서 API 키 로드
 FIREWORKS_API_KEY = os.getenv("FIREWORKS_API_KEY", "")
+FAL_KEY = os.getenv("FAL_KEY", "")
 
-def enhance_prompt(prompt: str, api_key: str) -> str:
+def enhance_prompt(prompt: str, fireworks_key: str) -> str:
     """
     Fireworks AI LLM API를 사용하여 프롬프트를 증강합니다.
     """
     if not prompt.strip():
         return "❌ 프롬프트를 입력해주세요."
     
-    if not api_key.strip():
+    if not fireworks_key.strip():
         return "❌ Fireworks API 키를 입력해주세요."
     
     system_message = """You are a professional prompt engineer specializing in AI image and video generation. 
@@ -55,7 +56,7 @@ Respond ONLY with the enhanced prompt, no explanations or additional text."""
     headers = {
         "Accept": "application/json",
         "Content-Type": "application/json",
-        "Authorization": f"Bearer {api_key}"
+        "Authorization": f"Bearer {fireworks_key}"
     }
     
     try:
@@ -73,13 +74,147 @@ Respond ONLY with the enhanced prompt, no explanations or additional text."""
     except requests.exceptions.Timeout:
         return "❌ API 요청 시간 초과. 다시 시도해주세요."
     except requests.exceptions.HTTPError as e:
-        return f"❌ API 오류: {e.response.status_code} - {e.response.text}"
+        return f"❌ API 오류: {e.response.status_code}"
     except requests.exceptions.RequestException as e:
         return f"❌ 네트워크 오류: {str(e)}"
     except Exception as e:
         return f"❌ 오류 발생: {str(e)}"
 
 
+def generate_image(prompt: str, fal_key: str, aspect_ratio: str = "1:1", resolution: str = "1K") -> str:
+    """
+    FAL AI의 nano-banana-pro 모델을 사용하여 이미지를 생성합니다.
+    """
+    if not prompt.strip():
+        return None
+    
+    if not fal_key.strip():
+        return None
+    
+    url = "https://queue.fal.run/fal-ai/nano-banana-pro"
+    
+    headers = {
+        "Authorization": f"Key {fal_key}",
+        "Content-Type": "application/json"
+    }
+    
+    payload = {
+        "prompt": prompt,
+        "num_images": 1,
+        "aspect_ratio": aspect_ratio,
+        "output_format": "png",
+        "resolution": resolution
+    }
+    
+    try:
+        # Submit request
+        response = requests.post(url, headers=headers, json=payload, timeout=120)
+        response.raise_for_status()
+        
+        data = response.json()
+        
+        # 직접 결과 반환되는 경우
+        if "images" in data:
+            images = data.get("images", [])
+            if images and len(images) > 0:
+                return images[0].get("url", None)
+        
+        # Queue 방식인 경우 - request_id로 결과 조회
+        request_id = data.get("request_id")
+        if request_id:
+            return poll_for_result(request_id, fal_key)
+        
+        return None
+        
+    except Exception as e:
+        print(f"이미지 생성 오류: {e}")
+        return None
+
+
+def poll_for_result(request_id: str, fal_key: str, max_attempts: int = 60) -> str:
+    """
+    Queue 결과를 폴링하여 이미지 URL을 가져옵니다.
+    """
+    status_url = f"https://queue.fal.run/fal-ai/nano-banana-pro/requests/{request_id}/status"
+    result_url = f"https://queue.fal.run/fal-ai/nano-banana-pro/requests/{request_id}"
+    
+    headers = {
+        "Authorization": f"Key {fal_key}",
+        "Content-Type": "application/json"
+    }
+    
+    for attempt in range(max_attempts):
+        try:
+            # 상태 확인
+            status_response = requests.get(status_url, headers=headers, timeout=30)
+            status_data = status_response.json()
+            
+            status = status_data.get("status", "")
+            
+            if status == "COMPLETED":
+                # 결과 가져오기
+                result_response = requests.get(result_url, headers=headers, timeout=30)
+                result_data = result_response.json()
+                
+                images = result_data.get("images", [])
+                if images and len(images) > 0:
+                    return images[0].get("url", None)
+                return None
+            
+            elif status == "FAILED":
+                print(f"이미지 생성 실패: {status_data}")
+                return None
+            
+            # IN_QUEUE 또는 IN_PROGRESS - 대기
+            import time
+            time.sleep(2)
+            
+        except Exception as e:
+            print(f"폴링 오류: {e}")
+            import time
+            time.sleep(2)
+    
+    return None
+
+
+def process_comparison(
+    original_prompt: str,
+    fireworks_key: str,
+    fal_key: str,
+    aspect_ratio: str,
+    resolution: str
+):
+    """
+    원본 프롬프트 증강 → 두 이미지 생성 → 비교 결과 반환
+    """
+    # 1. 프롬프트 증강
+    yield "🔄 프롬프트 증강 중...", None, None, None
+    
+    enhanced = enhance_prompt(original_prompt, fireworks_key)
+    
+    if enhanced.startswith("❌"):
+        yield enhanced, None, None, None
+        return
+    
+    yield f"✅ 프롬프트 증강 완료\n\n🔄 원본 프롬프트로 이미지 생성 중...", enhanced, None, None
+    
+    # 2. 원본 프롬프트로 이미지 생성
+    original_image = generate_image(original_prompt, fal_key, aspect_ratio, resolution)
+    
+    if not original_image:
+        yield "✅ 프롬프트 증강 완료\n\n❌ 원본 이미지 생성 실패", enhanced, None, None
+    else:
+        yield f"✅ 프롬프트 증강 완료\n✅ 원본 이미지 생성 완료\n\n🔄 증강 프롬프트로 이미지 생성 중...", enhanced, original_image, None
+    
+    # 3. 증강 프롬프트로 이미지 생성
+    enhanced_image = generate_image(enhanced, fal_key, aspect_ratio, resolution)
+    
+    if not enhanced_image:
+        yield "✅ 프롬프트 증강 완료\n✅ 원본 이미지 생성 완료\n\n❌ 증강 이미지 생성 실패", enhanced, original_image, None
+    else:
+        yield "✅ 모든 작업 완료!", enhanced, original_image, enhanced_image
+
+
 # Example prompts
 examples = [
     ["한복을 입은 여성이 전통 한옥 마당에서 부채를 들고 있다"],
@@ -90,7 +225,7 @@ examples = [
 ]
 
 # Build the Gradio interface
-with gr.Blocks(title="AI 프롬프트 증강기", theme=gr.themes.Soft(primary_hue="teal", secondary_hue="cyan")) as demo:
+with gr.Blocks(title="AI 프롬프트 증강기 + 이미지 비교", theme=gr.themes.Soft(primary_hue="teal", secondary_hue="cyan")) as demo:
     
     gr.HTML("""
         <div style="text-align: center; padding: 20px 0;">
@@ -106,43 +241,45 @@ with gr.Blocks(title="AI 프롬프트 증강기", theme=gr.themes.Soft(primary_h
                     font-size: 24px;
                 ">✨</div>
                 <h1 style="
-                    font-size: 2.5rem;
+                    font-size: 2.2rem;
                     font-weight: 800;
                     background: linear-gradient(135deg, #4ecdc4 0%, #45b7d1 50%, #96e6a1 100%);
                     -webkit-background-clip: text;
                     -webkit-text-fill-color: transparent;
                     margin: 0;
-                ">AI 프롬프트 증강기</h1>
+                ">AI 프롬프트 증강기 + 이미지 비교</h1>
             </div>
             <p style="color: #6b7280; font-size: 1.1rem; margin: 0;">
-                이미지 & 영상 생성을 위한 프롬프트를 AI가 자동으로 증강합니다
+                원본 vs 증강 프롬프트의 이미지 생성 결과를 비교합니다
             </p>
         </div>
     """)
     
-    with gr.Row():
-        with gr.Column(scale=1):
-            api_key_input = gr.Textbox(
-                label="🔑 Fireworks API Key",
+    # API Keys Section
+    with gr.Accordion("🔑 API Keys 설정", open=not (FIREWORKS_API_KEY and FAL_KEY)):
+        with gr.Row():
+            fireworks_key_input = gr.Textbox(
+                label="Fireworks API Key (프롬프트 증강용)",
                 placeholder="fw_...",
                 type="password",
                 value=FIREWORKS_API_KEY,
-                info="Fireworks AI API 키를 입력하세요 (환경변수 FIREWORKS_API_KEY로도 설정 가능)"
+                scale=1
+            )
+            fal_key_input = gr.Textbox(
+                label="FAL API Key (이미지 생성용)",
+                placeholder="your-fal-key",
+                type="password",
+                value=FAL_KEY,
+                scale=1
             )
     
+    # Input Section
     with gr.Row():
-        with gr.Column(scale=1):
+        with gr.Column(scale=2):
             prompt_input = gr.Textbox(
                 label="📝 원본 프롬프트",
-                placeholder="이미지나 영상으로 생성하고 싶은 장면을 설명해주세요...\n예: 한복을 입은 여성이 벚꽃 아래에서 걷고 있다",
-                lines=5,
-                max_lines=10
-            )
-            
-            enhance_btn = gr.Button(
-                "✨ 프롬프트 증강하기",
-                variant="primary",
-                size="lg"
+                placeholder="이미지로 생성하고 싶은 장면을 설명해주세요...\n예: 한복을 입은 여성이 벚꽃 아래에서 걷고 있다",
+                lines=4
             )
             
             gr.Examples(
@@ -150,38 +287,76 @@ with gr.Blocks(title="AI 프롬프트 증강기", theme=gr.themes.Soft(primary_h
                 inputs=prompt_input,
                 label="💡 예시 프롬프트"
             )
+        
+        with gr.Column(scale=1):
+            aspect_ratio = gr.Dropdown(
+                label="📐 비율",
+                choices=["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3", "21:9"],
+                value="1:1"
+            )
+            resolution = gr.Dropdown(
+                label="🖼️ 해상도",
+                choices=["1K", "2K", "4K"],
+                value="1K"
+            )
+            
+            generate_btn = gr.Button(
+                "🚀 증강 + 이미지 생성",
+                variant="primary",
+                size="lg"
+            )
+    
+    # Status
+    status_text = gr.Textbox(
+        label="📊 진행 상태",
+        interactive=False,
+        lines=3
+    )
+    
+    # Enhanced Prompt Output
+    enhanced_output = gr.Textbox(
+        label="✨ 증강된 프롬프트",
+        placeholder="증강된 프롬프트가 여기에 표시됩니다...",
+        lines=4,
+        interactive=True
+    )
+    
+    # Image Comparison Section
+    gr.HTML("""
+        <div style="text-align: center; padding: 15px 0 5px 0;">
+            <h3 style="color: #374151; margin: 0;">🖼️ 이미지 비교</h3>
+        </div>
+    """)
     
     with gr.Row():
-        with gr.Column(scale=1):
-            output_text = gr.Textbox(
-                label="🚀 증강된 프롬프트 (복사하여 사용하세요)",
-                placeholder="증강된 프롬프트가 여기에 표시됩니다...",
-                lines=8,
-                max_lines=15,
-                interactive=True
+        with gr.Column():
+            gr.HTML("<p style='text-align: center; font-weight: 600; color: #6b7280;'>📌 원본 프롬프트 결과</p>")
+            original_image_output = gr.Image(
+                label="원본 프롬프트 이미지",
+                type="filepath",
+                show_label=False
+            )
+        
+        with gr.Column():
+            gr.HTML("<p style='text-align: center; font-weight: 600; color: #10b981;'>✨ 증강 프롬프트 결과</p>")
+            enhanced_image_output = gr.Image(
+                label="증강 프롬프트 이미지", 
+                type="filepath",
+                show_label=False
             )
     
     gr.HTML("""
         <div style="text-align: center; padding: 20px 0; color: #6b7280; font-size: 0.9rem;">
-            <p>💡 <strong>팁:</strong> 증강된 프롬프트는 영어로 출력되어 대부분의 AI 이미지/영상 생성 모델과 호환됩니다.</p>
-            <p style="margin-top: 8px;">Powered by Fireworks AI (gpt-oss-120b) | Built with Gradio</p>
+            <p>💡 <strong>팁:</strong> 증강된 프롬프트는 더 상세한 시각적 요소를 포함하여 더 풍부한 이미지를 생성합니다.</p>
+            <p style="margin-top: 8px;">Powered by Fireworks AI (gpt-oss-120b) + FAL AI (nano-banana-pro)</p>
         </div>
     """)
     
-    # Event handlers
-    enhance_btn.click(
-        fn=enhance_prompt,
-        inputs=[prompt_input, api_key_input],
-        outputs=output_text,
-        show_progress=True
-    )
-    
-    # Allow Enter key to submit (Shift+Enter for new line)
-    prompt_input.submit(
-        fn=enhance_prompt,
-        inputs=[prompt_input, api_key_input],
-        outputs=output_text,
-        show_progress=True
+    # Event handler
+    generate_btn.click(
+        fn=process_comparison,
+        inputs=[prompt_input, fireworks_key_input, fal_key_input, aspect_ratio, resolution],
+        outputs=[status_text, enhanced_output, original_image_output, enhanced_image_output]
     )