Update app.py

app.py

@@ -9,6 +9,8 @@ import json
 import io
 
 import torch
+from PIL import Image
+
 from safetensors.torch import load_file
 from huggingface_hub import hf_hub_download
 
@@ -20,14 +22,10 @@ from diffusers import FluxPipeline
 from google import genai
 from google.genai import types
 
-# PIL 이미지 처리를 위해 추가
-from PIL import Image
-
 #######################################
 # 0. 환경설정
 #######################################
 
-# 모델 캐시 디렉토리 설정
 BASE_DIR = path.dirname(path.abspath(__file__)) if "__file__" in globals() else os.getcwd()
 CACHE_PATH = path.join(BASE_DIR, "models")
 
@@ -35,9 +33,6 @@ os.environ["TRANSFORMERS_CACHE"] = CACHE_PATH
 os.environ["HF_HUB_CACHE"] = CACHE_PATH
 os.environ["HF_HOME"] = CACHE_PATH
 
-# (예시) Google GenAI 사용:
-# export GAPI_TOKEN="<YOUR_GOOGLE_GENAI_API_KEY>"
-
 # 간단한 타이머 클래스
 class timer:
     def __init__(self, method_name="timed process"):
@@ -56,18 +51,15 @@ class timer:
 if not path.exists(CACHE_PATH):
     os.makedirs(CACHE_PATH, exist_ok=True)
 
-# FLUX 파이프라인 로드
 pipe = FluxPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-dev", 
+    "black-forest-labs/FLUX.1-dev",
     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)
 
-# GPU로 옮기기
 pipe.to(device="cuda", dtype=torch.bfloat16)
 
 #######################################
@@ -86,7 +78,7 @@ def generate_by_google_genai(text, file_name, model="gemini-2.0-flash-exp"):
     - file_name: 원본 이미지(예: .png) 경로
     - model: 사용할 gemini 모델 이름
     """
-    # (1) 환경 변수에서 API 키 가져오기 (필수)
+    # GAPI_TOKEN 환경변수에서 키를 가져옴 (필수)
     api_key = os.getenv("GAPI_TOKEN", None)
     if not api_key:
         raise ValueError(
@@ -94,13 +86,12 @@ def generate_by_google_genai(text, file_name, model="gemini-2.0-flash-exp"):
             "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 준비 (이미지 + 프롬프트)
+    # gemini에 전달할 Content 준비
     contents = [
         types.Content(
             role="user",
@@ -114,7 +105,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,
@@ -127,10 +117,9 @@ 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
-        # 5) 스트림 형태로 응답 받아서 이미지/텍스트 구분 처리
         for chunk in client.models.generate_content_stream(
             model=model,
             contents=contents,
@@ -140,35 +129,32 @@ def generate_by_google_genai(text, file_name, model="gemini-2.0-flash-exp"):
                 continue
             candidate = chunk.candidates[0].content.parts[0]
 
-            # inline_data가 있으면 이미지 응답
+            # 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
                 break
             else:
-                # 이미지 없이 텍스트만 반환되는 경우(희박) 누적
                 text_response += chunk.text + "\n"
     
-    # 업로드한 File 객체 제거
     del files
-
     return image_path, text_response
 
 #######################################
-# 3. Gradio 함수: 
-#    (1) FLUX로 이미지 생성 -> 
-#    (2) Google GenAI로 텍스트 교체
+# 3. Gradio 함수
 #######################################
 
 def generate_initial_image(prompt, text, height, width, steps, scale, seed):
     """
-    FLUX 파이프라인을 사용해 '텍스트가 포함된 이미지를' 먼저 생성하는 함수.
-    prompt: 이미지 배경/장면/스타일 묘사를 위한 프롬프트
-    text: 실제로 이미지에 들어가야 할 문구(예: "안녕하세요", "Hello world" 등)
+    FLUX를 이용해 텍스트가 포함된 이미지를 생성
+    - prompt 내에 <text>라는 특수 구분자가 있으면, 거기에 text가 치환됨.
+    - 그렇지 않은 경우, 기존처럼 prompt 뒤에 “with clear readable text that says ...”�� 추가.
     """
-    combined_prompt = f"{prompt} with clear readable text that says '{text}'"
+    if "<text>" in prompt:
+        combined_prompt = prompt.replace("<text>", text)
+    else:
+        combined_prompt = f"{prompt} with clear readable text that says '{text}'"
 
     with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
         result = pipe(
@@ -185,70 +171,70 @@ def generate_initial_image(prompt, text, height, width, steps, scale, seed):
 
 def change_text_in_image(original_image, new_text):
     """
