Update app.py

app.py

@@ -1,22 +1,40 @@
-import spaces
-import argparse
 import os
 import time
 from os import path
+import tempfile
+import uuid
+import base64
+import mimetypes
+import json
+
+import torch
 from safetensors.torch import load_file
 from huggingface_hub import hf_hub_download
 
-cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
-os.environ["TRANSFORMERS_CACHE"] = cache_path
-os.environ["HF_HUB_CACHE"] = cache_path
-os.environ["HF_HOME"] = cache_path
-
+# Diffusers 관련 라이브러리
 import gradio as gr
-import torch
 from diffusers import FluxPipeline
 
-torch.backends.cuda.matmul.allow_tf32 = True
+# Google GenAI 라이브러리
+from google import genai
+from google.genai import types
+
+#######################################
+# 0. 환경설정
+#######################################
+
+# 모델 캐시 디렉토리 설정
+BASE_DIR = path.dirname(path.abspath(__file__)) if "__file__" in globals() else os.getcwd()
+CACHE_PATH = path.join(BASE_DIR, "models")
 
+os.environ["TRANSFORMERS_CACHE"] = CACHE_PATH
+os.environ["HF_HUB_CACHE"] = CACHE_PATH
+os.environ["HF_HOME"] = CACHE_PATH
+
+# Google GenAI 사용을 위해서는 다음과 같은 환경 변수가 필요합니다.
+# os.environ["GAPI_TOKEN"] = "<YOUR_GOOGLE_GENAI_API_KEY>"
+
+# 작업 시간 측정을 위한 간단한 타이머 클래스
 class timer:
     def __init__(self, method_name="timed process"):
         self.method = method_name
@@ -27,83 +45,234 @@ class timer:
         end = time.time()
         print(f"{self.method} took {str(round(end - self.start, 2))}s")
 
-if not path.exists(cache_path):
-    os.makedirs(cache_path, exist_ok=True)
+#######################################
+# 1. FLUX 파이프라인 로드
+#######################################
+
+if not path.exists(CACHE_PATH):
+    os.makedirs(CACHE_PATH, exist_ok=True)
 
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
-pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
+# FLUX 파이프라인 로드
+pipe = FluxPipeline.from_pretrained(
+    "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)
 
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        <div style="text-align: center; max-width: 650px; margin: 0 auto;">
-            <h1 style="font-size: 2.5rem; font-weight: 700; margin-bottom: 1rem; display: contents;">Hyper-FLUX-8steps-LoRA</h1>
-            <p style="font-size: 1rem; margin-bottom: 1.5rem;">AutoML team from ByteDance</p>
-        </div>
-        """
+#######################################
+# 2. Google GenAI를 통한 이미지 내 텍스트 변환 함수
+#######################################
+
+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 모델 이름
+    """
+    # 1) Google Client 초기화
+    client = genai.Client(api_key=os.getenv("GAPI_TOKEN"))
+    
+    # 2) 이미지 업로드
+    files = [client.files.upload(file=file_name)]
+    
+    # 3) gemini에 전달할 Content 준비 (이미지 + 프롬프트)
+    contents = [
+        types.Content(
+            role="user",
+            parts=[
+                types.Part.from_uri(
+                    file_uri=files[0].uri,
+                    mime_type=files[0].mime_type,
+                ),
+                types.Part.from_text(text=text),
+            ],
+        ),
+    ]
+
+    # 4) 생성/변환 설정
+    generate_content_config = types.GenerateContentConfig(
+        temperature=1,
+        top_p=0.95,
+        top_k=40,
+        max_output_tokens=8192,
+        response_modalities=["image", "text"],
+        response_mime_type="text/plain",
     )
 
-    with gr.Row():
-        with gr.Column(scale=3):
-            with gr.Group():
-                prompt = gr.Textbox(
-                    label="Your Image Description",
-                    placeholder="E.g., A serene landscape with mountains and a lake at sunset",
-                    lines=3
-                )
-                
-                with gr.Accordion("Advanced Settings", open=False):
-                    with gr.Group():
-                        with gr.Row():
-                            height = gr.Slider(label="Height", minimum=256, maximum=1152, step=64, value=1024)
-                            width = gr.Slider(label="Width", minimum=256, maximum=1152, step=64, value=1024)
-                        
-                        with gr.Row():
-                            steps = gr.Slider(label="Inference Steps", minimum=6, maximum=25, step=1, value=8)
-                            scales = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=5.0, step=0.1, value=3.5)
-                        
-                        seed = gr.Number(label="Seed (for reproducibility)", value=3413, precision=0)
-                
-                generate_btn = gr.Button("Generate Image", variant="primary", scale=1)
-
-        with gr.Column(scale=4):
-            output = gr.Image(label="Your Generated Image")
+    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,
+            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]
+
+            # 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"
     
