Update app.py

app.py

@@ -7,7 +7,10 @@ import numpy as np
 import spaces
 import torch
 from diffusers import DiffusionPipeline
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
+import requests
+import json
+import re
 
 # Create permanent storage directory
 SAVE_DIR = "saved_images"  # Gradio will handle the persistence
@@ -28,6 +31,126 @@ pipeline = pipeline.to(device)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
+def is_korean_only(text):
+    """Check if text contains only Korean characters (excluding spaces and punctuation)"""
+    # Remove spaces and common punctuation
+    cleaned_text = re.sub(r'[\s\.,!?]', '', text)
+    # Check if all remaining characters are Korean
+    return bool(cleaned_text) and all('\uAC00' <= char <= '\uD7A3' for char in cleaned_text)
+
+def augment_prompt_with_llm(prompt):
+    """Augment Korean prompt using Friendli LLM API"""
+    token = os.getenv("FRIENDLI_TOKEN")
+    if not token:
+        return prompt  # Return original if no token
+    
+    url = "https://api.friendli.ai/dedicated/v1/chat/completions"
+    headers = {
+        "Authorization": f"Bearer {token}",
+        "Content-Type": "application/json"
+    }
+    
+    # Create a system message for prompt augmentation
+    system_message = """You are an expert at creating detailed, artistic prompts for ebook cover generation. 
+    When given a Korean prompt, expand it into a detailed English description suitable for AI image generation.
+    Focus on visual elements, artistic style, composition, lighting, and mood.
+    Always end the prompt with '[trigger]' to activate the LoRA model."""
+    
+    payload = {
+        "model": "dep89a2fld32mcm",
+        "messages": [
+            {
+                "role": "system",
+                "content": system_message
+            },
+            {
+                "role": "user",
+                "content": f"다음 한국어 프롬프트를 전자책 표지 생성을 위한 상세한 영어 프롬프트로 확장해주세요: {prompt}"
+            }
+        ],
+        "max_tokens": 500,
+        "top_p": 0.8,
+        "stream": False
+    }
+    
+    try:
+        response = requests.post(url, json=payload, headers=headers, timeout=30)
+        if response.status_code == 200:
+            result = response.json()
+            augmented_prompt = result['choices'][0]['message']['content']
+            return augmented_prompt
+        else:
+            print(f"API Error: {response.status_code}")
+            return prompt
+    except Exception as e:
+        print(f"Error calling LLM API: {e}")
+        return prompt
+
+def add_text_overlay(image, title_ko, title_en, author_ko, author_en, 
+                    title_position, author_position, text_color):
+    """Add text overlay to the generated image"""
+    # Create a copy of the image to work with
+    img_with_text = image.copy()
+    draw = ImageDraw.Draw(img_with_text)
+    
+    # Try to load a better font, fallback to default if not available
+    try:
+        # You may need to adjust the font path based on your system
+        title_font_size = 48
+        author_font_size = 32
+        # For production, you'd want to include proper font files
+        title_font = ImageFont.truetype("/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf", title_font_size)
+        author_font = ImageFont.truetype("/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf", author_font_size)
+    except:
+        # Fallback to default font
+        title_font = ImageFont.load_default()
+        author_font = ImageFont.load_default()
+    
+    # Get image dimensions
+    img_width, img_height = img_with_text.size
+    
+    # Define position mappings
+    position_coords = {
+        "Top": (img_width // 2, img_height // 10),
+        "Center": (img_width // 2, img_height // 2),
+        "Bottom": (img_width // 2, img_height * 9 // 10)
+    }
+    
+    # Draw title
+    if title_ko or title_en:
+        title_text = f"{title_ko}\n{title_en}" if title_ko and title_en else title_ko or title_en
+        title_x, title_y = position_coords[title_position]
+        
+        # Get text bbox for centering
+        bbox = draw.textbbox((0, 0), title_text, font=title_font)
+        text_width = bbox[2] - bbox[0]
+        text_height = bbox[3] - bbox[1]
+        
+        # Draw text with shadow for better visibility
+        shadow_offset = 2
+        draw.text((title_x - text_width // 2 + shadow_offset, title_y - text_height // 2 + shadow_offset), 
+                 title_text, font=title_font, fill="black", align="center")
+        draw.text((title_x - text_width // 2, title_y - text_height // 2), 
+                 title_text, font=title_font, fill=text_color, align="center")
+    
+    # Draw author
+    if author_ko or author_en:
+        author_text = f"{author_ko}\n{author_en}" if author_ko and author_en else author_ko or author_en
+        author_x, author_y = position_coords[author_position]
+        
+        # Get text bbox for centering
+        bbox = draw.textbbox((0, 0), author_text, font=author_font)
+        text_width = bbox[2] - bbox[0]
+        text_height = bbox[3] - bbox[1]
+        
+        # Draw text with shadow
+        draw.text((author_x - text_width // 2 + shadow_offset, author_y - text_height // 2 + shadow_offset), 
+                 author_text, font=author_font, fill="black", align="center")
+        draw.text((author_x - text_width // 2, author_y - text_height // 2), 
+                 author_text, font=author_font, fill=text_color, align="center")
+    
+    return img_with_text
+
 def save_generated_image(image, prompt):
     # Generate unique filename with timestamp
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
@@ -70,6 +193,13 @@ def inference(
     guidance_scale: float,
     num_inference_steps: int,
     lora_scale: float,
+    title_ko: str,
+    title_en: str,
+    author_ko: str,
+    author_en: str,
+    title_position: str,
+    author_position: str,
+    text_color: str,
     progress: gr.Progress = gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
@@ -86,12 +216,24 @@ def inference(
         joint_attention_kwargs={"scale": lora_scale},
     ).images[0]
     
+    # Add text overlay if any text is provided
+    if any([title_ko, title_en, author_ko, author_en]):
+        image = add_text_overlay(image, title_ko, title_en, author_ko, author_en,
+                               title_position, author_position, text_color)
+    
     # Save the generated image
     filepath = save_generated_image(image, prompt)
     
     # Return the image, seed, and updated gallery
     return image, seed, load_generated_images()
 
+def augment_prompt(prompt):
+    """Handle prompt augmentation"""
+    if is_korean_only(prompt):
+        augmented = augment_prompt_with_llm(prompt)
+        return augmented
+    return prompt
+
 examples = [
    "An anime-style illustration of a handsome male character with long, dark, flowing hair tied back partially with a traditional hairpiece. He wears a flowing, light-colored traditional East Asian robe with dark accents. His expression is thoughtful and slightly troubled, with his hand near his temple. In the blurred background, there are other figures in similar traditional attire, suggesting a scene of action or conflict in a fantasy setting. The overall mood is serious and dramatic, reminiscent of wuxia or xianxia genres.",
    "A fierce, action-oriented anime illustration of a male knight in full, dark, intricate armor. He has long, flowing dark hair and a confident, determined expression with a slight smirk. He wields a massive, ornate sword with a red glow on its blade, held high above his head in a striking pose. The background is a dramatic, desolate landscape with jagged mountains and a stormy, overcast sky, conveying a sense of epic conflict and adventure.",
@@ -108,8 +250,6 @@ examples = [
    "A colossal steampunk clocktower pierces through storm clouds, its gears and mechanisms visible through crystalline walls. 'TIMEKEEPER'S LEGACY' is constructed from intricate brass and copper mechanisms that appear to be in constant motion. Lightning arcs between copper spires, while 'By Theodore Cogsworth' is etched in burnished bronze below. Mathematical equations and alchemical symbols float in the turbulent sky. [trigger]"
 ]
 
-
-
 with gr.Blocks(theme=gr.themes.Soft(), analytics_enabled=False) as demo:
     gr.HTML('<div class="title"> eBOOK Cover generation </div>')
     
@@ -128,6 +268,7 @@ with gr.Blocks(theme=gr.themes.Soft(), analytics_enabled=False) as demo:
                         placeholder="Enter your prompt",
                         container=False,
                     )
+                    augment_button = gr.Button("증강", scale=0)
                     run_button = gr.Button("Run", scale=0)
 
                 # Modified to include the default image
@@ -137,6 +278,30 @@ with gr.Blocks(theme=gr.themes.Soft(), analytics_enabled=False) as demo:
                     value=DEFAULT_IMAGE_PATH  # Set the default image
                 )
 
+                with gr.Accordion("Text Overlay Settings", open=False):
+                    with gr.Row():
+                        with gr.Column():
+                            title_ko = gr.Textbox(label="Title (Korean)", placeholder="한글 제목")
+                            title_en = gr.Textbox(label="Title (English)", placeholder="English Title")
+                            title_position = gr.Radio(
+                                label="Title Position",
+                                choices=["Top", "Center", "Bottom"],
+                                value="Top"
+                            )
+                        with gr.Column():
+                            author_ko = gr.Textbox(label="Author (Korean)", placeholder="지은이")
+                            author_en = gr.Textbox(label="Author (English)", placeholder="Author Name")
+                            author_position = gr.Radio(
+                                label="Author Position",
+                                choices=["Top", "Center", "Bottom"],
+                                value="Bottom"
+                            )
+                    
+                    text_color = gr.ColorPicker(
+                        label="Text Color",
+                        value="#FFFFFF"
+                    )
+
                 with gr.Accordion("Advanced Settings", open=False):
                     seed = gr.Slider(
                         label="Seed",
@@ -213,6 +378,26 @@ with gr.Blocks(theme=gr.themes.Soft(), analytics_enabled=False) as demo:
         inputs=None,
         outputs=generated_gallery,
     )
+    
+    # Augment button handler
+    augment_button.click(
+        fn=augment_prompt,
+        inputs=[prompt],
+        outputs=[prompt],
+    )
+
+    # Auto-augment Korean prompts
+    def handle_prompt_change(prompt_text):
+        if is_korean_only(prompt_text):
+            return augment_prompt_with_llm(prompt_text)
+        return prompt_text
+    
+    # Optional: Auto-augment on prompt change (commented out to avoid too many API calls)
+    # prompt.change(
+    #     fn=handle_prompt_change,
+    #     inputs=[prompt],
+    #     outputs=[prompt]
+    # )
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
@@ -226,6 +411,13 @@ with gr.Blocks(theme=gr.themes.Soft(), analytics_enabled=False) as demo:
             guidance_scale,
             num_inference_steps,
             lora_scale,
+            title_ko,
+            title_en,
+            author_ko,
+            author_en,
+            title_position,
+            author_position,
+            text_color,
         ],
         outputs=[result, seed, generated_gallery],
     )