Update app.py

app.py

@@ -35,16 +35,15 @@ os.environ["TRANSFORMERS_CACHE"] = CACHE_PATH
 os.environ["HF_HUB_CACHE"] = CACHE_PATH
 os.environ["HF_HOME"] = CACHE_PATH
 
-# 타이머 클래스
 class timer:
     def __init__(self, method_name="timed process"):
         self.method = method_name
     def __enter__(self):
         self.start = time.time()
-        print(f"{self.method} starts")
+        print(f"[TIMER] {self.method} starts")
     def __exit__(self, exc_type, exc_val, exc_tb):
         end = time.time()
-        print(f"{self.method} took {str(round(end - self.start, 2))}s")
+        print(f"[TIMER] {self.method} took {round(end - self.start, 2)}s")
 
 #######################################
 # 1. FLUX 파이프라인 로드
@@ -61,30 +60,23 @@ pipe = FluxPipeline.from_pretrained(
 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)
-
 pipe.to(device="cuda", dtype=torch.bfloat16)
 
 #######################################
-# 2. Google GenAI 모델로 텍스트 변환 함수
+# 2. Google GenAI (Gemini) - 이미지 변환 함수
 #######################################
 
 def save_binary_file(file_name, data):
-    """Google GenAI에서 응답받은 이진 데이터를 이미지 파일로 저장"""
     with open(file_name, "wb") as f:
         f.write(data)
 
 def generate_by_google_genai(text, file_name, model="gemini-2.0-flash-exp"):
-    """
-    Google GenAI(gemini) 모델을 통해 이미지/텍스트를 생성하거나 변환.
-    - text: 변경할 텍스트나 명령어 등 프롬프트
-    - file_name: 원본 이미지(예: .png) 경로
-    - model: 사용할 gemini 모델 이름
-    """
+    """Gemini 모델을 통해 이미지 내부 텍스트를 변경."""
     api_key = os.getenv("GAPI_TOKEN", None)
     if not api_key:
         raise ValueError(
             "GAPI_TOKEN 환경 변수가 설정되지 않았습니다. "
-            "Google GenAI API 사용을 위해서는 GAPI_TOKEN이 필요합니다."
+            "Google GenAI API를 사용하기 위해서는 GAPI_TOKEN이 필요합니다."
         )
 
     client = genai.Client(api_key=api_key)
@@ -122,13 +114,13 @@ 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 or not chunk.candidates[0].content.parts:
+            if not chunk.candidates or not chunk.candidates[0].content:
                 continue
-            candidate = chunk.candidates[0].content.parts[0]
 
+            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}")
+                print(f"[DEBUG] Gemini returned image -> {temp_path}")
                 image_path = temp_path
                 break
             else:
@@ -138,24 +130,43 @@ def generate_by_google_genai(text, file_name, model="gemini-2.0-flash-exp"):
     return image_path, text_response
 
 #######################################
-# 3. Diffusion + GoogleGenAI를 연결
+# 3. Diffusion (Flux)용 함수
 #######################################
 
-def generate_initial_image(prompt, text, height, width, steps, scale, seed):
+def generate_random_letters(length: int) -> str:
+    """length 길이만큼 대소문자 알파벳을 무작위로 생성."""
+    letters = string.ascii_lowercase + string.ascii_uppercase
+    return "".join(random.choice(letters) for _ in range(length))
+
+def fill_prompt_with_random_texts(prompt: str, r1: str, r2: str, r3: str) -> str:
     """
-    1) FLUX 파이프라인을 사용해 'text'가 들어간 이미지를 생성
-       - prompt 내 <text>가 있으면 치환, 없으면 자동 추가
+    프롬프트 내 <text1>, <text2>, <text3>를
+    각각 r1, r2, r3로 치환.
+    - <text1>은 필수 (없으면 자동으로 뒤에 붙임).
+    - <text2>, <text3>는 있으면 치환, 없으면 무시.
     """
-    if "<text>" in prompt:
-        combined_prompt = prompt.replace("<text>", text)
+    # 1) <text1>은 필수
+    if "<text1>" in prompt:
+        prompt = prompt.replace("<text1>", r1)
     else:
