Update app.py

app.py

@@ -9,41 +9,186 @@ import torch
 from diffusers import DiffusionPipeline
 from PIL import Image
 
-# Apply more comprehensive patches to Gradio's utility functions
+# -----------------------------
+# Gemini API & Text Rendering 관련 추가 모듈
+# -----------------------------
+import re
+import tempfile
+import io
+import logging
+import base64
+import string
+import requests
+from google import genai
+from google.genai import types
+
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+
+def maybe_translate_to_english(text: str) -> str:
+    """
+    텍스트에 한글이 포함되어 있으면 간단한 규칙에 따라 영어로 변환.
+    """
+    if not text or not re.search("[가-힣]", text):
+        return text
+    try:
+        translations = {
+            "안녕하세요": "Hello",
+            "환영합니다": "Welcome",
+            "아름다운 당신": "Beautiful You",
+            "안녕": "Hello",
+            "고양이": "Cat",
+            "배너": "Banner",
+            "썬글라스": "Sunglasses",
+            "착용한": "wearing",
+            "흰색": "white"
+        }
+        for kr, en in translations.items():
+            if kr in text:
+                text = text.replace(kr, en)
+        print(f"[TRANSLATE] Translated Korean text: '{text}'")
+        return text
+    except Exception as e:
+        print(f"[WARNING] Translation failed: {e}")
+        return text
+
+def save_binary_file(file_name, data):
+    with open(file_name, "wb") as f:
+        f.write(data)
+
+def generate_by_google_genai(text, file_name, model="gemini-2.0-flash-exp"):
+    """
+    Gemini API를 호출하여 텍스트 기반 이미지 편집/생성을 수행.
+    """
+    api_key = os.getenv("GAPI_TOKEN", None)
+    if not 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)]
+    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),
+            ],
+        ),
+    ]
+    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",
+    )
+    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,
+            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]
+            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"
+    del files
+    return image_path, text_response
+
+def change_text_in_image_two_times(original_image, instruction):
+    """
+    Gemini API를 두 번 호출하여 텍스트 수정된 이미지 2개(버전 A, B)를 반환.
+    """
+    if original_image is None:
+        raise gr.Error("처리할 이미지가 없습니다. 먼저 이미지를 생성해주세요.")
+    results = []
+    for version_tag in ["(A)", "(B)"]:
+        mod_instruction = f"{instruction} {version_tag}"
+        try:
+            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
+                original_path = tmp.name
+                if isinstance(original_image, Image.Image):
+                    original_image.save(original_path, format="PNG")
+                    print(f"[DEBUG] Saved image to temporary file: {original_path}")
+                else:
+                    raise gr.Error(f"예상된 PIL Image가 아닌 {type(original_image)} 타입이 제공되었습니다.")
+            print(f"[DEBUG] Google Gemini API에 보내는 지시사항: {mod_instruction}")
+            image_path, text_response = generate_by_google_genai(
+                text=mod_instruction,
+                file_name=original_path
+            )
+            if image_path:
+                try:
+                    with open(image_path, "rb") as f:
+                        image_data = f.read()
+                    new_img = Image.open(io.BytesIO(image_data))
+                    results.append(new_img)
+                except Exception as img_err:
+                    print(f"[ERROR] Failed to process Gemini image: {img_err}")
+                    results.append(original_image)
+            else:
+                print(f"[WARNING] 이미지가 반환되지 않았습니다. 텍스트 응답: {text_response}")
+                results.append(original_image)
+        except Exception as e:
+            logging.exception(f"Text modification error: {e}")
+            results.append(original_image)
+    return results
+
+def gemini_text_rendering(image, rendering_text):
+    """
+    주어진 이미지에 대해 Gemini API를 사용해 텍스트 렌더링을 적용.
+    """
+    rendering_text_en = maybe_translate_to_english(rendering_text)
+    instruction = f"Render the following text on the image in a clear, visually appealing manner: {rendering_text_en}."
+    rendered_images = change_text_in_image_two_times(image, instruction)
+    if rendered_images and len(rendered_images) > 0:
+        return rendered_images[0]
+    return image
+
+def apply_text_rendering(image, rendering_text):
+    """
+    입력된 텍스트가 있으면 Gemini API를 통해 이미지에 텍스트 렌더링을 적용하고, 없으면 원본 이미지를 그대로 반환.
+    """
+    if rendering_text and rendering_text.strip():
+        return gemini_text_rendering(image, rendering_text)
+    return image
+
+# -----------------------------
+# 기존 Diffusion Pipeline 관련 코드
+# -----------------------------
 import gradio_client.utils
 import types
 
