Update app.py

app.py

@@ -35,7 +35,7 @@ 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
@@ -65,7 +65,7 @@ pipe.fuse_lora(lora_scale=0.125)
 pipe.to(device="cuda", dtype=torch.bfloat16)
 
 #######################################
-# 2. Google GenAI를 통한 이미지 내 텍스트 변환 함수
+# 2. Google GenAI 모델로 텍스트 변환 함수
 #######################################
 
 def save_binary_file(file_name, data):
@@ -80,21 +80,16 @@ def generate_by_google_genai(text, file_name, model="gemini-2.0-flash-exp"):
     - file_name: 원본 이미지(예: .png) 경로
     - model: 사용할 gemini 모델 이름
     """
-    # (1) 환경 변수에서 API 키 가져오기 (필수)
     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이 필요합니다."
         )
 
-    # (2) Google Client 초기화
     client = genai.Client(api_key=api_key)
-    
-    # (3) 이미지 업로드
     files = [client.files.upload(file=file_name)]
     
-    # (4) gemini에 전달할 Content 준비 (이미지 + 프롬프트)
     contents = [
         types.Content(
             role="user",
@@ -108,7 +103,6 @@ def generate_by_google_genai(text, file_name, model="gemini-2.0-flash-exp"):
         ),
     ]
 
-    # (5) 생성/변환 설정
     generate_content_config = types.GenerateContentConfig(
         temperature=1,
         top_p=0.95,
@@ -121,10 +115,8 @@ 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(
             model=model,
             contents=contents,
@@ -134,60 +126,34 @@ def generate_by_google_genai(text, file_name, model="gemini-2.0-flash-exp"):
                 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"File of mime type {candidate.inline_data.mime_type} saved to: {temp_path}")
                 image_path = temp_path
                 break
             else:
-                # 이미지 없이 텍스트만 반환되는 경우
                 text_response += chunk.text + "\n"
     
-    # 업로드한 File 객체 제거
     del files
-
     return image_path, text_response
 
 #######################################
-# 3. 임의 알파벳 생성용 함수
-#######################################
-
-def generate_random_letters(length: int) -> str:
-    """
-    length 길이의 임의의 알파벳 문자열을 생성해 반환.
-    대문자+소문자(a-z, A-Z) 범위에서 랜덤 뽑기.
-    """
-    letters = string.ascii_uppercase + string.ascii_lowercase
-    return "".join(random.choice(letters) for _ in range(length))
-
-def fill_text_input_with_random(new_text: str) -> str:
-    """
-    "새로 바꿀 텍스트"의 길이를 세어, 그 길이만큼 임의의 알파벳 문자열 생성 후 반환.
-    이 반환값이 곧 "이미지 안에 들어갈 텍스트" UI(=text_input)에 표시됨.
-    """
-    length = len(new_text)
-    random_letters = generate_random_letters(length)
-    return random_letters
-
-#######################################
-# 4. Gradio 함수 
+# 3. Diffusion + GoogleGenAI를 연결
 #######################################
 
 def generate_initial_image(prompt, text, height, width, steps, scale, seed):
     """
-    FLUX 파이프라인을 사용해 '텍스트가 포함된 이미지를' 먼저 생성.
-    - prompt 내 <text>를 text로 치환
-    - <text>가 없다면 "with clear readable text that says '<text>'"를 자동 붙임
+    1) FLUX 파이프라인을 사용해 'text'가 들어간 이미지를 생성
+       - prompt 내 <text>가 있으면 치환, 없으면 자동 추가
     """
     if "<text>" in prompt:
         combined_prompt = prompt.replace("<text>", text)
     else:
         combined_prompt = f"{prompt} with clear readable text that says '{text}'"
 
-    print(f"[DEBUG] Final combined_prompt: {combined_prompt}")
+    print("[DEBUG] combined_prompt:", combined_prompt)
 
-    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("GenerateInitialImage"):
         result = pipe(
             prompt=[combined_prompt],
             generator=torch.Generator().manual_seed(int(seed)),
@@ -202,9 +168,8 @@ def generate_initial_image(prompt, text, height, width, steps, scale, seed):
 
 def change_text_in_image(original_image, new_text):
     """
