Update app.py

app.py

@@ -9,109 +9,138 @@ import torch
 from diffusers import DiffusionPipeline
 from PIL import Image
 
-# -----------------------------
-# Gemini API & Text Rendering 관련 추가 모듈
-# -----------------------------
 import re
 import tempfile
 import io
 import logging
-import base64
-import string
-import requests
-from google import genai
-from google.genai import types
-
-import numpy as np
-
 
+# -----------------------------
+# Google Gemini API 관련
+# -----------------------------
+import google.generativeai as genai
+import google.generativeai.types as genai_types
 
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
+###############################################################################
+# 1. 텍스트(한글 → 영어) 변환 보조 함수
+###############################################################################
+
 def maybe_translate_to_english(text: str) -> str:
     """
-    텍스트에 한글이 포함되어 있으면 간단한 규칙에 따라 영어로 변환.
+    텍스트에 한국어가 있으면 간단한 치환 규칙에 따라 영어로 변환.
     """
-    if not text or not re.search("[가-힣]", text):
-        return text
-    try:
-        translations = {
-            "안녕하세요": "Hello",
-            "환영합니다": "Welcome",
-            "아름다운 당신": "Beautiful You",
-            "안녕": "Hello",
-            "고양이": "Cat",
-            "배너": "Banner",
-            "썬글라스": "Sunglasses",
-            "착용한": "wearing",
-            "흰색": "white"
-        }
-        for kr, en in translations.items():
-            if kr in text:
-                text = text.replace(kr, en)
-        print(f"[TRANSLATE] Translated Korean text: '{text}'")
-        return text
-    except Exception as e:
-        print(f"[WARNING] Translation failed: {e}")
-        return text
+    translations = {
+        "안녕하세요": "Hello",
+        "환영합니다": "Welcome",
+        "안녕": "Hello",
+        "배너": "Banner",
+        # 필요에 따라 추가
+    }
+    for kr, en in translations.items():
+        if kr in text:
+            text = text.replace(kr, en)
+    return text
+
+###############################################################################
+# 2. Gemini API 호출을 위한 준비
+###############################################################################
 
 def save_binary_file(file_name, data):
+    """ 이진 파일을 저장하는 헬퍼 함수 """
     with open(file_name, "wb") as f:
         f.write(data)
 
 def generate_by_google_genai(text, file_name, model="gemini-2.0-flash-exp"):
     """
-    Gemini API를 호출하여 텍스트 기반 이미지 편집/생성을 수행.
+    Google Gemini API를 호출해 텍스트 기반 이미지 편집/생성을 수행.
+    file_name: 원본 이미지를 임시 업로드하여 API로 전달
+    text: 적용할 텍스트 지시사항
     """
-    api_key = os.getenv("GAPI_TOKEN", None)
+    api_key = os.getenv("GAPI_TOKEN")
     if not api_key:
         raise ValueError("GAPI_TOKEN is missing. Please set an API key.")
