app.py

import gradio as gr
import requests
from PIL import Image, ImageDraw, ImageFont
import io
import openai
from text_overlay import create_pattern_template, get_pattern_template, add_text_to_image, get_font_path

# Image Generation Functions
def generate_ideogram_image(api_key, prompt, aspect_ratio="ASPECT_1_1"):
    """Generate image using Ideogram API"""
    if not api_key or not prompt:
        return None, "Please provide both API key and prompt", None
    
    url = "https://api.ideogram.ai/generate"
    headers = {
        "Api-Key": api_key,
        "Content-Type": "application/json"
    }
    
    payload = {
        "image_request": {
            "prompt": prompt,
            "aspect_ratio": aspect_ratio,
            "model": "V_2",
            "magic_prompt_option": "AUTO"
        }
    }
    
    try:
        response = requests.post(url, json=payload, headers=headers)
        if response.status_code == 200:
            result = response.json()
            if result.get("data"):
                image_url = result["data"][0]["url"]
                img_response = requests.get(image_url)
                image = Image.open(io.BytesIO(img_response.content))
                safe_copy = safe_image_copy(image)
                return image, "Image generated successfully!", safe_copy
            else:
                return None, "No image data received", None
        else:
            return None, f"Error: {response.status_code} - {response.text}", None
    except Exception as e:
        return None, f"Error generating image: {str(e)}", None

def generate_dalle_image(api_key, prompt, size="1024x1024"):
    """Generate image using DALL-E 3 API"""
    if not api_key or not prompt:
        return None, "Please provide both API key and prompt", None
    
    try:
        client = openai.OpenAI(api_key=api_key)
        response = client.images.generate(
            model="dall-e-3",
            prompt=prompt,
            size=size,
            quality="standard",
            n=1,
        )
        
        image_url = response.data[0].url
        img_response = requests.get(image_url)
        image = Image.open(io.BytesIO(img_response.content))
        safe_copy = safe_image_copy(image)
        return image, "Image generated successfully with DALL-E 3!", safe_copy
    
    except Exception as e:
        return None, f"Error generating image: {str(e)}", None

def process_uploaded_image(image):
    """Process uploaded image"""
    if image is None:
        return None, "No image uploaded", None
    try:
        safe_copy = safe_image_copy(image)
        return image, "Image uploaded successfully!", safe_copy
    except Exception as e:
        print(f"Error processing uploaded image: {e}")
        return image, f"Image uploaded with warning: {e}", image

# Fix the image passing with better error handling
def safe_image_copy(img):
    """Safely copy an image to avoid corruption issues"""
    if img is None:
        return None
    try:
        # Convert to RGB and create a fresh copy
        if hasattr(img, 'mode'):
            # It's a PIL image
            rgb_img = img.convert('RGB')
            # Create a new image by copying pixel data
            new_img = Image.new('RGB', rgb_img.size)
            new_img.paste(rgb_img)
            return new_img
        else:
            print(f"Warning: Unexpected image type: {type(img)}")
            return img
    except Exception as e:
        print(f"Error copying image: {e}")
        return img