+    # 업로드 파일(google.genai.files.File) 객체 제거
+    del files
+
+    return image_path, text_response
+
+#######################################
+# 3. Gradio 함수: (1) FLUX로 이미지 생성 -> (2) Google GenAI로 텍스트 교체
+#######################################
+
+def generate_initial_image(prompt, text, height, width, steps, scale, seed):
+    """
+    FLUX 파이프라인을 사용해 '텍스트가 포함된 이미지를' 먼저 생성하는 함수.
+    prompt: 이미지 배경/장면/스타일 묘사를 위한 프롬프트
+    text: 실제로 이미지에 들어가야 할 문구(예: "안녕하세요", "Hello world" 등)
+    """
+    # 이미지에 텍스트를 포함시키려면 프롬프트에 직접 문구 요청을 넣는 것이 중요.
+    # Diffusion 모델에 따라 잘 반영되지 않을 수도 있으니, 구체적으로 기재할수록 유리.
+    # 예: "A poster with large bold Korean text that says '안녕하세요' in red color ..."
+    # 여기서는 간단히 prompt 뒤에 텍스트 삽입 예시를 보여줌
+    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(
+            prompt=[combined_prompt],
+            generator=torch.Generator().manual_seed(int(seed)),
+            num_inference_steps=int(steps),
+            guidance_scale=float(scale),
+            height=int(height),
+            width=int(width),
+            max_sequence_length=256
+        ).images[0]
+
+    return result
+
+def change_text_in_image(original_image, new_text):
+    """
+    Google GenAI의 gemini 모델을 통해,
+    업로드된 이미지 내부의 문구를 `new_text`로 변경해주는 함수.
+    """
+    try:
+        # 임시 파일에 먼저 저장
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
+            original_path = tmp.name
+            original_image.save(original_path)
+
+        # 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:
+            modified_img = gr.processing_utils.decode_base64_to_image(
+                base64.b64encode(open(image_path, "rb").read())
+            )
+            return modified_img, ""  # (결과 이미지, 빈 텍스트)
+        else:
+            # 이미지가 없이 텍스트만 응답으로 온 경우
+            return None, text_response
+
+    except Exception as e:
+        raise gr.Error(f"Error: {e}")
+
+
+#######################################
+# 4. Gradio 인터페이스 구성
+#######################################
+
+with gr.Blocks(title="Flux + Google GenAI Text Replacement") as demo:
     gr.Markdown(
         """
-        <div style="max-width: 650px; margin: 2rem auto; padding: 1rem; border-radius: 10px; background-color: #f0f0f0;">
-            <h2 style="font-size: 1.5rem; margin-bottom: 1rem;">How to Use</h2>
-            <ol style="padding-left: 1.5rem;">
-                <li>Enter a detailed description of the image you want to create.</li>
-                <li>Adjust advanced settings if desired (tap to expand).</li>
-                <li>Tap "Generate Image" and wait for your creation!</li>
-            </ol>
-            <p style="margin-top: 1rem; font-style: italic;">Tip: Be specific in your description for best results!</p>
-        </div>
+        # Flux 기반 이미지 생성 + Google GenAI를 통한 텍스트 변환
+        **이 데모는 아래 두 단계를 보여줍니다.**
+
+        1) **Diffusion 모델(FluxPipeline)을 이용해** 이미지 생성.  
+           - 이때, 사용자가 지정한 텍스트를 이미지 안에 표시하도록 시도합니다.  
+        2) **생성된 이미지를 Google GenAI(gemini) 모델에 전달**하여,  
+           - 이미지 내 텍스트 부분만 다른 문자열로 변경.
+
+        ---
         """
     )
 
-    @spaces.GPU
-    def process_image(height, width, steps, scales, prompt, seed):
-        global pipe
-        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
-            return pipe(
-                prompt=[prompt],
-                generator=torch.Generator().manual_seed(int(seed)),
-                num_inference_steps=int(steps),
-                guidance_scale=float(scales),
-                height=int(height),
-                width=int(width),
-                max_sequence_length=256
-            ).images[0]
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("## 1) Step 1: FLUX로 텍스트 포함 이미지 생성")
+            prompt_input = gr.Textbox(
+                lines=3,
+                label="이미지 장면/배경 Prompt",
+                placeholder="예) A poster with futuristic neon style..."
+            )
+            text_input = gr.Textbox(
+                lines=1,
+                label="이미지 안에 들어갈 텍스트",
+                placeholder="예) 안녕하세요"
+            )
+            with gr.Accordion("고급 설정 (확장)", 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=5.0, step=0.1, 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"
+            )
+
+        with gr.Column():
+            gr.Markdown("## 2) Step 2: 생성된 이미지 내 텍스트 수정")
+            new_text_input = gr.Textbox(
+                label="새로 바꿀 텍스트",
+                placeholder="예) Hello world"
+            )
+            modify_btn = gr.Button("Change Text in Image via Gemini", variant="secondary")
+            output_img = gr.Image(label="Modified Image", type="pil")
+            output_txt = gr.Textbox(label="(If only text returned)")
 
+    # 버튼 액션 연결
     generate_btn.click(
-        process_image,
-        inputs=[height, width, steps, scales, prompt, seed],
-        outputs=output
+        fn=generate_initial_image,
+        inputs=[prompt_input, text_input, height, width, steps, scale, seed],
+        outputs=[generated_image]
+    )
+
+    modify_btn.click(
+        fn=change_text_in_image,
+        inputs=[generated_image, new_text_input],
+        outputs=[output_img, output_txt]
     )
 
+# 실제 실행 시에는 아래와 같이 demo.launch()를 호출합니다.
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue(concurrency_count=1, max_size=20).launch()