-    Google GenAI의 gemini 모델을 통해,
-    업로드된 이미지 내부의 문구를 `new_text`로 변경해주는 함수.
-    (여기서는 new_text가 이미 임의의 알파벳으로 채워진 상태가 됨)
+    2) Gemini 모델을 통해,
+       original_image 내 텍스트를 `new_text`로 변경한 이미지 반환
     """
     try:
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
@@ -220,44 +185,84 @@ def change_text_in_image(original_image, new_text):
             with open(image_path, "rb") as f:
                 image_data = f.read()
             modified_img = Image.open(io.BytesIO(image_data))
-            return modified_img, ""
+            return modified_img
         else:
-            return None, text_response
+            # 이미지가 없이 텍스트만 응답된 경우
+            return None
 
     except Exception as e:
         raise gr.Error(f"Error: {e}")
 
 #######################################
-# 5. Gradio 인터페이스 구성
+# 4. 임의 알파벳 생성
 #######################################
 
-with gr.Blocks(title="Flux + Google GenAI + Random Text Replacement") as demo:
-    gr.Markdown(
-        """
-        # Flux 기반 이미지 생성 + Google GenAI + 임의 알파벳 치환
+def generate_random_letters(length: int) -> str:
+    """
+    length 길이의 임의 알파벳(대소문자) 문자열 생성
+    """
+    letters = string.ascii_lowercase + string.ascii_uppercase
+    return "".join(random.choice(letters) for _ in range(length))
 
-        **기능 요약**  
-        1) "새로 바꿀 텍스트"에 입력된 문자열 길이를 측정  
-        2) 그 길이만큼 **무작위 알파벳**(대소문자)으로 구성된 문자열을 만들어  
-        3) "이미지 안에 들어갈 텍스트" 칸에 자동으로 입력 (즉, 실제 바뀔 텍스트는 완전히 랜덤)  
-        4) 그 후 Google GenAI(gemini)로 이미지 문자를 교체
+#######################################
+# 5. 최종 함수: 버튼 한 번으로
+#    (1) 무작위 알파벳으로 1차 이미지 생성
+#    (2) 진짜 "새로 바꿀 텍스트"로 2차 이미지 생성
+#######################################
+
+def run_full_process(prompt, final_text, height, width, steps, scale, seed):
+    """
+    - final_text의 길이에 맞춰 랜덤 알파벳을 생성 -> 1차 이미지
+    - 그 1차 이미지를 바탕으로, final_text로 교체 -> 2차 최종 이미지
+    """
+    # (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)
+
+    # (C) 2차 이미지: 실제 final_text로 교체
+    final_image = change_text_in_image(random_image, final_text)
+    
+    return [random_image, final_image]
+
+#######################################
+# 6. Gradio UI
+#######################################
+
+with gr.Blocks(title="Flux + Google GenAI (Random & Then Real Text)") as demo:
+    gr.Markdown(
+        """
+        # Flux + Google GenAI: 두 단계에 걸친 텍스트 교체
+
+        **사용 흐름**  
+        1) Prompt에 장면이나 스타일을 작성 (필요하면 `<text>` 구분자 사용)  
+        2) "새로 바꿀 텍스트" 에 최종 원하는 문자열을 입력 (예: "안녕하세요")  
+        3) "Generate Images" 버튼을 누르면,  
+           - 먼저 "새로 바꿀 텍스트" 길이에 맞는 **무작위 알파벳**을 넣어 이미지 생성 (1차 이미지)  
+           - 이어서 **진짜** "새로 바꿀 텍스트"로 다시 교체(2차 최종 이미지)  
+        4) 결과로 두 장의 이미지를 확인할 수 있습니다.  
+
+        **주요 포인트**  
+        - "이미지 안에 들어갈 텍스트"는 **UI에 노출되지 않으며**(사용자 입력 불가), 오직 내부에서 자동 설정됩니다.
+        - 1차 이미지는 완전히 임의의 알파벳 텍스트를 포함합니다.
+        - 2차 이미지는 최종적으로 사용자가 입력한 "새로 바꿀 텍스트"를 포함합니다.
         """
     )
 
     with gr.Row():
         with gr.Column():
