Update app.py

app.py

@@ -3,10 +3,6 @@ import re
 import time
 from os import path
 import tempfile
-import uuid
-import base64
-import mimetypes
-import json
 import io
 import random
 import string
@@ -18,7 +14,6 @@ from transformers import pipeline
 from safetensors.torch import load_file
 from huggingface_hub import hf_hub_download
 
-# Diffusers
 import gradio as gr
 from diffusers import FluxPipeline
 
@@ -49,6 +44,7 @@ def maybe_translate_to_english(text: str) -> str:
     If the prompt contains any Korean characters, translate to English.
     Otherwise, return as-is.
     """
+    import re
     if re.search("[가-힣]", text):
         translated = translator(text)[0]["translation_text"]
         print(f"[TRANSLATE] Detected Korean -> '{text}' -> '{translated}'")
@@ -78,6 +74,7 @@ pipe = FluxPipeline.from_pretrained(
     torch_dtype=torch.bfloat16
 )
 
+# 예시용 LoRA 다운로드 & 합치기
 lora_path = hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors")
 pipe.load_lora_weights(lora_path)
 pipe.fuse_lora(lora_scale=0.125)
@@ -93,14 +90,13 @@ def save_binary_file(file_name, data):
 
 def generate_by_google_genai(text, file_name, model="gemini-2.0-flash-exp"):
     """
-    Internally modifies text within an image, returning a new image path.
-    (Screen instructions do not mention 'Google'.)
+    - 추가 지시사항(AIP)을 전달해 이미지 기반 편집을 수행.
+    - 응답이 '이미지'면 저장, '텍스트'면 누적하여 반환.
     """
+    # 기존 API 키 로직 유지 (환경 변수 GAPI_TOKEN 사용)
     api_key = os.getenv("GAPI_TOKEN", None)
     if not api_key:
-        raise ValueError(
-            "GAPI_TOKEN is missing. Please set an 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)]
@@ -130,6 +126,7 @@ def generate_by_google_genai(text, file_name, model="gemini-2.0-flash-exp"):
     text_response = ""
     image_path = None
 
+    # 임시 파일에 이미지 저장 가능하도록 준비
     with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
         temp_path = tmp.name
         for chunk in client.models.generate_content_stream(
@@ -137,22 +134,24 @@ def generate_by_google_genai(text, file_name, model="gemini-2.0-flash-exp"):
             contents=contents,
             config=generate_content_config,
         ):
-            if not chunk.candidates or not chunk.candidates[0].content:
+            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]
+            # 만약 inline_data(이미지 데이터)가 있다면 -> 실제 이미지 편집 결과
             if candidate.inline_data:
                 save_binary_file(temp_path, candidate.inline_data.data)
-                print(f"[DEBUG] Returned new image -> {temp_path}")
+                print(f"File of mime type {candidate.inline_data.mime_type} saved to: {temp_path}")
                 image_path = temp_path
+                # 이미지 한 장만 확보하면 중단
                 break
             else:
+                # inline_data가 없으면 텍스트 데이터이므로 누적
                 text_response += chunk.text + "\n"
     
     del files
     return image_path, text_response
 
-
 #######################################
 # 3. Diffusion Utility
 #######################################
@@ -167,8 +166,7 @@ def generate_random_letters(length: int) -> str:
 def is_all_english(text: str) -> bool:
     """
     Check if text consists only of English letters (a-z, A-Z), digits, spaces,
-    and a few basic punctuation characters. If so, return True.
-    Otherwise, False (includes Korean or other characters).
+    and basic punctuation. If so, return True; otherwise False.
     """
     import re
     return bool(re.match(r'^[a-zA-Z0-9\s\.,!\?\']*$', text))
@@ -176,8 +174,7 @@ def is_all_english(text: str) -> bool:
 def maybe_use_random_or_original(final_text: str) -> str:
     """
     If final_text is strictly English/allowed chars, use it as-is.
-    If it contains other chars (like Korean, etc.), 
-    replace with random letters of the same length.
+    Else replace with random letters of the same length.
     """
     if not final_text:
         return ""
@@ -188,8 +185,7 @@ def maybe_use_random_or_original(final_text: str) -> str:
 
 def fill_prompt_with_random_texts(prompt: str, r1: str, r2: str, r3: str) -> str:
     """