-# Patch 1: Fix the _json_schema_to_python_type function
 original_json_schema = gradio_client.utils._json_schema_to_python_type
-
 def patched_json_schema(schema, defs=None):
-    # Handle boolean values directly
     if isinstance(schema, bool):
         return "bool"
-    
-    # Handle cases where 'additionalProperties' is a boolean
     try:
         if "additionalProperties" in schema and isinstance(schema["additionalProperties"], bool):
             schema["additionalProperties"] = {"type": "any"}
     except (TypeError, KeyError):
         pass
-    
-    # Call the original function
     try:
         return original_json_schema(schema, defs)
     except Exception as e:
-        # Fallback to a safe value when the schema can't be parsed
         return "any"
-
-# Replace the original function with our patched version
 gradio_client.utils._json_schema_to_python_type = patched_json_schema
 
-# Create permanent storage directory
-SAVE_DIR = "saved_images"  # Gradio will handle the persistence
+SAVE_DIR = "saved_images"
 if not os.path.exists(SAVE_DIR):
     os.makedirs(SAVE_DIR, exist_ok=True)
 
-# Safe settings for model loading
 device = "cuda" if torch.cuda.is_available() else "cpu"
 repo_id = "black-forest-labs/FLUX.1-dev"
 adapter_id = "openfree/flux-chatgpt-ghibli-lora"
@@ -72,37 +217,27 @@ def load_model_with_retry(max_retries=5):
             else:
                 raise Exception(f"Failed to load model after {max_retries} attempts: {e}")
 
-# Load the model
 pipeline = load_model_with_retry()
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 def save_generated_image(image, prompt):
-    # Generate unique filename with timestamp
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     unique_id = str(uuid.uuid4())[:8]
     filename = f"{timestamp}_{unique_id}.png"
     filepath = os.path.join(SAVE_DIR, filename)
-    
-    # Save the image
     image.save(filepath)
-    
-    # Save metadata
     metadata_file = os.path.join(SAVE_DIR, "metadata.txt")
     with open(metadata_file, "a", encoding="utf-8") as f:
         f.write(f"{filename}|{prompt}|{timestamp}\n")
-    
     return filepath
 
 def load_generated_images():
     if not os.path.exists(SAVE_DIR):
         return []
-    
-    # Load all images from the directory
     image_files = [os.path.join(SAVE_DIR, f) for f in os.listdir(SAVE_DIR) 
                   if f.endswith(('.png', '.jpg', '.jpeg', '.webp'))]
-    # Sort by creation time (newest first)
     image_files.sort(key=lambda x: os.path.getctime(x), reverse=True)
     return image_files
 
