app.py

import random
import os
import uuid
from datetime import datetime
import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import DiffusionPipeline
from PIL import Image, ImageDraw, ImageFont
import requests
import json
import re

# At the beginning of the script, after imports
def check_available_fonts():
    """Check which font files are available"""
    font_files = {
        "나눔고딕": "NanumGothic-Regular.ttf",
        "흑백사진": "BlackAndWhitePicture-Regular.ttf",
        "치론성": "ChironSungHK-VariableFont_wght.ttf",
        "독도": "Dokdo-Regular.ttf",
        "싱글데이": "SingleDay-Regular.ttf"
    }
    
    available_fonts = []
    for font_name, font_file in font_files.items():
        paths = [
            font_file,
            f"./{font_file}",
            os.path.join(os.getcwd(), font_file),
            os.path.join(os.path.dirname(os.path.abspath(__file__)), font_file)
        ]
        
        font_found = False
        for path in paths:
            if os.path.exists(path):
                available_fonts.append(font_name)
                print(f"✓ Found {font_name} at: {path}")
                font_found = True
                break
        
        if not font_found:
            print(f"✗ {font_name} ({font_file}) not found")
    
    return available_fonts

# Check fonts at startup
print("Checking available fonts...")
available_fonts = check_available_fonts()
if not available_fonts:
    print("WARNING: No Korean fonts found! Please ensure font files are in the same directory as app.py")
    available_fonts = ["나눔고딕"]  # Default fallback
else:
    print(f"Available fonts: {', '.join(available_fonts)}")

# Create permanent storage directory
SAVE_DIR = "saved_images"  # Gradio will handle the persistence
if not os.path.exists(SAVE_DIR):
    os.makedirs(SAVE_DIR, exist_ok=True)

# Load the default image
DEFAULT_IMAGE_PATH = "cover1.webp"

device = "cuda" if torch.cuda.is_available() else "cpu"
repo_id = "black-forest-labs/FLUX.1-dev"
adapter_id = "prithivMLmods/EBook-Creative-Cover-Flux-LoRA"

pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
pipeline.load_lora_weights(adapter_id)
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 get_korean_font(font_name, font_size):
    """Load Korean font from the same directory as app.py"""
    font_files = {
        "나눔고딕": "NanumGothic-Regular.ttf",
        "흑백사진": "BlackAndWhitePicture-Regular.ttf",
        "치론성": "ChironSungHK-VariableFont_wght.ttf",
        "독도": "Dokdo-Regular.ttf",
        "싱글데이": "SingleDay-Regular.ttf"
    }
    
    # Get the font file name
    font_file = font_files.get(font_name, "NanumGothic-Regular.ttf")
    print(f"Trying to load font: {font_name} -> {font_file}")  # Debug info
    
    # List of paths to try
    paths_to_try = [
        font_file,  # Current directory
        f"./{font_file}",  # Explicit current directory
        os.path.join(os.getcwd(), font_file),  # Full path from current working directory
        os.path.join(os.path.dirname(os.path.abspath(__file__)), font_file),  # Same directory as script
    ]
    
    # Try each path
    for path in paths_to_try:
        try:
            font = ImageFont.truetype(path, font_size)
            print(f"Successfully loaded font from: {path}")  # Debug info
            return font
        except Exception as e:
            print(f"Failed to load font from {path}: {e}")  # Debug info
            continue
    
    # If specific font not found, try to load NanumGothic as fallback
    print(f"Warning: {font_file} not found. Trying NanumGothic fallback...")  # Debug info
    for path in ["NanumGothic-Regular.ttf", "./NanumGothic-Regular.ttf"]:
        try:
            font = ImageFont.truetype(path, font_size)
            print(f"Loaded fallback font from: {path}")  # Debug info
            return font
        except:
            continue
    
    # Final fallback to default font
    print("Warning: No Korean fonts found. Using default font.")  # Debug info
    return ImageFont.load_default()