-    client = genai.Client(api_key=api_key)
-    files = [client.files.upload(file=file_name)]
+
+    # Gemini API 인증 설정
+    genai.configure(api_key=api_key)
+
+    # 이미지 파일 업로드
+    uploaded_file = genai.upload_file(path=file_name)
+
+    # API에 전달할 content 구성
     contents = [
-        types.Content(
+        genai_types.Content(
             role="user",
             parts=[
-                types.Part.from_uri(
-                    file_uri=files[0].uri,
-                    mime_type=files[0].mime_type,
+                # 먼저 업로드된 파일 URI를 포함
+                genai_types.Part.from_uri(
+                    file_uri=uploaded_file.uri,
+                    mime_type=uploaded_file.mime_type,
                 ),
-                types.Part.from_text(text=text),
+                # 이어서 text 지시사항을 포함
+                genai_types.Part.from_text(text=text),
             ],
         ),
     ]
-    generate_content_config = types.GenerateContentConfig(
+
+    # 생성(편집) 설정
+    generation_config = genai_types.GenerationConfig(
         temperature=1,
         top_p=0.95,
         top_k=40,
-        max_output_tokens=8192,
-        response_modalities=["image", "text"],
+        max_output_tokens=8192,      # 출력 토큰 제한
         response_mime_type="text/plain",
     )
-    text_response = ""
-    image_path = None
+
+    text_response = ""   # API가 반환한 텍스트 누적
+    image_path = None    # API가 반환한 이미지 파일의 로컬 경로
+
+    # 임시 파일에 편집된 이미지 저장
     with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
         temp_path = tmp.name
-        for chunk in client.models.generate_content_stream(
+
+        # 스트리밍 형태로 응답을 받음
+        response = genai.generate_content(
             model=model,
             contents=contents,
-            config=generate_content_config,
-        ):
-            if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
-                continue
-            candidate = chunk.candidates[0].content.parts[0]
-            if candidate.inline_data:
-                save_binary_file(temp_path, candidate.inline_data.data)
-                print(f"File of mime type {candidate.inline_data.mime_type} saved to: {temp_path}")
-                image_path = temp_path
+            generation_config=generation_config,
+            stream=True
+        )
+
+        # 스트리밍된 chunk들에서 이미지와 텍스트를 추출
+        for chunk in response:
+            for candidate in chunk.candidates:
+                for part in candidate.content.parts:
+                    # 이미지인 경우
+                    if hasattr(part, 'inline_data') and part.inline_data:
+                        save_binary_file(temp_path, part.inline_data.data)
+                        image_path = temp_path
+                        break
+                    # 텍스트인 경우
+                    elif hasattr(part, 'text'):
+                        text_response += part.text + "\n"
+
+                if image_path:
+                    break
+            if image_path:
                 break
-            else:
-                text_response += chunk.text + "\n"
-    del files
+
+    # 업로드된 임시 파일 삭제
+    genai.delete_file(uploaded_file.name)
+
     return image_path, text_response
 
+###############################################################################
+# 3. 이미지에 텍스트를 삽입/수정하는 함수 (Gemini API 2회 호출)
+###############################################################################
+
 def change_text_in_image_two_times(original_image, instruction):
-    # 만약 이미지가 numpy.ndarray 타입이면 PIL Image로 변환
+    """
+    Gemini API를 두 번 호출하여 두 개의 버전을 생성한다.
+    """
+    import numpy as np
+
+    # 만약 이미지가 numpy.ndarray 타입이면 PIL로 변환
     if isinstance(original_image, np.ndarray):
         original_image = Image.fromarray(original_image)
     
@@ -123,38 +152,46 @@ def change_text_in_image_two_times(original_image, instruction):
                 original_path = tmp.name
                 if isinstance(original_image, Image.Image):
                     original_image.save(original_path, format="PNG")
-                    print(f"[DEBUG] Saved image to temporary file: {original_path}")
+                    logging.debug(f"[DEBUG] Saved image to temporary file: {original_path}")
                 else:
                     raise gr.Error(f"예상된 PIL Image가 아닌 {type(original_image)} 타입이 제공되었습니다.")
-            # 이후 Gemini API 호출 로직 유지
+            # Gemini API 호출
             image_path, text_response = generate_by_google_genai(
                 text=mod_instruction,
                 file_name=original_path
             )
             if image_path:
+                # 반환된 이미지 로드
                 try:
                     with open(image_path, "rb") as f:
                         image_data = f.read()
                     new_img = Image.open(io.BytesIO(image_data))
                     results.append(new_img)
                 except Exception as img_err:
-                    print(f"[ERROR] Failed to process Gemini image: {img_err}")
+                    logging.error(f"[ERROR] Failed to process Gemini image: {img_err}")
                     results.append(original_image)
             else:
-                print(f"[WARNING] 이미지가 반환되지 않았습니다. 텍스트 응답: {text_response}")
+                logging.warning(f"[WARNING] 이미지가 반환되지 않았습니다. 텍스트 응답: {text_response}")
                 results.append(original_image)
         except Exception as e:
             logging.exception(f"Text modification error: {e}")
             results.append(original_image)
     return results
 
+###############################################################################
+# 4. 텍스트 렌더링(문자 삽입)용 함수
+###############################################################################
 
 def gemini_text_rendering(image, rendering_text):
     """
-    주어진 이미지에 대해 Gemini API를 사용해 텍스트 렌더링을 적용.
+    주어진 image에 대해 Gemini API로 text_rendering을 적용
     """
     rendering_text_en = maybe_translate_to_english(rendering_text)
-    instruction = f"Render the following text on the image in a clear, visually appealing manner: {rendering_text_en}."
+    instruction = (
+        f"Render the following text on the image in a clear, visually appealing manner: "
+        f"{rendering_text_en}."
+    )
+    # 이미지에 텍스트 삽입(A/B 버전 2회 생성) → 여기서는 2회 중 첫 번째만 반환
     rendered_images = change_text_in_image_two_times(image, instruction)
     if rendered_images and len(rendered_images) > 0:
         return rendered_images[0]
@@ -162,32 +199,16 @@ def gemini_text_rendering(image, rendering_text):
 
 def apply_text_rendering(image, rendering_text):
     """
-    입력된 텍스트가 있으면 Gemini API를 통해 이미지에 텍스트 렌더링을 적용하고, 없으면 원본 이미지를 그대로 반환.
+    rendering_text가 존재하면 Gemini API로 텍스트 삽입을 적용.
+    없으면 원본 이미지를 그대로 반환.
     """
     if rendering_text and rendering_text.strip():
         return gemini_text_rendering(image, rendering_text)
     return image
 
-# -----------------------------
-# 기존 Diffusion Pipeline 관련 코드
-# -----------------------------
-import gradio_client.utils
-import types
-
-original_json_schema = gradio_client.utils._json_schema_to_python_type
-def patched_json_schema(schema, defs=None):
-    if isinstance(schema, bool):
-        return "bool"
-    try:
-        if "additionalProperties" in schema and isinstance(schema["additionalProperties"], bool):
-            schema["additionalProperties"] = {"type": "any"}
-    except (TypeError, KeyError):
-        pass
-    try:
-        return original_json_schema(schema, defs)
-    except Exception as e:
-        return "any"
-gradio_client.utils._json_schema_to_python_type = patched_json_schema
+###############################################################################
+# 5. Diffusion Pipeline 로드 및 기본 세팅
+###############################################################################
 
 SAVE_DIR = "saved_images"
 if not os.path.exists(SAVE_DIR):
@@ -198,24 +219,27 @@ repo_id = "black-forest-labs/FLUX.1-dev"
 adapter_id = "openfree/flux-chatgpt-ghibli-lora"
 
 def load_model_with_retry(max_retries=5):
+    """
+    로컬 또는 Hugging Face로부터 모델(FLUX.1-dev) + LoRA 어댑터(weights)를 불러온다.
+    """
     for attempt in range(max_retries):
         try:
-            print(f"Loading model attempt {attempt+1}/{max_retries}...")
+            logging.info(f"Loading model attempt {attempt+1}/{max_retries}...")
             pipeline = DiffusionPipeline.from_pretrained(
-                repo_id, 
+                repo_id,
                 torch_dtype=torch.bfloat16,
                 use_safetensors=True,
                 resume_download=True
             )
-            print("Model loaded successfully, loading LoRA weights...")
+            logging.info("Model loaded successfully, loading LoRA weights...")
             pipeline.load_lora_weights(adapter_id)
             pipeline = pipeline.to(device)
-            print("Pipeline ready!")
+            logging.info("Pipeline ready!")
             return pipeline
         except Exception as e:
             if attempt < max_retries - 1:
                 wait_time = 10 * (attempt + 1)
-                print(f"Error loading model: {e}. Retrying in {wait_time} seconds...")
+                logging.error(f"Error loading model: {e}. Retrying in {wait_time} seconds...")
                 import time
                 time.sleep(wait_time)
             else:
@@ -227,21 +251,31 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 def save_generated_image(image, prompt):
+    """
+    생성된 이미지를 저장하면서 메타 정보를 기록한다.
+    """
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     unique_id = str(uuid.uuid4())[:8]
     filename = f"{timestamp}_{unique_id}.png"
     filepath = os.path.join(SAVE_DIR, filename)
     image.save(filepath)
+
     metadata_file = os.path.join(SAVE_DIR, "metadata.txt")
     with open(metadata_file, "a", encoding="utf-8") as f:
         f.write(f"{filename}|{prompt}|{timestamp}\n")
     return filepath
 
 def load_generated_images():
+    """
+    저장된 이미지를 최신순으로 불러온다.
+    """
     if not os.path.exists(SAVE_DIR):
         return []
-    image_files = [os.path.join(SAVE_DIR, f) for f in os.listdir(SAVE_DIR) 
-                  if f.endswith(('.png', '.jpg', '.jpeg', '.webp'))]
+    image_files = [
+        os.path.join(SAVE_DIR, f)
+        for f in os.listdir(SAVE_DIR)
+        if f.endswith(('.png', '.jpg', '.jpeg', '.webp'))
+    ]
     image_files.sort(key=lambda x: os.path.getctime(x), reverse=True)
     return image_files
 
@@ -257,9 +291,13 @@ def inference(
     lora_scale: float,
     progress: gr.Progress = gr.Progress(track_tqdm=True),
 ):
+    """
+    Diffusion Pipeline을 사용해 이미지를 생성. (LoRA 스케일, Steps 등 설정 가능)
+    """
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device=device).manual_seed(seed)
+
     try:
         image = pipeline(
             prompt=prompt,
@@ -270,16 +308,19 @@ def inference(
             generator=generator,
             joint_attention_kwargs={"scale": lora_scale},
         ).images[0]
+
         filepath = save_generated_image(image, prompt)
         return image, seed, load_generated_images()
+
     except Exception as e:
-        print(f"Error during inference: {e}")
+        logging.error(f"Error during inference: {e}")
         error_img = Image.new('RGB', (width, height), color='red')
         return error_img, seed, load_generated_images()
 
-# -----------------------------
-# Gradio UI (입력 프롬프트 아래에 Text Rendering 입력란 추가)
-# -----------------------------
+###############################################################################
+# 6. Gradio UI
+###############################################################################
+
 examples = [
     "Ghibli style futuristic stormtrooper with glossy white armor and a sleek helmet, standing heroically on a lush alien planet, vibrant flowers blooming around, soft sunlight illuminating the scene, a gentle breeze rustling the leaves. The armor reflects the pink and purple hues of the alien sunset, creating an ethereal glow around the figure. [trigger]",
     "Ghibli style young mechanic girl in a floating workshop, surrounded by hovering tools and glowing mechanical parts, her blue overalls covered in oil stains, tinkering with a semi-transparent robot companion. Magical sparks fly as she works, while floating islands with waterfalls drift past her open workshop window. [trigger]",
@@ -506,7 +547,7 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
                         placeholder="Describe your Ghibli-style image here...",
                         lines=3
                     )
-                    # ★ 새롭게 추가된 Text Rendering 입력란
+                    # Text Rendering 입력란
                     text_rendering = gr.Textbox(
                         label="Text Rendering (Multilingual: English, Korean...)",
                         placeholder="Man saying '안녕' in 'speech bubble'",
@@ -568,15 +609,14 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
                 
                 with gr.Group(elem_classes="container"):
                     gr.Markdown("### ✨ Example Prompts")
-                    examples_html = '\n'.join([f'<div class="example-prompt">{example}</div>' for example in examples])
+                    examples_html = '\n'.join([f'<div class="example-prompt">{ex}</div>' for ex in examples])
                     example_container = gr.HTML(examples_html)
             
             with gr.Column(scale=4):
                 with gr.Group(elem_classes="container"):
-                    with gr.Group():
-                        generation_status = gr.HTML('<div class="status-complete">Ready to generate</div>')
-                        result = gr.Image(label="Generated Image", elem_id="result-image")
-                        seed_text = gr.Number(label="Used Seed", value=42)
+                    generation_status = gr.HTML('<div class="status-complete">Ready to generate</div>')
+                    result = gr.Image(label="Generated Image", elem_id="result-image")
+                    seed_text = gr.Number(label="Used Seed", value=42)
     
     with gr.Tabs(elem_classes="tabs") as tabs:
         with gr.TabItem("Gallery"):
@@ -592,6 +632,9 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
                     elem_classes="gallery-item"
                 )
 
+    ###########################################################################
+    # Gradio Helper Functions
+    ###########################################################################
     def refresh_gallery():
         return load_generated_images()
 
@@ -601,9 +644,12 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
     def before_generate():
         return '<div class="status-generating">Generating image...</div>'
     
-    def after_generate(image, seed, gallery):
-        return image, seed, gallery, '<div class="status-complete">Generation complete!</div>'
+    def after_generate(image, seed_num, gallery):
+        return image, seed_num, gallery, '<div class="status-complete">Generation complete!</div>'
 
+    ###########################################################################
+    # Gradio Event Wiring
+    ###########################################################################
     refresh_btn.click(
         fn=refresh_gallery,
         inputs=None,
@@ -616,7 +662,10 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
         outputs=[prompt, result, seed_text, generation_status]
     )
     
-    # 체인에 마지막에 텍스트 렌더링 적용 (text_rendering 입력값이 있으면)
+    # 1) 상태 표시
+    # 2) 이미지 생성
+    # 3) 상태 업데이트
+    # 4) 텍스트 렌더링(있다면)
     run_button.click(
         fn=before_generate,
         inputs=None,
@@ -641,9 +690,10 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
     ).then(
         fn=apply_text_rendering,
         inputs=[result, text_rendering],
-        outputs=result,
+        outputs=result
     )
     
+    # prompt submit 시에도 동일한 체인 실행
     prompt.submit(
         fn=before_generate,
         inputs=None,
@@ -668,9 +718,10 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
     ).then(
         fn=apply_text_rendering,
         inputs=[result, text_rendering],
-        outputs=result,
+        outputs=result
     )
     
+    # JS로 예시 prompt 클릭 시 자동 채우기
     gr.HTML("""
     <script>
     document.addEventListener('DOMContentLoaded', function() {
@@ -689,10 +740,13 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
     </script>
     """)
 
+###############################################################################
+# 7. 실행
+###############################################################################
 try:
     demo.queue(concurrency_count=1, max_size=20)
     demo.launch(debug=True, show_api=False)
 except Exception as e:
-    print(f"Error during launch: {e}")
-    print("Trying alternative launch configuration...")
+    logging.error(f"Error during launch: {e}")
+    logging.info("Trying alternative launch configuration...")
     demo.launch(debug=True, show_api=False, share=False)