-        combined_prompt = f"{prompt} with clear readable text that says '{text}'"
-
-    print("[DEBUG] combined_prompt:", combined_prompt)
+        # 자동 덧붙임
+        prompt = f"{prompt} with clear readable text that says '{r1}'"
+
+    # 2) <text2>, <text3>는 선택
+    if "<text2>" in prompt:
+        prompt = prompt.replace("<text2>", r2)
+    if "<text3>" in prompt:
+        prompt = prompt.replace("<text3>", r3)
+    
+    return prompt
 
-    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("GenerateInitialImage"):
+def generate_initial_image(prompt, random1, random2, random3, height, width, steps, scale, seed):
+    """
+    Flux 파이프라인을 이용해 (r1, r2, r3)가 들어간 이미지를 생성.
+    """
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("Flux Generation"):
         result = pipe(
-            prompt=[combined_prompt],
+            prompt=[prompt],
             generator=torch.Generator().manual_seed(int(seed)),
             num_inference_steps=int(steps),
             guidance_scale=float(scale),
@@ -163,108 +174,163 @@ def generate_initial_image(prompt, text, height, width, steps, scale, seed):
             width=int(width),
             max_sequence_length=256
         ).images[0]
-
     return result
 
-def change_text_in_image(original_image, new_text):
+
+def change_multi_text_in_image(original_image, random1, final1, random2, final2, random3, final3):
     """
-    2) Gemini 모델을 통해,
-       original_image 내 텍스트를 `new_text`로 변경한 이미지 반환
+    Gemini를 통해, 이미지 안의 r1->final1, r2->final2, r3->final3 식으로 텍스트 교체.
+    - r2, final2 (또는 r3, final3)가 빈 문자열이면 해당 교체는 건너뜀.
     """
+    # 교체 지시문 만들기
+    instructions = []
+    if random1 and final1:
+        instructions.append(f"Change any text reading '{random1}' in this image to '{final1}'.")
+    if random2 and final2:
+        instructions.append(f"Change any text reading '{random2}' in this image to '{final2}'.")
+    if random3 and final3:
+        instructions.append(f"Change any text reading '{random3}' in this image to '{final3}'.")
+
+    # 만약 교체 지시문이 없다면 그냥 return original_image
+    if not instructions:
+        print("[WARN] No text changes requested!")
+        return original_image
+
+    full_instruction = " ".join(instructions)
     try:
+        # 임시 파일에 original_image 저장
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
             original_path = tmp.name
             original_image.save(original_path)
 
         image_path, text_response = generate_by_google_genai(
-            text=f"Change the text in this image to: '{new_text}'",
+            text=full_instruction,
             file_name=original_path
         )
-
         if image_path:
             with open(image_path, "rb") as f:
                 image_data = f.read()
-            modified_img = Image.open(io.BytesIO(image_data))
-            return modified_img
+            new_img = Image.open(io.BytesIO(image_data))
+            return new_img
         else:
-            # 이미지가 없이 텍스트만 응답된 경우
-            return None
-
+            # 이미지 없이 텍스트만 온 경우
+            print("[WARN] Gemini returned only text:", text_response)
+            return original_image
     except Exception as e:
         raise gr.Error(f"Error: {e}")
 
 #######################################
-# 4. 임의 알파벳 생성
-#######################################
-
-def generate_random_letters(length: int) -> str:
-    """
-    length 길이의 임의 알파벳(대소문자) 문자열 생성
-    """
-    letters = string.ascii_lowercase + string.ascii_uppercase
-    return "".join(random.choice(letters) for _ in range(length))
-
-#######################################
-# 5. 최종 함수: 버튼 한 번으로
-#    (1) 무작위 알파벳으로 1차 이미지 생성
-#    (2) 진짜 "새로 바꿀 텍스트"로 2차 이미지 생성
+# 4. 메인 프로세스 함수
 #######################################
 