@@ -121,8 +256,6 @@ def inference(
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device=device).manual_seed(seed)
-    
-    # Error handling for the inference process
     try:
         image = pipeline(
             prompt=prompt,
@@ -133,33 +266,25 @@ def inference(
             generator=generator,
             joint_attention_kwargs={"scale": lora_scale},
         ).images[0]
-        
-        # Save the generated image
         filepath = save_generated_image(image, prompt)
-        
-        # Return the image, seed, and updated gallery
         return image, seed, load_generated_images()
     except Exception as e:
-        # Log the error and return a simple error image
         print(f"Error during inference: {e}")
         error_img = Image.new('RGB', (width, height), color='red')
         return error_img, seed, load_generated_images()
 
+# -----------------------------
+# Gradio UI (입력 프롬프트 아래에 Text Rendering 입력란 추가)
+# -----------------------------
 examples = [
     "Ghibli style futuristic stormtrooper with glossy white armor and a sleek helmet, standing heroically on a lush alien planet, vibrant flowers blooming around, soft sunlight illuminating the scene, a gentle breeze rustling the leaves. The armor reflects the pink and purple hues of the alien sunset, creating an ethereal glow around the figure. [trigger]",
-    
     "Ghibli style young mechanic girl in a floating workshop, surrounded by hovering tools and glowing mechanical parts, her blue overalls covered in oil stains, tinkering with a semi-transparent robot companion. Magical sparks fly as she works, while floating islands with waterfalls drift past her open workshop window. [trigger]",
-    
     "Ghibli style ancient forest guardian robot, covered in moss and flowering vines, sitting peacefully in a crystal-clear lake. Its gentle eyes glow with soft blue light, while bioluminescent dragonflies dance around its weathered metal frame. Ancient tech symbols on its surface pulse with a gentle rhythm. [trigger]",
-    
     "Ghibli style sky whale transport ship, its metallic skin adorned with traditional Japanese patterns, gliding through cotton candy clouds at sunrise. Small floating gardens hang from its sides, where workers in futuristic kimonos tend to glowing plants. Rainbow auroras shimmer in the background. [trigger]",
-    
     "Ghibli style cyber-shrine maiden with flowing holographic robes, performing a ritual dance among floating lanterns and digital cherry blossoms. Her traditional headdress emits soft light patterns, while spirit-like AI constructs swirl around her in elegant patterns. The scene is set in a modern shrine with both ancient wood and sleek chrome elements. [trigger]",
-    
     "Ghibli style robot farmer tending to floating rice paddies in the sky, wearing a traditional straw hat with advanced sensors. Its gentle movements create ripples in the water as it plants glowing rice seedlings. Flying fish leap between the terraced fields, leaving trails of sparkles in their wake, while future Tokyo's spires gleam in the distance. [trigger]"
 ]
 
-# Enhanced CSS for a more visually refined UI
 css = """
 :root {
     --primary-color: #6a92cc;
@@ -172,17 +297,14 @@ css = """
     --shadow: 0 4px 12px rgba(0,0,0,0.08);
     --font-main: 'Poppins', -apple-system, BlinkMacSystemFont, sans-serif;
 }
-
 body {
     background-color: var(--background-color);
     font-family: var(--font-main);
 }
-
 .gradio-container {
     margin: 0 auto;
     max-width: 1200px !important;
 }