def add_text_overlay(image, title_ko, author_ko, 
                    title_position, author_position, text_color, font_name,
                    title_size, author_size):
    """Add Korean text overlay to the generated image"""
    print(f"add_text_overlay called with font_name: {font_name}")  # Debug info
    
    # Create a copy of the image to work with
    img_with_text = image.copy()
    draw = ImageDraw.Draw(img_with_text)
    
    # Load Korean fonts with custom sizes
    title_font = get_korean_font(font_name, title_size)
    author_font = get_korean_font(font_name, author_size)
    
    # Get image dimensions
    img_width, img_height = img_with_text.size
    
    # Define position mappings for horizontal text
    horizontal_position_coords = {
        "상단": (img_width // 2, img_height // 10),
        "중앙": (img_width // 2, img_height // 2),
        "하단": (img_width // 2, img_height * 9 // 10)
    }
    
    # Define position mappings for vertical text
    vertical_position_coords = {
        "좌측상단": (img_width // 10, img_height // 10),
        "중앙상단": (img_width // 2, img_height // 10),
        "우측상단": (img_width * 9 // 10, img_height // 10)
    }
    
    shadow_offset = 2
    
    # Draw title (Korean only)
    if title_ko:
        # Check if title position is vertical
        if title_position in vertical_position_coords:
            # Vertical text rendering
            start_x, start_y = vertical_position_coords[title_position]
            
            # For vertical text, we need to draw each character separately
            for i, char in enumerate(title_ko):
                # Get character bbox
                try:
                    bbox = draw.textbbox((0, 0), char, font=title_font)
                    char_width = bbox[2] - bbox[0]
                    char_height = bbox[3] - bbox[1]
                except:
                    # Fallback for older PIL versions
                    char_width, char_height = draw.textsize(char, font=title_font)
                
                # Calculate position for this character
                char_x = start_x - char_width // 2
                # Use fixed spacing based on font size for consistent vertical spacing
                vertical_spacing = int(title_size * 1.2)  # 120% of font size for better spacing
                char_y = start_y + i * vertical_spacing
                
                # Draw shadow
                draw.text((char_x + shadow_offset, char_y + shadow_offset), 
                         char, font=title_font, fill="black")
                # Draw character
                draw.text((char_x, char_y), 
                         char, font=title_font, fill=text_color)
        else:
            # Horizontal text rendering (existing code)
            title_x, title_y = horizontal_position_coords[title_position]
            
            # Get text bbox for centering
            try:
                bbox = draw.textbbox((0, 0), title_ko, font=title_font)
                text_width = bbox[2] - bbox[0]
                text_height = bbox[3] - bbox[1]
            except:
                # Fallback for older PIL versions
                text_width, text_height = draw.textsize(title_ko, font=title_font)
            
            # Draw text with shadow for better visibility
            draw.text((title_x - text_width // 2 + shadow_offset, title_y - text_height // 2 + shadow_offset), 
                     title_ko, font=title_font, fill="black")
            draw.text((title_x - text_width // 2, title_y - text_height // 2), 
                     title_ko, font=title_font, fill=text_color)
    
    # Draw author (Korean only) - keeping horizontal only for now
    if author_ko:
        # Add "지은이: " prefix if not already present
        if not author_ko.startswith("지은이"):
            author_text = f"지은이: {author_ko}"
        else:
            author_text = author_ko
        
        # Check if author position is vertical
        if author_position in vertical_position_coords:
            # Vertical text rendering for author
            start_x, start_y = vertical_position_coords[author_position]
            
            # Adjust start_y for author (place it lower than title if both are vertical)
            if title_position in vertical_position_coords:
                # Calculate how much space the title takes with new spacing
                title_length = len(title_ko) if title_ko else 0
                vertical_spacing = int(title_size * 1.2)
                start_y += title_length * vertical_spacing + author_size
            
            for i, char in enumerate(author_text):
                # Get character bbox
                try:
                    bbox = draw.textbbox((0, 0), char, font=author_font)
                    char_width = bbox[2] - bbox[0]
                    char_height = bbox[3] - bbox[1]
                except:
                    # Fallback for older PIL versions
                    char_width, char_height = draw.textsize(char, font=author_font)
                
                # Calculate position for this character
                char_x = start_x - char_width // 2
                # Use fixed spacing based on font size for consistent vertical spacing
                vertical_spacing = int(author_size * 1.2)  # 120% of font size for better spacing
                char_y = start_y + i * vertical_spacing
                
                # Draw shadow
                draw.text((char_x + shadow_offset, char_y + shadow_offset), 
                         char, font=author_font, fill="black")
                # Draw character
                draw.text((char_x, char_y), 
                         char, font=author_font, fill=text_color)
        else:
            # Horizontal text rendering (existing code)
            author_x, author_y = horizontal_position_coords[author_position]
            
            # Get text bbox for centering
            try:
                bbox = draw.textbbox((0, 0), author_text, font=author_font)
                text_width = bbox[2] - bbox[0]
                text_height = bbox[3] - bbox[1]
            except:
                # Fallback for older PIL versions
                text_width, text_height = draw.textsize(author_text, font=author_font)
            
            # 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")
            draw.text((author_x - text_width // 2, author_y - text_height // 2), 
                     author_text, font=author_font, fill=text_color)
    
    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")
    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

def load_predefined_images():
    # Return empty list since we're not using predefined images
    return []

@spaces.GPU(duration=120)
def inference(
    prompt: str,
    seed: int,
    randomize_seed: bool,
    width: int,
    height: int,
    guidance_scale: float,
    num_inference_steps: int,
    lora_scale: float,
    title_ko: str,
    author_ko: str,
    title_position: str,
    author_position: str,
    text_color: str,
    font_name: str,
    title_size: int,
    author_size: int,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
):
    print(f"inference called with font_name: {font_name}")  # Debug info
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(seed)
    
    image = pipeline(
        prompt=prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
        joint_attention_kwargs={"scale": lora_scale},
    ).images[0]
    
    # Add text overlay if any Korean text is provided
    if title_ko or author_ko:
        print(f"Adding text overlay with font: {font_name}")  # Debug info
        image = add_text_overlay(image, title_ko, author_ko,
                               title_position, author_position, text_color, font_name,
                               title_size, author_size)
    
    # 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.",
   "A haunting cathedral ruins bathed in ethereal moonlight, with ancient stone archways stretching toward a starlit sky. The title 'WHISPERS OF ETERNITY' appears in weathered silver lettering that seems to float between the pillars. Ghostly wisps of fog curl around crumbling gothic sculptures, while 'By Alexander Blackwood' is inscribed in elegant script that glows with a subtle blue luminescence. Delicate patterns of celestial symbols and arcane runes border the edges. [trigger]",
   
   "A massive ancient tree with crystalline leaves dominates the composition, its translucent branches reaching across a sunset sky streaked with impossible colors. 'THE LUMINOUS Crown' is written in intricate golden calligraphy that intertwines with the branches. Mysterious glowing orbs float among the leaves, casting prismatic light. 'By Isabella Moonshadow' appears to be carved into the tree's bark. Sacred geometry patterns shimmer in the background. [trigger]",
   
   "A dramatic spiral staircase made of weathered copper and stained glass descends into swirling cosmic depths. The title 'CHRONICLES OF THE INFINITE' spans the spiral in bold art deco typography that seems to be crafted from constellations. Nebulae and galaxies swirl in the background, while 'By Marcus Starweaver' appears to be formed from falling stardust. Complex mechanical clockwork elements frame the corners. [trigger]",
   
   "An intricate doorway carved from ancient jade stands solitary in a field of shimmering black sand. 'GATES OF THE IMMORTAL' is emblazoned across the top in powerful metallic letters that seem to be forged from liquid mercury. Ethereal phoenix feathers drift across the scene, leaving trails of golden light. 'By Victoria Jade' flows along the bottom in brushstrokes that resemble living smoke. Sacred Chinese characters appear to float in the background. [trigger]",
   
   "A magnificent underwater city of pearl and coral rises from abyssal depths, illuminated by bioluminescent sea life. 'DEPTHS OF WONDER' ripples across the scene in iridescent letters that appear to be formed from living water. Schools of ethereal fish create flowing patterns of light, while 'By Neptune Rivers' shimmers like mother-of-pearl below. Ancient Atlantean symbols pulse with a subtle aqua glow around the borders. [trigger]",
   
   "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>')
    
    gr.HTML("""<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fginigen-Book-Cover.hf.space">
               <img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Fginigen-Book-Cover.hf.space&countColor=%23263759" />
               </a>""")
    
    with gr.Tabs() as tabs:
        with gr.Tab("Generation"):
            with gr.Column(elem_id="col-container"):
                with gr.Row():
                    prompt = gr.Text(
                        label="Prompt",
                        show_label=False,
                        max_lines=1,
                        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
                result = gr.Image(
                    label="Result",
                    show_label=False,
                    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="제목", placeholder="한글 제목을 입력하세요")
                            title_position = gr.Radio(
                                label="제목 위치",
                                choices=["상단", "중앙", "하단", "좌측상단", "중앙상단", "우측상단"],
                                value="상단",
                                info="좌측상단, 중앙상단, 우측상단은 세로쓰기입니다"
                            )
                            title_size = gr.Slider(
                                label="제목 글자 크기",
                                minimum=20,
                                maximum=200,
                                value=100,
                                step=2
                            )
                        with gr.Column():
                            author_ko = gr.Textbox(label="지은이", placeholder="지은이 이름을 입력하세요")
                            author_position = gr.Radio(
                                label="지은이 위치",
                                choices=["상단", "중앙", "하단", "좌측상단", "중앙상단", "우측상단"],
                                value="하단",
                                info="좌측상단, 중앙상단, 우측상단은 세로쓰기입니다"
                            )
                            author_size = gr.Slider(
                                label="지은이 글자 크기",
                                minimum=16,
                                maximum=60,
                                value=32,
                                step=2
                            )
                    
                    with gr.Row():
                        with gr.Column():
                            font_name = gr.Dropdown(
                                label="폰트 선택",
                                choices=available_fonts,
                                value=available_fonts[0] if available_fonts else "나눔고딕"
                            )
                        with gr.Column():
                            gr.Markdown("### 글자 색상 선택")
                            color_buttons = gr.Radio(
                                label="미리 정의된 색상",
                                choices=[
                                    ("흰색", "#FFFFFF"),
                                    ("검정색", "#000000"),
                                    ("빨간색", "#FF0000"),
                                    ("파란색", "#0000FF"),
                                    ("초록색", "#00FF00"),
                                    ("노란색", "#FFFF00"),
                                    ("주황색", "#FFA500"),
                                    ("보라색", "#800080"),
                                    ("회색", "#808080"),
                                    ("금색", "#FFD700"),
                                    ("은색", "#C0C0C0")
                                ],
                                value="#FFFFFF",
                                type="value"
                            )
                            
                            custom_color = gr.ColorPicker(
                                label="커스텀 색상 선택",
                                value="#FFFFFF",
                                visible=False
                            )
                            
                            text_color = gr.Textbox(
                                value="#FFFFFF",
                                visible=False
                            )

                with gr.Accordion("Advanced Settings", open=False):
                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=42,
                    )
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

                    with gr.Row():
                        width = gr.Slider(
                            label="Width",
                            minimum=256,
                            maximum=MAX_IMAGE_SIZE,
                            step=32,
                            value=768,
                        )
                        height = gr.Slider(
                            label="Height",
                            minimum=256,
                            maximum=MAX_IMAGE_SIZE,
                            step=32,
                            value=1024,
                        )

                    with gr.Row():
                        guidance_scale = gr.Slider(
                            label="Guidance scale",
                            minimum=0.0,
                            maximum=10.0,
                            step=0.1,
                            value=3.5,
                        )
                        num_inference_steps = gr.Slider(
                            label="Number of inference steps",
                            minimum=1,
                            maximum=50,
                            step=1,
                            value=30,
                        )
                        lora_scale = gr.Slider(
                            label="LoRA scale",
                            minimum=0.0,
                            maximum=1.0,
                            step=0.1,
                            value=1.0,
                        )

                gr.Examples(
                    examples=examples,
                    inputs=[prompt],
                    outputs=[result, seed],
                )

        with gr.Tab("Gallery"):
            gallery_header = gr.Markdown("### Generated Images Gallery")
            generated_gallery = gr.Gallery(
                label="Generated Images",
                columns=6,
                show_label=False,
                value=load_generated_images(),
                elem_id="generated_gallery",
                height="auto"
            )
            refresh_btn = gr.Button("🔄 Refresh Gallery")

    # Event handlers
    def refresh_gallery():
        return load_generated_images()
    
    def update_color_visibility(selected_color):
        if selected_color == "custom":
            return gr.update(visible=True), gr.update(value="#FFFFFF")
        else:
            return gr.update(visible=False), gr.update(value=selected_color)
    
    def update_final_color(selected_color, custom_color_value):
        if selected_color == "custom":
            return custom_color_value
        else:
            return selected_color
    
    def on_font_change(font_name):
        print(f"Font changed to: {font_name}")
        return font_name

    refresh_btn.click(
        fn=refresh_gallery,
        inputs=None,
        outputs=generated_gallery,
    )
    
    augment_button.click(
        fn=augment_prompt,
        inputs=[prompt],
        outputs=[prompt],
    )
    
    font_name.change(
        fn=on_font_change,
        inputs=[font_name],
        outputs=[font_name]
    )
    
    color_buttons.change(
        fn=update_color_visibility,
        inputs=[color_buttons],
        outputs=[custom_color, text_color]
    )
    
    custom_color.change(
        fn=update_final_color,
        inputs=[color_buttons, custom_color],
        outputs=[text_color]
    )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=inference,
        inputs=[
            prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            lora_scale,
            title_ko,
            author_ko,
            title_position,
            author_position,
            text_color,
            font_name,
            title_size,
            author_size,
        ],
        outputs=[result, seed, generated_gallery],
    )

demo.queue()
demo.launch()