-            gr.Markdown("## 1) Step 1: FLUX로 텍스트 포함 이미지 생성")
             prompt_input = gr.Textbox(
                 lines=3,
-                label="이미지 장면/배경 Prompt (use `<text>` placeholder if you like)",
-                placeholder="e.g. A white cat with speech bubble says <text>"
+                label="Prompt (use `<text>` if you want)",
+                placeholder="e.g. A white cat with a speech bubble <text>"
             )
-            text_input = gr.Textbox(
+            final_text_input = gr.Textbox(
                 lines=1,
-                label="이미지 안에 들어갈 텍스트",
-                placeholder="e.g. Hello or 안녕"
+                label="새로 바꿀 텍스트",
+                placeholder="예) 안녕하세요"
             )
             with gr.Accordion("고급 설정 (확장)", open=False):
                 height = gr.Slider(label="Height", minimum=256, maximum=1152, step=64, value=512)
@@ -266,52 +271,17 @@ with gr.Blocks(title="Flux + Google GenAI + Random Text Replacement") as demo:
                 scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.5, value=3.5)
                 seed = gr.Number(label="Seed (reproducibility)", value=1234, precision=0)
 
-            generate_btn = gr.Button("Generate Base Image", variant="primary")
-            generated_image = gr.Image(label="Generated Image (with text)", type="pil")
+            run_btn = gr.Button("Generate Images", variant="primary")
 
         with gr.Column():
-            gr.Markdown(
-                """
-                ## 2) Step 2: "새로 바꿀 텍스트" 길이를 기준으로  
-                임의 알파벳을 "이미지 안에 들어갈 텍스트"로 자동 설정 → 이미지 교체
-                """
-            )
-            new_text_input = gr.Textbox(
-                label="새로 바꿀 텍스트 (입력된 길이를 기준으로 무작위 알파벳 생성)",
-                placeholder="예) Hello123 (길이 8이면, 8글자 랜덤)"
-            )
-
-            # (A) 먼저 "새로 바꿀 텍스트" 글자수를 세어, 무작위 알파벳으로 text_input 채우기
-            # (B) 그런 뒤 change_text_in_image() 실행
-            modify_btn = gr.Button("Generate Random Letters & Change Text in Image", variant="secondary")
-
-            # 최종 결과
-            output_img = gr.Image(label="Modified Image", type="pil")
-            output_txt = gr.Textbox(label="(If only text returned)")
-
-    ###########################
-    # 1) Step1 버튼 액션
-    ###########################
-    generate_btn.click(
-        fn=generate_initial_image,
-        inputs=[prompt_input, text_input, height, width, steps, scale, seed],
-        outputs=[generated_image]
-    )
-
-    ###########################
-    # 2) Step2 버튼 액션
-    ###########################
-    # (A) 무작위 알파벳 생성해 text_input에 반영
-    chain = modify_btn.click(
-        fn=fill_text_input_with_random,
-        inputs=[new_text_input],
-        outputs=[text_input]
-    )
-    # (B) 방금 업데이트된 text_input 값을 이용해, 이미지 텍스트 교체
-    chain.then(
-        fn=change_text_in_image,
-        inputs=[generated_image, text_input],
-        outputs=[output_img, output_txt]
+            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],
+        outputs=[random_image_output, final_image_output]
     )
 
 demo.launch(max_threads=20)