-    Replace <text1>, <text2>, <text3> with r1, r2, r3 respectively.
-    <text1> is required; if missing, we append something.
+    Replace <text1>, <text2>, <text3> placeholders with r1, r2, r3.
     """
     if "<text1>" in prompt:
         prompt = prompt.replace("<text1>", r1)
@@ -200,7 +196,6 @@ def fill_prompt_with_random_texts(prompt: str, r1: str, r2: str, r3: str) -> str
         prompt = prompt.replace("<text2>", r2)
     if "<text3>" in prompt:
         prompt = prompt.replace("<text3>", r3)
-    
     return prompt
 
 def generate_initial_image(prompt, height, width, steps, scale, seed):
@@ -219,30 +214,13 @@ def generate_initial_image(prompt, height, width, steps, scale, seed):
         ).images[0]
     return result
 
-
 #######################################
 # 4. Creating 2 Final Images
 #######################################
 
-def build_multi_change_instruction(r1, f1, r2, f2, r3, f3):
-    """
-    Summarize instructions to replace (r1->f1), (r2->f2), (r3->f3).
-    """
-    instructions = []
-    if r1 and f1:
-        instructions.append(f"Change any text reading '{r1}' in this image to '{f1}'.")
-    if r2 and f2:
-        instructions.append(f"Change any text reading '{r2}' in this image to '{f2}'.")
-    if r3 and f3:
-        instructions.append(f"Change any text reading '{r3}' in this image to '{f3}'.")
-    if instructions:
-        return " ".join(instructions)
-    return "No text changes needed."
-
 def change_text_in_image_two_times(original_image, instruction):
     """
-    Call the text modification function twice, 
-    returning 2 final variations.
+    Call the text-modification API twice, returning 2 final variations.
     """
     results = []
     for version_tag in ["(A)", "(B)"]:
@@ -262,12 +240,13 @@ def change_text_in_image_two_times(original_image, instruction):
                 new_img = Image.open(io.BytesIO(image_data))
                 results.append(new_img)
             else:
+                # 만약 이미지 응답이 없고, 텍스트만 온 경우
+                print("[WARNING] No image returned. text_response=", text_response)
                 results.append(original_image)
         except Exception as e:
             raise gr.Error(f"Error: {e}")
     return results
 
-
 #######################################
 # 5. Main Process (Generation from Prompt)
 #######################################
@@ -284,51 +263,61 @@ def run_process(
     seed
 ):
     """
-    1) If prompt has Korean, translate to English
-    2) For each <textX>, if it's purely English, use as-is,
-       else generate random letters of the same length.
-    3) Generate initial image with these placeholders
-    4) Then produce 2 final images by replacing placeholders with real texts
+    1) Translate prompt if Korean -> English
+    2) For each text, if not English -> random
+    3) Generate initial image
+    4) Replace placeholders with real text via API (2 variations)
     """
+    # 1) Translate prompt if needed
     prompt_en = maybe_translate_to_english(prompt)
 
-    # Decide random vs original for each text
+    # 2) Decide placeholders
     r1 = maybe_use_random_or_original(final_text1)
     r2 = maybe_use_random_or_original(final_text2)
     r3 = maybe_use_random_or_original(final_text3)
-
     print(f"[DEBUG] Using placeholders: r1='{r1}', r2='{r2}', r3='{r3}'")
 
-    # Fill prompt
+    # 3) Fill placeholders in prompt
     final_prompt = fill_prompt_with_random_texts(prompt_en, r1, r2, r3)
     print(f"[DEBUG] final_prompt = {final_prompt}")
 
-    # Generate initial "random/original" image
+    # 4) Generate initial "random/original" image
     _random_image = generate_initial_image(final_prompt, height, width, steps, scale, seed)
 
-    # Build final instructions & call twice -> 2 final images
-    instruction = build_multi_change_instruction(r1, final_text1, r2, final_text2, r3, final_text3)
+    # Build final instructions (replace placeholders -> real text)
+    instructions = []
+    if r1 and final_text1:
+        instructions.append(f"Change any text reading '{r1}' in this image to '{final_text1}'.")
+    if r2 and final_text2:
+        instructions.append(f"Change any text reading '{r2}' in this image to '{final_text2}'.")
+    if r3 and final_text3:
+        instructions.append(f"Change any text reading '{r3}' in this image to '{final_text3}'.")
+    instruction = " ".join(instructions) if instructions else "No text changes needed."
+
+    # Call 2 variations
     final_imgs = change_text_in_image_two_times(_random_image, instruction)
-    # Return only the 2 final images
     return [final_imgs[0], final_imgs[1]]
 
 #######################################
 # 5-2. Process for Editing Uploaded Image
 #######################################
-
-def run_edit_process(input_image, final_text1, final_text2, final_text3):
+def run_edit_process(input_image, edit_prompt, final_text1):
     """