-def run_full_process(prompt, final_text, height, width, steps, scale, seed):
+def run_process(
+    prompt,
+    final_text1,
+    final_text2,
+    final_text3,
+    height,
+    width,
+    steps,
+    scale,
+    seed
+):
     """
-    - final_text의 길이에 맞춰 랜덤 알파벳을 생성 -> 1차 이미지
-    - 그 1차 이미지를 바탕으로, final_text로 교체 -> 2차 최종 이미지
+    1) final_text1(필수), final_text2, final_text3(옵션) 각각 길이에 맞춰 무작위 알파벳 만들기.
+    2) prompt 내 <text1>, <text2>, <text3> 치환 -> Flux로 1차(랜덤) 이미지.
+    3) Gemini 호출 -> r1->final_text1, r2->final_text2, r3->final_text3 교체 -> 최종 이미지.
     """
-    # (A) 새로 바꿀 텍스트(final_text) 글자수만큼 임의 알파벳 생성
-    random_len = len(final_text)
-    random_text = generate_random_letters(random_len)
-    print(f"[STEP] final_text='{final_text}' => random_text='{random_text}'")
-
-    # (B) 1차 이미지: 무작위 알파벳으로 생성
-    random_image = generate_initial_image(prompt, random_text, height, width, steps, scale, seed)
+    # (A) 무작위 알파벳
+    r1 = generate_random_letters(len(final_text1)) if final_text1 else ""
+    r2 = generate_random_letters(len(final_text2)) if final_text2 else ""
+    r3 = generate_random_letters(len(final_text3)) if final_text3 else ""
+
+    # (B) 프롬프트 치환
+    final_prompt = fill_prompt_with_random_texts(prompt, r1, r2, r3)
+    print(f"[DEBUG] final_prompt = {final_prompt}")
+
+    # (C) 1차 이미지 (랜덤 텍스트)
+    random_image = generate_initial_image(final_prompt, r1, r2, r3, height, width, steps, scale, seed)
+
+    # (D) 2차 이미지 (실제 텍스트)
+    final_image = change_multi_text_in_image(
+        random_image,
+        r1, final_text1,
+        r2, final_text2,
+        r3, final_text3
+    )
 
-    # (C) 2차 이미지: 실제 final_text로 교체
-    final_image = change_text_in_image(random_image, final_text)
-    
     return [random_image, final_image]
 
 #######################################
-# 6. Gradio UI
+# 5. Gradio UI
 #######################################
 
-with gr.Blocks(title="Flux + Google GenAI (Random & Then Real Text)") as demo:
+with gr.Blocks(title="Flux + Google GenAI (Up to 3 Text placeholders)") as demo:
     gr.Markdown(
         """
-        # Flux + Google GenAI: 두 단계에 걸친 텍스트 교체
-
-        **사용 흐름**  
-        1) Prompt에 장면이나 스타일을 작성 (필요하면 `<text>` 구분자 사용)  
-        2) "새로 바꿀 텍스트" 에 최종 원하는 문자열을 입력 (예: "안녕하세요")  
-        3) "Generate Images" 버튼을 누르면,  
-           - 먼저 "새로 바꿀 텍스트" 길이에 맞는 **무작위 알파벳**을 넣어 이미지 생성 (1차 이미지)  
-           - 이어서 **진짜** "새로 바꿀 텍스트"로 다시 교체(2차 최종 이미지)  
-        4) 결과로 두 장의 이미지를 확인할 수 있습니다.  
-
-        **주요 포인트**  
-        - "이미지 안에 들어갈 텍스트"는 **UI에 노출되지 않으며**(사용자 입력 불가), 오직 내부에서 자동 설정됩니다.
-        - 1차 이미지는 완전히 임의의 알파벳 텍스트를 포함합니다.
-        - 2차 이미지는 최종적으로 사용자가 입력한 "새로 바꿀 텍스트"를 포함합니다.
+        # Flux + Google GenAI: 최대 3개의 `<text>` 교체
+
+        ## 사용 방법
+        1. 아래 Prompt에 `<text1>`, `<text2>`, `<text3>`를 최대 3개까지 배치 가능.
+           - 예) "A poster with <text1> in large letters, also <text2> in the corner"
+           - **<text1>은 필수**(없으면 자동으로 문구가 뒤에 붙음)
+           - <text2>, <text3>는 넣어도 되고, 안 넣어도 됨.
+        2. "New Text #1" (필수), "New Text #2", "New Text #3"를 입력.
+           - #2, #3는 비워 두면 해당 자리 교체 없음.
+        3. "Generate Images" 버튼 →  
+           (1) `<text1>`, `<text2>`, `<text3>` 자리에 (또는 자동으로) **무작위 알파벳** 넣은 1차 이미지 생성  
+           (2) 이어 Gemini 모델을 통해 무작위 알파벳 → 실제 "New Text #1/2/3" 변경한 2차 이미지  
+           - **두 이미지**(랜덤 텍스트 → 최종 텍스트)가 순서대로 출력됩니다.
+
+        ---
         """
     )
 