# Create the Gradio Interface
with gr.Blocks(title="AI Image Generator & Text Overlay") as demo:
    gr.Markdown("# 🎨 AI Image Generator & Text Overlay")
    gr.Markdown("Generate images with AI or upload your own, then add custom text overlays!")
    
    # Shared state for the current image
    current_image = gr.State(value=None)
    
    # Tab 1: Image Generation
    with gr.Tab("🖼️ Image Generation", id="generation_tab"):
        gr.Markdown("## Generate or Upload Images")
        
        with gr.Row():
            with gr.Column():
                # API Key inputs
                ideogram_key = gr.Textbox(
                    label="Ideogram API Key", 
                    type="password", 
                    placeholder="Enter your Ideogram API key"
                )
                dalle_key = gr.Textbox(
                    label="OpenAI API Key", 
                    type="password", 
                    placeholder="Enter your OpenAI API key"
                )
                
                # Prompt input
                prompt = gr.Textbox(
                    label="Image Prompt", 
                    placeholder="Describe the image you want to generate...",
                    lines=3
                )
                
                # Generation options
                with gr.Row():
                    ideogram_aspect = gr.Dropdown(
                        choices=["ASPECT_1_1", "ASPECT_16_9", "ASPECT_9_16", "ASPECT_4_3", "ASPECT_3_4"],
                        value="ASPECT_1_1",
                        label="Ideogram Aspect Ratio"
                    )
                    dalle_size = gr.Dropdown(
                        choices=["1024x1024", "1792x1024", "1024x1792"],
                        value="1024x1024",
                        label="DALL-E Size"
                    )
                
                # Generation buttons
                with gr.Row():
                    ideogram_btn = gr.Button("🎨 Generate with Ideogram", variant="primary")
                    dalle_btn = gr.Button("🤖 Generate with DALL-E 3", variant="primary")
                
                # Upload option
                gr.Markdown("### Or Upload Your Own Image")
                upload_image = gr.Image(label="Upload Image", type="pil")
                upload_btn = gr.Button("📁 Process Uploaded Image")
            
            with gr.Column():
                # Output
                generated_image = gr.Image(label="Generated/Uploaded Image", type="pil")
                status_message = gr.Textbox(label="Status", interactive=False)
        
        # Event handlers for Tab 1
        ideogram_btn.click(
            generate_ideogram_image,
            inputs=[ideogram_key, prompt, ideogram_aspect],
            outputs=[generated_image, status_message, current_image]
        )
        
        dalle_btn.click(
            generate_dalle_image,
            inputs=[dalle_key, prompt, dalle_size],
            outputs=[generated_image, status_message, current_image]
        )
        
        upload_btn.click(
            process_uploaded_image,
            inputs=[upload_image],
            outputs=[generated_image, status_message, current_image]
        )
    
    # Tab 2: Add Text Overlay
    with gr.Tab("📝 Add Text", id="text_tab"):
        gr.Markdown("## Overlay Text on Images")
        gr.Markdown("Choose a layout pattern, type your text, and render it on the current image.")
        
        with gr.Row():
            with gr.Column():
                # Text pattern selection with template preview
                with gr.Row():
                    with gr.Column(scale=2):
                        # FIXED: Use proper tuple format and matching value
                        pattern = gr.Dropdown(
                            choices=[
                                ("2-lines-top", "📄 2 Lines - Top (Header style)"),
                                ("2-lines-center", "📄 2 Lines - Center (Title style)"), 
                                ("2-lines-bottom", "📄 2 Lines - Bottom (Caption style)"),
                                ("3-lines-top", "📋 3 Lines - Top (Header + subtitle)"),
                                ("3-lines-center", "📋 3 Lines - Center (Full title)"),
                                ("3-lines-bottom", "📋 3 Lines - Bottom (Credits style)")
                            ],
                            value="2-lines-top",
                            label="Layout Pattern"
                        )
                    
                    with gr.Column(scale=1):
                        # Template preview image
                        template_preview = gr.Image(
                            label="Template Preview",
                            value=None,
                            show_label=True,
                            interactive=False,
                            width=200,
                            height=120
                        )
                
                # Text input fields
                line1_inp = gr.Textbox(
                    label="Line 1",
                    placeholder="Enter first line of text..."
                )
                line2_inp = gr.Textbox(
                    label="Line 2", 
                    placeholder="Enter second line of text..."
                )
                line3_inp = gr.Textbox(
                    label="Line 3",
                    placeholder="Enter third line of text..."
                )
                
                # Styling options
                with gr.Row():
                    font_size = gr.Slider(
                        minimum=12,
                        maximum=120,
                        value=48,
                        step=2,
                        label="Font Size"
                    )
                    color = gr.ColorPicker(
                        label="Text Color",
                        value="#FFFFFF"
                    )
                    
                # NEW: Font selector
                with gr.Row():
                    font_selector = gr.Dropdown(
                        choices=[
                            "Impact",           # Bold, condensed - great for memes
                            "Arial Bold",       # Clean, readable
                            "Times Bold",       # Classic, serif
                            "Comic Sans",       # Casual, friendly
                            "Bebas Neue"        # Modern, tall
                        ],
                        value="Impact",
                        label="Font Style",
                        info="Choose font for text overlay"
                    )
                
                # Add outline option
                add_outline = gr.Checkbox(
                    label="Add black outline",
                    value=True,
                    info="Adds black outline for better text visibility"
                )
                
                add_text_btn = gr.Button("✨ Add Text to Image", variant="primary")
            
            with gr.Column():
                # Output
                text_image = gr.Image(label="Image with Text Overlay")
        
        # Event handlers for Tab 2
        
        # Update template when pattern changes
        pattern.change(
            get_pattern_template,
            inputs=[pattern],
            outputs=[template_preview]
        )
        
        # Add text to image
        add_text_btn.click(
            add_text_to_image,
            inputs=[current_image, pattern, line1_inp, line2_inp, line3_inp, font_size, color, add_outline, font_selector],
            outputs=[text_image, status_message]
        )

    # FIXED: Move demo.load() INSIDE the Blocks context
    demo.load(
        lambda: get_pattern_template("2-lines-top"),
        outputs=[template_preview]
    )

# Launch the application - OUTSIDE the Blocks context
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860
    )