Update app.py

app.py

@@ -1,24 +1,13 @@
 import gradio as gr
 import google.generativeai as genai
-from google.generativeai import types
 from PIL import Image
 import os
 import json
-import tempfile
-import logging
-import re
-import time
-import base64
-from io import BytesIO
-
-# 로깅 설정
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
 
 # Gemini API 키 설정 (환경 변수에서 가져오거나 직접 입력)
 GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
 
-# Gemini API 초기화 (아래 함수에서 재설정됨)
+# Gemini API 초기화
 genai.configure(api_key=GEMINI_API_KEY)
 
 # 배경 JSON 파일 경로 설정 - 상대 경로 사용
@@ -160,7 +149,7 @@ def generate_prompt_with_gemini(product_name, background_info, additional_info="
         )
         response = model.generate_content(
             prompt_request,
-            generation_config=genai.GenerationConfig(
+            generation_config=genai.types.GenerationConfig(
                 temperature=0.7,
                 top_p=0.95,
                 top_k=64,
@@ -174,214 +163,10 @@ def generate_prompt_with_gemini(product_name, background_info, additional_info="
     except Exception as e:
         return f"프롬프트 생성 중 오류가 발생했습니다: {str(e)}"
 
-# --- 이미지 생성 관련 함수들 (수정 적용) ---
-
-def save_binary_file(file_name, data):
-    with open(file_name, "wb") as f:
-        f.write(data)
-
-def translate_prompt_to_english(prompt):
-    # 기존 함수를 수정
-    if not re.search("[가-힣]", prompt):
-        return prompt
-    
-    prompt = prompt.replace("#1", "IMAGE_TAG_ONE")
-    
-    try:
-        api_key = os.environ.get("GEMINI_API_KEY")
-        if not api_key:
-            logger.error("Gemini API 키가 설정되지 않았습니다.")
-            prompt = prompt.replace("IMAGE_TAG_ONE", "#1")
-            return prompt
-            
-        genai.configure(api_key=api_key)
-        model = genai.GenerativeModel('gemini-2.0-flash')
-        
-        translation_prompt = f"""
-        Translate the following Korean text to English:
-        
-        {prompt}
-        
-        IMPORTANT: The token IMAGE_TAG_ONE is a special tag 
-        and must be preserved exactly as is in your translation. Do not translate this token.
-        """
-        
-        logger.info(f"Translation prompt: {translation_prompt}")
-        response = model.generate_content(translation_prompt)
-        
-        translated_text = response.text
-        
-        if translated_text.strip():
-            translated_text = translated_text.replace("IMAGE_TAG_ONE", "#1")
-            logger.info(f"Translated text: {translated_text.strip()}")
-            return translated_text.strip()
-        else:
-            logger.warning("번역 결과가 없습니다. 원본 프롬프트 사용")
-            prompt = prompt.replace("IMAGE_TAG_ONE", "#1")
-            return prompt
-    except Exception as e:
-        logger.exception("번역 중 오류 발생:")
-        prompt = prompt.replace("IMAGE_TAG_ONE", "#1")
-        return prompt
-
-def preprocess_prompt(prompt, image1, image2, image3):
-    has_img1 = image1 is not None
-    
-    if "#1" in prompt and not has_img1:
-        prompt = prompt.replace("#1", "첫 번째 이미지(없음)")
-    else:
-        prompt = prompt.replace("#1", "첫 번째 이미지")
-    
-    prompt += " 이미지를 생성해주세요. 이미지에 텍스트나 글자를 포함하지 마세요."
-    return prompt
-
-# --- 수정된 이미지 생성 함수 ---
-def generate_with_images(prompt, images, variation_index=0):
-    try:
-        api_key = os.environ.get("GEMINI_API_KEY")
-        if not api_key:
-            return None, "API 키가 설정되지 않았습니다. 환경변수를 확인해주세요."
-        genai.configure(api_key=api_key)
-        model = genai.GenerativeModel('gemini-2.0-flash-exp-image-generation')
-        logger.info(f"Gemini API 요청 시작 - 프롬프트: {prompt}, 변형 인덱스: {variation_index}")
-
-        variation_suffixes = [
-            " Create this as the first variation. Do not add any text, watermarks, or labels to the image.",
-            " Create this as the second variation with more vivid colors. Do not add any text, watermarks, or labels to the image.",
-            " Create this as the third variation with a more creative style. Do not add any text, watermarks, or labels to the image.",
-            " Create this as the fourth variation with enhanced details. Do not add any text, watermarks, or labels to the image."
-        ]
-        
-        if variation_index < len(variation_suffixes):
-            prompt = prompt + variation_suffixes[variation_index]
-        else:
-            prompt = prompt + " Do not add any text, watermarks, or labels to the image."
-        
-        # 이미지 객체를 그대로 contents에 추가 (base64 인코딩 없이)
-        contents = [prompt]
-        for idx, img in enumerate(images, 1):
-            if img is not None:
-                contents.append(img)
-                logger.info(f"이미지 #{idx} 추가됨")
-        
-        response = model.generate_content(
-            contents,
-            generation_config=genai.GenerationConfig(
-                temperature=1,
-                top_p=0.95,
-                top_k=40,
-                max_output_tokens=8192
-            )
-        )
-        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
-            temp_path = tmp.name
-            result_text = ""
-            image_found = False
-            for part in response.candidates[0].content.parts:
-                if hasattr(part, 'text') and part.text:
-                    result_text += part.text
-                    logger.info(f"응답 텍스트: {part.text}")
-                elif hasattr(part, 'inline_data') and part.inline_data:
-                    save_binary_file(temp_path, part.inline_data.data)
-                    image_found = True
-                    logger.info("응답에서 이미지 추출 성공")
-            if not image_found:
-                return None, f"API에서 이미지를 생성하지 못했습니다. 응답 텍스트: {result_text}"
-            
-            result_img = Image.open(temp_path)
-            if result_img.mode == "RGBA":
-                result_img = result_img.convert("RGB")
-            result_img.save(temp_path, format="JPEG", quality=95)
-            
-            return temp_path, f"이미지가 성공적으로 생성되었습니다. {result_text}"
-    except Exception as e:
-        logger.exception("이미지 생성 중 오류 발생:")
-        return None, f"오류 발생: {str(e)}"
-
-def process_images_with_prompt(image1, image2, image3, prompt, variation_index=0, max_retries=3):
-    retry_count = 0
-    last_error = None
-    
-    while retry_count < max_retries:
-        try:
-            images = [image1, image2, image3]
-            valid_images = [img for img in images if img is not None]
-            if not valid_images:
-                return None, "적어도 하나의 이미지를 업로드해주세요.", ""
-
-            if prompt and prompt.strip():
-                processed_prompt = preprocess_prompt(prompt, image1, image2, image3)
-                if re.search("[가-힣]", processed_prompt):
-                    final_prompt = translate_prompt_to_english(processed_prompt)
-                else:
-                    final_prompt = processed_prompt
-            else:
-                if len(valid_images) == 1:
-                    final_prompt = "Please creatively transform this image into a more vivid and artistic version. Do not include any text or watermarks in the generated image."
-                    logger.info("Default prompt generated for single image")
-                elif len(valid_images) == 2:
-                    final_prompt = "Please seamlessly composite these two images, integrating their key elements harmoniously into a single image. Do not include any text or watermarks in the generated image."
-                    logger.info("Default prompt generated for two images")
-                else:
-                    final_prompt = "Please creatively composite these three images, combining their main elements into a cohesive and natural scene. Do not include any text or watermarks in the generated image."
-                    logger.info("Default prompt generated for three images")
-
-            result_img, status = generate_with_images(final_prompt, valid_images, variation_index)
-            if result_img is not None:
-                return result_img, status, final_prompt
-            else:
-                last_error = status
-                retry_count += 1
-                logger.warning(f"이미지 생성 실패, 재시도 {retry_count}/{max_retries}: {status}")
-                time.sleep(1)
-        except Exception as e:
-            last_error = str(e)
-            retry_count += 1
-            logger.exception(f"이미지 처리 중 오류 발생, 재시도 {retry_count}/{max_retries}:")
-            time.sleep(1)
-    
-    return None, f"최대 재시도 횟수({max_retries}회) 초과 후 실패: {last_error}", prompt
-
-def generate_multiple_images(image1, image2, image3, prompt, progress=gr.Progress()):
-    results = []
-    statuses = []
-    prompts = []
-    
-    num_images = 4
-    max_retries = 3
-    
-    progress(0, desc="이미지 생성 준비 중...")
-    
-    for i in range(num_images):
-        progress((i / num_images), desc=f"{i+1}/{num_images} 이미지 생성 중...")
-        result_img, status, final_prompt = process_images_with_prompt(image1, image2, image3, prompt, i, max_retries)
-        
-        if result_img is not None:
-            results.append(result_img)
-            statuses.append(f"이미지 #{i+1}: {status}")
-            prompts.append(f"이미지 #{i+1}: {final_prompt}")
-        else:
-            results.append(None)
-            statuses.append(f"이미지 #{i+1} 생성 실패: {status}")
-            prompts.append(f"이미지 #{i+1}: {final_prompt}")
-        
-        time.sleep(1)
-    
-    progress(1.0, desc="이미지 생성 완료!")
-    
-    while len(results) < 4:
-        results.append(None)
-    
-    combined_status = "\n".join(statuses)
-    combined_prompts = "\n".join(prompts)
-    
-    return results[0], results[1], results[2], results[3], combined_status, combined_prompts
-
-# --- 앱 UI 구성 ---
 def create_app():
-    with gr.Blocks(title="이미지 생성 애플리케이션") as demo:
-        gr.Markdown("# 상품 이미지 배경 프롬프트 생성기 & 이미지 생성기")
-        gr.Markdown("상품 이미지를 업로드하고 옵션을 선택하면 고품질 상업용 프롬프트가 생성되고 이미지가 생성됩니다.")
+    with gr.Blocks(title="고급 상품 이미지 배경 프롬프트 생성기") as demo:
+        gr.Markdown("# 고급 상품 이미지 배경 프롬프트 생성기")
+        gr.Markdown("상품 이미지를 업로드하고 옵션을 선택하면 고품질 상업용 프롬프트가 생성됩니다.")
         
         # 상품명 섹션
         with gr.Row():
@@ -453,26 +238,13 @@ def create_app():
                     interactive=True
                 )
                 
-                # 버튼 영역
-                with gr.Row():
-                    generate_prompt_btn = gr.Button("프롬프트 생성")
-                    submit_single_btn = gr.Button("이미지 생성 (1장)")
-                    submit_btn = gr.Button("이미지 생성 (4장)")
+                submit_btn = gr.Button("프롬프트 생성", variant="primary")
             
             # 우측 컬럼 - 출력 결과
             with gr.Column(scale=1):
                 prompt_output = gr.Textbox(label="생성된 프롬프트", lines=10)
-                
-                # 이미지 출력 영역
-                gr.Markdown("## 생성된 이미지")
-                with gr.Row():
-                    output_image1 = gr.Image(label="이미지 #1", type="pil")
-                    output_image2 = gr.Image(label="이미지 #2", type="pil")
-                with gr.Row():
-                    output_image3 = gr.Image(label="이미지 #3", type="pil")
-                    output_image4 = gr.Image(label="이미지 #4", type="pil")
-                    
-                output_text = gr.Textbox(label="상태 메시지", lines=2)
+                image_preview = gr.Image(label="업로드된 이미지 (#1)", type="pil")
+                preview_html = gr.HTML("프롬프트 생성 시 여기에 미리보기가 표시됩니다.")
         
         # 배경 유형에 따라 드롭다운 표시 업데이트
         def update_dropdowns(bg_type):
@@ -531,9 +303,12 @@ def create_app():
                 }
         
         # 프롬프트 생성 함수
-        def generate_prompt(bg_type, simple, studio, nature, indoor, abstract, product_text, additional_text):
-            if not product_text.strip():
-                product_text = "제품"
+        def generate_output(image, bg_type, simple, studio, nature, indoor, abstract, product_text, additional_text):
+            if image is None:
+                gr.Warning("이미지를 업로드해주세요.")
+                return "이미지를 업로드해주세요.", None, "이미지를 업로드 후 프롬프트를 생성해주세요."
+            
+            product_text = product_text.strip() or "제품"
             
             # 배경 정보 가져오기
             background_info = get_selected_background_info(bg_type, simple, studio, nature, indoor, abstract)
@@ -550,17 +325,30 @@ API 키 설정 방법:
 2. 코드 내 직접 입력: GEMINI_API_KEY = "your-api-key"
 키 발급: https://makersuite.google.com/
 """
+                    return prompt, image, "API 키를 설정해야 합니다."
                 
-                return prompt
+                # 프롬프트 요약 HTML 생성
+                preview = f"""
+                <div style="padding:10px; border:1px solid #ddd; border-radius:8px; margin-top:10px;">
+                    <h3>프롬프트 요약</h3>
+                    <p><strong>총 길이:</strong> {len(prompt)} 글자</p>
+                    <p><strong>배경:</strong> {background_info['category']} &gt; {background_info['name']}</p>
+                    <p><strong>주요 요소:</strong> {", ".join([kw for kw in ["commercial photography", "product", "square format", "centered", "detailed"] if kw.lower() in prompt.lower()])}</p>
+                    <p><strong>미드저니 파라미터:</strong> {" ".join([param for param in ["--ar 1:1", "--s 750", "--q 2"] if param in prompt])}</p>
+                </div>
+                """
+                
+                return prompt, image, preview
             except Exception as e:
                 error_msg = f"프롬프트 생성 중 오류 발생: {str(e)}"
                 gr.Error(error_msg)
-                return error_msg
+                return error_msg, image, "오류가 발생했습니다."
         
         # 프롬프트 생성 버튼 클릭 이벤트
-        generate_prompt_btn.click(
-            fn=generate_prompt,
+        submit_btn.click(
+            fn=generate_output,
             inputs=[
+                image_input,
                 background_type,
                 simple_dropdown,
                 studio_dropdown,
@@ -570,25 +358,15 @@ API 키 설정 방법:
                 product_name,
                 additional_info
             ],
-            outputs=prompt_output
-        )
-        
-        # 단일 이미지 생성 버튼 이벤트 연결
-        submit_single_btn.click(
-            fn=lambda image, prompt: process_images_with_prompt(image, None, None, prompt, 0),
-            inputs=[image_input, prompt_output],
-            outputs=[output_text],
-        )
-        
-        # 4장 이미지 생성 버튼 이벤트 연결
-        submit_btn.click(
-            fn=lambda image, prompt: generate_multiple_images(image, None, None, prompt),
-            inputs=[image_input, prompt_output],
-            outputs=[output_text],
+            outputs=[
+                prompt_output,
+                image_preview,
+                preview_html
+            ]
         )
     
     return demo
 
 if __name__ == "__main__":
     app = create_app()
-    app.launch()
+    app.launch()