+    # 예시 5개
+    examples = [
+        [
+            "A futuristic billboard shows <text1> and a small sign <text2> on the left side. <text3> is a hidden watermark.",
+            "HELLO", "WELCOME", "2025"
+        ],
+        [
+            "A fantasy poster with <text1> and <text2> in stylized letters, plus a tiny note <text3> at the bottom.",
+            "Dragons", "MagicRealm", "Beware!"
+        ],
+        [
+            "A neon sign reading <text1>, with a secondary text <text2> below. <text3> might appear in the corner.",
+            "OPEN", "24HOUR", "NoSmoking"
+        ],
+        [
+            "A big invitation card with main text <text1>, subtitle <text2>, signature <text3> in cursive.",
+            "Birthday Party", "Today Only", "From Your Friend"
+        ],
+        [
+            "A large graffiti wall with <text1> in bold letters, plus <text2> and <text3> near the edges.",
+            "FREEDOM", "HOPE", "LOVE"
+        ]
+    ]
+
     with gr.Row():
         with gr.Column():
             prompt_input = gr.Textbox(
                 lines=3,
-                label="Prompt (use `<text>` if you want)",
-                placeholder="e.g. A white cat with a speech bubble <text>"
+                label="Prompt (use `<text1>`, `<text2>`, `<text3>` as needed)",
+                placeholder="Ex) A poster with <text1>, plus a line <text2>, etc."
+            )
+            final_text1 = gr.Textbox(
+                label="New Text #1 (Required)",
+                placeholder="Ex) HELLO"
             )
-            final_text_input = gr.Textbox(
-                lines=1,
-                label="새로 바꿀 텍스트",
-                placeholder="예) 안녕하세요"
+            final_text2 = gr.Textbox(
+                label="New Text #2 (Optional)",
+                placeholder="Ex) WELCOME"
             )
-            with gr.Accordion("고급 설정 (확장)", open=False):
+            final_text3 = gr.Textbox(
+                label="New Text #3 (Optional)",
+                placeholder="Ex) 2025 or anything"
+            )
+
+            with gr.Accordion("Advanced Settings", 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)
@@ -273,14 +339,30 @@ with gr.Blocks(title="Flux + Google GenAI (Random & Then Real Text)") as demo:
 
             run_btn = gr.Button("Generate Images", variant="primary")
 
+            gr.Examples(
+                examples=examples,
+                inputs=[prompt_input, final_text1, final_text2, final_text3],
+                label="Click to load example"
+            )
+
         with gr.Column():
             random_image_output = gr.Image(label="1) Random Text Image", type="pil")
             final_image_output = gr.Image(label="2) Final Text Image", type="pil")
 
-    # 버튼 액션: 위 단계들을 모두 실행 -> 결과 2장 출력
+    # 버튼 액션
     run_btn.click(
-        fn=run_full_process,
-        inputs=[prompt_input, final_text_input, height, width, steps, scale, seed],
+        fn=run_process,
+        inputs=[
+            prompt_input,
+            final_text1,
+            final_text2,
+            final_text3,
+            height,
+            width,
+            steps,
+            scale,
+            seed
+        ],
         outputs=[random_image_output, final_image_output]
     )