-
 .main-header {
     text-align: center;
     padding: 2rem 1rem 1rem;
@@ -192,32 +314,27 @@ body {
     border-radius: var(--border-radius);
     box-shadow: var(--shadow);
 }
-
 .main-header h1 {
     font-size: 2.5rem;
     margin-bottom: 0.5rem;
     font-weight: 700;
     text-shadow: 0 2px 4px rgba(0,0,0,0.2);
 }
-
 .main-header p {
     font-size: 1rem;
     margin-bottom: 0.5rem;
     opacity: 0.9;
 }
-
 .main-header a {
     color: var(--secondary-color);
     text-decoration: none;
     font-weight: 600;
     transition: all 0.2s ease;
 }
-
 .main-header a:hover {
     text-decoration: underline;
     opacity: 0.9;
 }
-
 .container {
     background-color: var(--panel-background);
     padding: 1.5rem;
@@ -225,7 +342,6 @@ body {
     box-shadow: var(--shadow);
     margin-bottom: 1.5rem;
 }
-
 button.primary {
     background: var(--primary-color) !important;
     border: none !important;
@@ -236,13 +352,11 @@ button.primary {
     box-shadow: 0 2px 5px rgba(0,0,0,0.1) !important;
     transition: all 0.2s ease !important;
 }
-
 button.primary:hover {
     background: var(--primary-hover) !important;
     transform: translateY(-2px) !important;
     box-shadow: 0 4px 8px rgba(0,0,0,0.15) !important;
 }
-
 button.secondary {
     background: white !important;
     border: 1px solid #ddd !important;
@@ -253,64 +367,51 @@ button.secondary {
     box-shadow: 0 2px 5px rgba(0,0,0,0.05) !important;
     transition: all 0.2s ease !important;
 }
-
 button.secondary:hover {
     background: #f5f5f5 !important;
     transform: translateY(-2px) !important;
 }
-
 .gr-box {
     border-radius: var(--border-radius) !important;
     border: 1px solid #e0e0e0 !important;
 }
-
 .gr-panel {
     border-radius: var(--border-radius) !important;
 }
-
 .gr-input {
     border-radius: 8px !important;
     border: 1px solid #ddd !important;
     padding: 12px !important;
 }
-
 .gr-form {
     border-radius: var(--border-radius) !important;
     background-color: var(--panel-background) !important;
 }
-
 .gr-accordion {
     border-radius: var(--border-radius) !important;
     overflow: hidden !important;
 }
-
 .gr-button {
     border-radius: 8px !important;
 }
-
 .gallery-item {
     border-radius: var(--border-radius) !important;
     transition: all 0.3s ease !important;
 }
-
 .gallery-item:hover {
     transform: scale(1.02) !important;
     box-shadow: 0 6px 15px rgba(0,0,0,0.1) !important;
 }
-
 .tabs {
     border-radius: var(--border-radius) !important;
     overflow: hidden !important;
 }
-
 footer {
     display: none !important;
 }
-
 .settings-accordion legend span {
     font-weight: 600 !important;
 }
-
 .example-prompt {
     font-size: 0.9rem;
     color: #555;
@@ -322,12 +423,9 @@ footer {
     cursor: pointer;
     transition: all 0.2s;
 }
-
 .example-prompt:hover {
     background: #eef2f8;
 }
-
-/* Status indicators */
 .status-generating {
     color: #ffa200;
     font-weight: 500;
@@ -335,7 +433,6 @@ footer {
     align-items: center;
     gap: 8px;
 }
-
 .status-generating::before {
     content: "";
     display: inline-block;
@@ -345,7 +442,6 @@ footer {
     background-color: #ffa200;
     animation: pulse 1.5s infinite;
 }
-
 .status-complete {
     color: #00c853;
     font-weight: 500;
@@ -353,7 +449,6 @@ footer {
     align-items: center;
     gap: 8px;
 }
-
 .status-complete::before {
     content: "";
     display: inline-block;
@@ -362,42 +457,27 @@ footer {
     border-radius: 50%;
     background-color: #00c853;
 }
-
 @keyframes pulse {
-    0% {
-        opacity: 0.6;
-    }
-    50% {
-        opacity: 1;
-    }
-    100% {
-        opacity: 0.6;
-    }
+    0% { opacity: 0.6; }
+    50% { opacity: 1; }
+    100% { opacity: 0.6; }
 }
-
-/* Custom accordions and tabs styling */
 .gr-accordion-title {
     font-weight: 600 !important;
     color: var(--text-color) !important;
 }
-
 .tabs button {
     font-weight: 500 !important;
     padding: 10px 16px !important;
 }
-
 .tabs button.selected {
     font-weight: 600 !important;
     color: var(--primary-color) !important;
     background: rgba(106, 146, 204, 0.1) !important;
 }
-
-/* Improve slider appearance */
 .gr-slider-container {
     padding: 10px 0 !important;
 }
-
-/* Enhanced Markdown content */
 .gr-prose h3 {
     font-weight: 600 !important;
     color: var(--primary-color) !important;
@@ -405,10 +485,8 @@ footer {
 }
 """
 
-# Use a cleaner, more visual UI configuration
 with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
     with gr.Column():
-        # Custom header with improved styling
         gr.HTML('''
         <div class="main-header">
             <h1>✨ FLUX Ghibli LoRA Generator ✨</h1>
@@ -424,6 +502,12 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
                         placeholder="Describe your Ghibli-style image here...",
                         lines=3
                     )
+                    # ★ 새롭게 추가된 Text Rendering 입력란
+                    text_rendering = gr.Textbox(
+                        label="Text Rendering (Multilingual: English, Korean...)",
+                        placeholder="Man saying '안녕' in 'speech bubble'",
+                        lines=1
+                    )
                     
                     with gr.Row():
                         run_button = gr.Button("✨ Generate Image", elem_classes="primary")
@@ -439,7 +523,6 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
                                 value=42,
                             )
                             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
                         with gr.Row():
                             width = gr.Slider(
                                 label="Width",
@@ -455,7 +538,6 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
                                 step=32,
                                 value=768,
                             )
-
                         with gr.Row():
                             guidance_scale = gr.Slider(
                                 label="Guidance scale",
@@ -464,7 +546,6 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
                                 step=0.1,
                                 value=3.5,
                             )
-                            
                         with gr.Row():
                             num_inference_steps = gr.Slider(
                                 label="Steps",
@@ -483,13 +564,11 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
                 
                 with gr.Group(elem_classes="container"):
                     gr.Markdown("### ✨ Example Prompts")
-                    # Create HTML for examples manually
                     examples_html = '\n'.join([f'<div class="example-prompt">{example}</div>' for example in examples])
                     example_container = gr.HTML(examples_html)
             
             with gr.Column(scale=4):
                 with gr.Group(elem_classes="container"):
-                    # Image result container with status indicator
                     with gr.Group():
                         generation_status = gr.HTML('<div class="status-complete">Ready to generate</div>')
                         result = gr.Image(label="Generated Image", elem_id="result-image")
@@ -509,7 +588,6 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
                     elem_classes="gallery-item"
                 )
 
-    # Event handlers
     def refresh_gallery():
         return load_generated_images()
 
@@ -534,7 +612,7 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
         outputs=[prompt, result, seed_text, generation_status]
     )
     
-    # Update with status indicators for generation process
+    # 체인에 마지막에 텍스트 렌더링 적용 (text_rendering 입력값이 있으면)
     run_button.click(
         fn=before_generate,
         inputs=None,
@@ -556,6 +634,10 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
         fn=after_generate,
         inputs=[result, seed_text, generated_gallery],
         outputs=[result, seed_text, generated_gallery, generation_status],
+    ).then(
+        fn=apply_text_rendering,
+        inputs=[result, text_rendering],
+        outputs=result,
     )
     
     prompt.submit(
@@ -579,36 +661,34 @@ with gr.Blocks(css=css, analytics_enabled=False, theme="soft") as demo:
         fn=after_generate,
         inputs=[result, seed_text, generated_gallery],
         outputs=[result, seed_text, generated_gallery, generation_status],
+    ).then(
+        fn=apply_text_rendering,
+        inputs=[result, text_rendering],
+        outputs=result,
     )
     
-    # Custom JavaScript for handling example prompts
     gr.HTML("""
     <script>
     document.addEventListener('DOMContentLoaded', function() {
-        // Add click handlers to example prompts
         setTimeout(() => {
             const examples = document.querySelectorAll('.example-prompt');
             const promptInput = document.querySelector('textarea');
-            
             examples.forEach(example => {
                 example.addEventListener('click', function() {
                     promptInput.value = this.textContent.trim();
-                    // Trigger input event to update Gradio's state
                     const event = new Event('input', { bubbles: true });
                     promptInput.dispatchEvent(event);
                 });
             });
-        }, 1000); // Small delay to ensure elements are loaded
+        }, 1000);
     });
     </script>
     """)
 
-# Launch with fallback options
 try:
     demo.queue(concurrency_count=1, max_size=20)
     demo.launch(debug=True, show_api=False)
 except Exception as e:
     print(f"Error during launch: {e}")
     print("Trying alternative launch configuration...")
-    # Skip queue and simplify launch parameters
-    demo.launch(debug=True, show_api=False, share=False)
+    demo.launch(debug=True, show_api=False, share=False)