-    Process for editing an uploaded image.
-    1) For each provided text, if it's purely English, use as-is;
-       otherwise, replace with random letters.
-    2) Build instructions and call the text modification function twice.
+    1) If final_text1 is empty => skip text replacement
+    2) Otherwise, combine edit_prompt + text-change instructions
+    3) Call 2 times for final images
     """
     r1 = maybe_use_random_or_original(final_text1)
-    r2 = maybe_use_random_or_original(final_text2)
-    r3 = maybe_use_random_or_original(final_text3)
-    print(f"[DEBUG] Editing image with placeholders: r1='{r1}', r2='{r2}', r3='{r3}'")
-    instruction = build_multi_change_instruction(r1, final_text1, r2, final_text2, r3, final_text3)
-    print(f"[DEBUG] Editing instruction: {instruction}")
+    print(f"[DEBUG] Editing image with placeholder r1='{r1}'")
+
+    # *** 수정 핵심 ***
+    # final_text1이 비어 있으면 텍스트 치환을 생략, 
+    # 아니면 "Change any text reading 'r1' => final_text1" 명령 추가
+    if not final_text1.strip():
+        instruction = f"{edit_prompt}"
+    else:
+        instruction = f"{edit_prompt}\nChange any text reading '{r1}' in this image to '{final_text1}'."
+
     final_imgs = change_text_in_image_two_times(input_image, instruction)
     return [final_imgs[0], final_imgs[1]]
 
@@ -372,16 +361,14 @@ with gr.Blocks(title="Eevery Text Imaginator: FLUX") as demo:
           color: #d500f9;
         }
         </style>
-
         <h2 style="text-align:center; margin-bottom: 15px;">
             <strong>Eevery Text Imaginator: FLUX</strong>
         </h2>
         
         <p style="text-align:center;">
             This tool generates <b>two final images</b> from a prompt
-            containing placeholders <code>&lt;text1&gt;</code>, <code>&lt;text2&gt;</code>, <code>&lt;text3&gt;</code>.
-            <br>If your chosen text is purely English, it will appear directly;
-            otherwise it becomes <i>random letters</i> in the initial phase.
+            or an uploaded image, optionally containing placeholders 
+            <code>&lt;text1&gt;</code>, <code>&lt;text2&gt;</code>, <code>&lt;text3&gt;</code>.
         </p>
         
         <hr style="margin: 15px 0;">
@@ -389,6 +376,9 @@ with gr.Blocks(title="Eevery Text Imaginator: FLUX") as demo:
     )
     
     with gr.Tabs():
+        ###############################################
+        # Tab 1) Generate from Prompt
+        ###############################################
         with gr.TabItem("Generate from Prompt"):
             with gr.Row():
                 with gr.Column():
@@ -411,12 +401,41 @@ with gr.Blocks(title="Eevery Text Imaginator: FLUX") as demo:
                             placeholder="(Leave blank if not used)"
                         )
                     with gr.Accordion("Advanced Settings (optional)", open=False):
-                        height = gr.Slider(label="Height", minimum=256, maximum=1152, step=64, value=512)
-                        width = gr.Slider(label="Width", minimum=256, maximum=1152, step=64, value=512)
-                        steps = gr.Slider(label="Inference Steps", minimum=6, maximum=25, step=1, value=8)
-                        scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.5, value=3.5)
-                        seed = gr.Number(label="Seed", value=1234, precision=0)
+                        height = gr.Slider(
+                            label="Height", 
+                            minimum=256, 
+                            maximum=1152, 
+                            step=64, 
+                            value=512
+                        )
+                        width = gr.Slider(
+                            label="Width", 
+                            minimum=256, 
+                            maximum=1152, 
+                            step=64, 
+                            value=512
+                        )
+                        steps = gr.Slider(
+                            label="Inference Steps", 
+                            minimum=6, 
+                            maximum=25, 
+                            step=1, 
+                            value=8
+                        )
+                        scale = gr.Slider(
+                            label="Guidance Scale", 
+                            minimum=0.0, 
+                            maximum=10.0, 
+                            step=0.5, 
+                            value=3.5
+                        )
+                        seed = gr.Number(
+                            label="Seed", 
+                            value=1234, 
+                            precision=0
+                        )
                     run_btn = gr.Button("Generate 2 Final Images", variant="primary")
+
                     gr.Examples(
                         examples=[
                             [
@@ -428,11 +447,11 @@ with gr.Blocks(title="Eevery Text Imaginator: FLUX") as demo:
                                 "환영합니다.", "", ""
                             ],
                             [
-                                "A classical poster reading <text1> in bold,  as a subtitle",
+                                "A classical poster reading <text1> in bold, as a subtitle",
                                 "错觉", "", ""
                             ],
                             [
-                                "In a cartoon style, a speech bubble with <text1> and another text ",
+                                "In a cartoon style, a speech bubble with <text1> and another text",
                                 "안녕", "", ""
                             ],
                             [
@@ -448,8 +467,16 @@ with gr.Blocks(title="Eevery Text Imaginator: FLUX") as demo:
                         label="Example Prompts"
                     )
                 with gr.Column():
-                    final_image_output1 = gr.Image(label="Final Image #1", type="pil")
-                    final_image_output2 = gr.Image(label="Final Image #2", type="pil")
+                    final_image_output1 = gr.Image(
+                        label="Final Image #1", 
+                        type="pil"
+                    )
+                    final_image_output2 = gr.Image(
+                        label="Final Image #2", 
+                        type="pil"
+                    )
+
+            # 버튼 클릭 시 처리
             run_btn.click(
                 fn=run_process,
                 inputs=[
@@ -466,33 +493,41 @@ with gr.Blocks(title="Eevery Text Imaginator: FLUX") as demo:
                 outputs=[final_image_output1, final_image_output2]
             )
         
+        ###############################################
+        # Tab 2) Edit Uploaded Image
+        ###############################################
         with gr.TabItem("Edit Uploaded Image"):
             with gr.Row():
                 with gr.Column():
+                    # Gradio 구버전 호환을 위해 source="upload"는 제거
                     uploaded_image = gr.Image(
-                        label="Upload Image for Editing", 
+                        label="Upload Image for Editing",
                         type="pil"
                     )
+                    edit_prompt = gr.Textbox(
+                        label="Additional Instruction Prompt",
+                        placeholder="(예: Make the background black, add sparkles, etc.)"
+                    )
                     final_text1_edit = gr.Textbox(
-                        label="New Text #1 (Required)",
+                        label="Replace Text",
                         placeholder="Example: HELLO or 안녕하세요"
                     )
-                    final_text2_edit = gr.Textbox(
-                        label="New Text #2 (Optional)",
-                        placeholder="Example: WORLD or 반갑습니다"
-                    )
-                    final_text3_edit = gr.Textbox(
-                        label="New Text #3 (Optional)",
-                        placeholder="(Leave blank if not used)"
-                    )
                     run_edit_btn = gr.Button("Edit Image", variant="primary")
                 with gr.Column():
-                    edited_image_output1 = gr.Image(label="Edited Image #1", type="pil")
-                    edited_image_output2 = gr.Image(label="Edited Image #2", type="pil")
+                    edited_image_output1 = gr.Image(
+                        label="Edited Image #1", 
+                        type="pil"
+                    )
+                    edited_image_output2 = gr.Image(
+                        label="Edited Image #2", 
+                        type="pil"
+                    )
+
+            # 업로드 이미지 편집 시 처리
             run_edit_btn.click(
                 fn=run_edit_process,
-                inputs=[uploaded_image, final_text1_edit, final_text2_edit, final_text3_edit],
+                inputs=[uploaded_image, edit_prompt, final_text1_edit],
                 outputs=[edited_image_output1, edited_image_output2]
             )
 
-demo.launch(max_threads=20)
+demo.launch(max_threads=20)