-    Google GenAI의 gemini 모델을 통해,
+    Gemini 모델을 통해,
     업로드된 이미지 내부의 문구를 `new_text`로 변경해주는 함수.
     """
     try:
-        # 이미지를 임시 파일로 저장
+        # 임시 파일에 먼저 저장
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
             original_path = tmp.name
             original_image.save(original_path)
 
-        # Gemini 모델 호출하여 텍스트 변경
+        # Gemini 모델 호출
         image_path, text_response = generate_by_google_genai(
             text=f"Change the text in this image to: '{new_text}'",
             file_name=original_path
         )
 
-        # 결과가 이미지로 반환된 경우
         if image_path:
-            # Gradio 구버전(4.26.0 이하)에는 decode_base64_to_image가 없으므로 PIL을 직접 사용
+            # Gradio 구버전에는 decode_base64_to_image가 없으므로 PIL을 직접 사용
             with open(image_path, "rb") as f:
                 image_data = f.read()
-
-            # base64를 거치지 않고, PIL Image로 디코딩
             modified_img = Image.open(io.BytesIO(image_data))
             return modified_img, ""
-
         else:
-            # 이미지 없이 텍스트만 응답으로 온 경우
             return None, text_response
 
     except Exception as e:
         raise gr.Error(f"Error: {e}")
 
-
 #######################################
-# 4. Gradio 인터페이스 구성
+# 4. Gradio 인터페이스
 #######################################
 
 with gr.Blocks(title="Flux + Google GenAI Text Replacement") as demo:
     gr.Markdown(
         """
-        # Flux 기반 이미지 생성 + Google GenAI를 통한 텍스트 변환
-        **이 데모는 아래 두 단계를 보여줍니다.**
-        1) **Diffusion 모델(FluxPipeline)을 이용해** 이미지 생성.  
-           - 이때, 사용자가 지정한 텍스트를 이미지 안에 표시하도록 시도합니다.  
-        2) **생성된 이미지를 Google GenAI(gemini) 모델에 전달**하여,  
-           - 이미지 내 텍스트 부분만 다른 문자열로 변경.
+        # Flux Image Generation + Google GenAI Text Replacement
+
+        **Usage Instructions (in English)**  
+        1. Write a prompt that may contain the special placeholder `<text>`.  
+           - Example: `A white cat says <text> in a cartoon style`.  
+        2. Enter the actual text in the "Text to Include in the Image" field.  
+           - Example: `안녕`  
+        3. Click the "Generate Base Image" button.  
+           - The prompt will be transformed so that `<text>` is replaced with your actual text.  
+           - If `<text>` is **not** found, the text will be appended automatically as `with clear readable text that says ...`.  
+        4. (Optional) If you want to change the text again, use the "Change Text in Image" button.
+
         ---
         """
     )
 
     with gr.Row():
         with gr.Column():
-            gr.Markdown("## 1) Step 1: FLUX로 텍스트 포함 이미지 생성")
+            gr.Markdown("## 1) Generate the Base Image (FLUX)")
             prompt_input = gr.Textbox(
                 lines=3,
-                label="이미지 장면/배경 Prompt",
-                placeholder="예) A poster with futuristic neon style..."
+                label="Prompt (with optional `<text>` placeholder)",
+                placeholder="e.g. A white cat says <text> in a cartoon style"
             )
             text_input = gr.Textbox(
                 lines=1,
-                label="이미지 안에 들어갈 텍스트",
-                placeholder="예) 안녕하세요"
+                label="Text to Include in the Image",
+                placeholder="e.g. 안녕"
             )
-            with gr.Accordion("고급 설정 (확장)", open=False):
+            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)
@@ -256,18 +242,13 @@ with gr.Blocks(title="Flux + Google GenAI Text Replacement") as demo:
                 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"
-            )
+            generated_image = gr.Image(label="Generated Image", type="pil")
 
         with gr.Column():
-            gr.Markdown("## 2) Step 2: 생성된 이미지 내 텍스트 수정")
+            gr.Markdown("## 2) (Optional) Change Text in the Generated Image (Gemini)")
             new_text_input = gr.Textbox(
-                label="새로 바꿀 텍스트",
-                placeholder="예) Hello world"
+                label="New Text to Insert",
+                placeholder="e.g. Hello"
             )
             modify_btn = gr.Button("Change Text in Image via Gemini", variant="secondary")
             output_img = gr.Image(label="Modified Image", type="pil")
@@ -286,5 +267,4 @@ with gr.Blocks(title="Flux + Google GenAI Text Replacement") as demo:
         outputs=[output_img, output_txt]
     )
 
-# Gradio 실행
 demo.launch(max_threads=20)