app.py

# app.py - TextDiffuser-2 implementation with focus on layout planning
import os
import torch
import gradio as gr
import numpy as np
import json
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoTokenizer, AutoModelForCausalLM
from diffusers import StableDiffusionPipeline
import time
import random

# Try to import fastchat - may need to install with pip if not available
try:
    from fastchat.model import get_conversation_template
except ImportError:
    # Fallback implementation if fastchat is not available
    print("FastChat not found. Installing...")
    os.system("pip install fschat")
    from fastchat.model import get_conversation_template

# Check for GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Define global storage for user interactions
global_dict = {}

class TextDiffuserLayoutPlanner:
    """
    Implementation focused on the layout planning aspect of TextDiffuser-2
    """
    def __init__(self):
        # Load the layout planner model
        self.layout_model_path = "JingyeChen22/textdiffuser2_layout_planner"
        
        print(f"Loading layout planner model from {self.layout_model_path}...")
        
        try:
            # Initialize the tokenizer and model
            self.layout_tokenizer = AutoTokenizer.from_pretrained(
                self.layout_model_path, 
                use_fast=False
            )
            
            # Load the model with half precision if GPU is available
            model_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
            self.layout_model = AutoModelForCausalLM.from_pretrained(
                self.layout_model_path,
                torch_dtype=model_dtype,
                low_cpu_mem_usage=True
            ).to(device)
            
            print("Layout planner model loaded successfully")
        except Exception as e:
            print(f"Error loading layout planner: {e}")
            print("Falling back to simpler implementation...")
            # Set models to None to indicate fallback mode
            self.layout_model = None
            self.layout_tokenizer = None
        
        # Initialize a simple diffusion model for context visualization
        # This is optional and could be removed if you only need layout
        self.diffusion_model = None
        if torch.cuda.is_available():
            try:
                self.diffusion_model = StableDiffusionPipeline.from_pretrained(
                    "runwayml/stable-diffusion-v1-5",
                    torch_dtype=torch.float16
                ).to(device)
                print("Diffusion model loaded for context visualization")
            except Exception as e:
                print(f"Could not load diffusion model: {e}")
                print("Will use placeholder images instead")
    
    def generate_layout(self, prompt, keywords="", image_size=(512, 512), temperature=0.7):
        """
        Generate a text layout based on the prompt using the layout planner model
        
        Args:
            prompt: Description of the image to generate
            keywords: Optional keywords to include in the layout (format: "word1/word2/...")
            image_size: Size of the target image (width, height)
            temperature: Temperature for layout generation (higher = more diverse)
            
        Returns:
            layout_elements: List of text elements with positions
            layout_text: Raw output from the layout planner
            layout_image: Visualization of the layout
        """
        width, height = image_size
        
        # Only proceed with the layout planner if available
        if self.layout_model is not None and self.layout_tokenizer is not None:
            # Format the prompt for layout generation
            if len(keywords.strip()) == 0:
                template = f'Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is {width//4}x{height//4}. Therefore, all properties of the positions should not exceed {width//4}, including the coordinates of top, left, right, and bottom. All keywords are included in the caption. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {prompt}'
            else:
                keywords_list = keywords.split('/')
                keywords_list = [k.strip() for k in keywords_list]
                template = f'Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is {width//4}x{height//4}. Therefore, all properties of the positions should not exceed {width//4}, including the coordinates of top, left, right, and bottom. In addition, we also provide all keywords at random order for reference. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {prompt}. Keywords: {str(keywords_list)}'
            
            # Use FastChat's conversation template
            conv = get_conversation_template(self.layout_model_path)
            conv.append_message(conv.roles[0], template)
            conv.append_message(conv.roles[1], None)
            prompt_text = conv.get_prompt()
            
            # Generate the layout
            time_start = time.time()
            print(f"Generating layout for prompt: {prompt}")
            
            # Tokenize and prepare inputs
            inputs = self.layout_tokenizer([prompt_text], return_token_type_ids=False)
            inputs = {k: torch.tensor(v).to(device) for k, v in inputs.items()}
            
            # Generate layout with the model
            with torch.no_grad():
                output_ids = self.layout_model.generate(
                    **inputs,
                    do_sample=True,
                    temperature=temperature,
                    repetition_penalty=1.0,
                    max_new_tokens=512,
                )
            
            # Process the output
            if self.layout_model.config.is_encoder_decoder:
                output_ids = output_ids[0]
            else:
                output_ids = output_ids[0][len(inputs["input_ids"][0]):]
                
            layout_text = self.layout_tokenizer.decode(
                output_ids, skip_special_tokens=True, spaces_between_special_tokens=False
            )
            
            time_end = time.time()
            print(f"Layout generation took {time_end - time_start:.2f} seconds")
            print(f"Layout output: {layout_text}")
            
            # Parse the layout text to extract text elements
            layout_elements = self.parse_layout_text(layout_text, image_size)
            
            # Create a visualization of the layout
            layout_image = self.visualize_layout(layout_elements, image_size)
            
        else:
            # Fallback: Generate a simple layout
            print("Using fallback layout generation")
            layout_elements = self.generate_fallback_layout(prompt, keywords, image_size)
            layout_text = "Fallback layout generation - Layout planner model not available"
            layout_image = self.visualize_layout(layout_elements, image_size)
        
        return layout_elements, layout_text, layout_image
    
    def parse_layout_text(self, layout_text, image_size=(512, 512)):
        """
        Parse the layout text from the layout planner to extract text elements
        
        Args:
            layout_text: Output text from the layout planner
            image_size: Size of the target image
            
        Returns:
            layout_elements: List of text elements with positions
        """
        layout_elements = []
        lines = layout_text.strip().split('\n')
        
        for line in lines:
            line = line.strip()
            if not line or '###' in line or '.com' in line:
                continue
            
            try:
                # Parse the line to extract text and position
                parts = line.split()
                if len(parts) < 5:  # Need at least text and 4 coordinates
                    continue
                
                # Last 4 parts should be coordinates, everything else is text
                coords = parts[-1]
                text = ' '.join(parts[:-1])
                
                # Parse coordinates (left, top, right, bottom)
                try:
                    l, t, r, b = map(int, coords.split(','))
                    
                    # Scale coordinates to image size (they are given in 1/4 scale)
                    l, t, r, b = l*4, t*4, r*4, b*4
                    
                    # Create text element
                    element = {
                        "text": text,
                        "position": (l, t),
                        "size": (r-l, b-t),
                        "box": (l, t, r, b),
                        "style": {
                            "font": "Arial",
                            "size": 24,
                            "color": (0, 0, 0)
                        }
                    }
                    layout_elements.append(element)
                except ValueError:
                    print(f"Could not parse coordinates in line: {line}")
                    continue
                
            except Exception as e:
                print(f"Error parsing layout line: {e}")
                continue
        
        return layout_elements
    
    def generate_fallback_layout(self, prompt, keywords="", image_size=(512, 512)):
        """
        Generate a fallback layout when the layout planner is not available
        
        Args:
            prompt: Description of the image
            keywords: Optional keywords to include
            image_size: Size of the target image
            
        Returns:
            layout_elements: List of text elements with positions
        """
        width, height = image_size
        layout_elements = []
        
        # Extract keywords from the prompt or use provided keywords
        if keywords:
            keywords_list = keywords.split('/')
            keywords_list = [k.strip() for k in keywords_list]
        else:
            # Extract potential keywords from the prompt
            words = prompt.split()
            keywords_list = [word for word in words if len(word) > 3 and word.isalpha()]
            keywords_list = keywords_list[:3]  # Limit to 3 keywords
        
        # Generate positions for the keywords
        for i, keyword in enumerate(keywords_list):
            # Calculate a position based on the index
            row = i // 2
            col = i % 2
            
            l = 50 + col * (width // 2)
            t = 50 + row * (height // 3)
            r = l + 200
            b = t + 50
            
            element = {
                "text": keyword,
                "position": (l, t),
                "size": (r-l, b-t),
                "box": (l, t, r, b),
                "style": {
                    "font": "Arial",
                    "size": 24,
                    "color": (0, 0, 0)
                }
            }
            layout_elements.append(element)
        
        return layout_elements
    
    def visualize_layout(self, layout_elements, image_size=(512, 512)):
        """
        Create a visualization of the text layout
        
        Args:
            layout_elements: List of text elements with positions
            image_size: Size of the target image
            
        Returns:
            layout_image: Visualization of the layout
        """
        width, height = image_size
        image = Image.new("RGB", image_size, (240, 240, 240))
        draw = ImageDraw.Draw(image)
        
        # Draw grid lines
        for x in range(0, width, 32):
            alpha = 255 if x % 128 == 0 else 100
            draw.line([(x, 0), (x, height)], fill=(200, 200, 200, alpha), width=1)
        
        for y in range(0, height, 32):
            alpha = 255 if y % 128 == 0 else 100
            draw.line([(0, y), (width, y)], fill=(200, 200, 200, alpha), width=1)
        
        # Try to load a font
        try:
            font_large = ImageFont.truetype("Arial.ttf", 20)
            font_small = ImageFont.truetype("Arial.ttf", 12)
        except IOError:
            try:
                font_large = ImageFont.truetype("DejaVuSans.ttf", 20)
                font_small = ImageFont.truetype("DejaVuSans.ttf", 12)
            except IOError:
                font_large = ImageFont.load_default()
                font_small = ImageFont.load_default()
        
        # Draw text elements
        for i, element in enumerate(layout_elements):
            box = element.get("box", (0, 0, 0, 0))
            text = element["text"]
            
            # Draw bounding box
            draw.rectangle(box, outline=(255, 0, 0), width=2)
            
            # Draw text label
            draw.text(
                (box[0] + 5, box[1] - 20), 
                f"{i+1}: {text}", 
                font=font_small, 
                fill=(0, 0, 0)
            )
            
            # Draw coordinates
            coord_text = f"({box[0]},{box[1]}) to ({box[2]},{box[3]})"
            draw.text(
                (box[0] + 5, box[3] + 5), 
                coord_text, 
                font=font_small, 
                fill=(0, 0, 255)
            )
        
        return image
    
    def generate_context_image(self, prompt, image_size=(512, 512)):
        """
        Generate a context image based on the prompt
        
        Args:
            prompt: Description of the image
            image_size: Size of the target image
            
        Returns:
            image: Generated or placeholder image
        """
        if self.diffusion_model is not None:
            # Generate image using the diffusion model
            try:
                images = self.diffusion_model(
                    prompt=prompt,
                    height=image_size[1],
                    width=image_size[0],
                    num_inference_steps=20
                ).images
                return images[0]
            except Exception as e:
                print(f"Error generating image: {e}")
                print("Using placeholder image instead")
        
        # Create a placeholder gradient image
        width, height = image_size
        image = Image.new("RGB", image_size, (240, 240, 240))
        
        # Add a subtle gradient background
        for y in range(height):
            for x in range(width):
                r = int(240 - 30 * (y / height))
                g = int(240 - 20 * (x / width))
                b = int(240 - 40 * ((x + y) / (width + height)))
                image.putpixel((x, y), (r, g, b))
        
        return image
    
    def process_request(self, prompt, keywords="", width=512, height=512, temperature=0.7, generate_image=False):
        """
        Process a user request to generate a layout
        
        Args:
            prompt: Description of the image
            keywords: Optional keywords to include
            width: Width of the target image
            height: Height of the target image
            temperature: Temperature for layout generation
            generate_image: Whether to generate a context image
            
        Returns:
            layout_elements: List of text elements with positions
            layout_text: Raw output from the layout planner
            layout_image: Visualization of the layout
            context_image: Generated or placeholder image (if requested)
        """
        image_size = (width, height)
        
        # Generate layout
        layout_elements, layout_text, layout_image = self.generate_layout(
            prompt, keywords, image_size, temperature
        )
        
        # Generate context image if requested
        context_image = None
        if generate_image:
            context_image = self.generate_context_image(prompt, image_size)
        
        # Format the layout data for display
        layout_data = {
            "prompt": prompt,
            "keywords": keywords,
            "image_size": image_size,
            "text_elements": layout_elements,
        }
        
        return layout_elements, layout_text, layout_image, context_image, layout_data

# Initialize the model
model = TextDiffuserLayoutPlanner()

# Create the Gradio interface
with gr.Blocks(title="TextDiffuser-2 Layout Planner") as demo:
    gr.Markdown("""
    # TextDiffuser-2 Layout Planner
    
    This application focuses on the layout planning aspect of TextDiffuser-2. It allows you to:
    
    1. Generate text layouts for images based on prompts
    2. Visualize the layout with text positions and bounding boxes
    3. Export the layout information for use in your own HTML5 Canvas UI editor
    
    Based on the paper "[TextDiffuser-2: Unleashing the Power of Language Models for Text Rendering](https://arxiv.org/abs/2311.16465)" by Jingye Chen et al.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(
                label="Prompt",
                value="A beautiful city skyline stamp of Shanghai",
                lines=3,
                placeholder="Describe the image you want to generate with text elements"
            )
            
            keywords_input = gr.Textbox(
                label="Optional Keywords (separated by /)",
                placeholder="keyword1/keyword2/keyword3",
                info="If provided, the layout planner will try to use these keywords"
            )
            
            with gr.Row():
                width_input = gr.Number(label="Width", value=512, minimum=256, maximum=1024, step=64)
                height_input = gr.Number(label="Height", value=512, minimum=256, maximum=1024, step=64)
            
            temperature_input = gr.Slider(
                label="Temperature",
                minimum=0.1,
                maximum=2.0,
                value=0.7,
                step=0.1,
                info="Controls randomness in layout generation. Higher values produce more diverse layouts."
            )
            
            show_image_input = gr.Checkbox(
                label="Generate Context Image",
                value=False,
                info="Generate a simple image to provide context (this is just for visualization)"
            )
            
            generate_button = gr.Button("Generate Layout", variant="primary")
            
            gr.Markdown("""
            ## Tips for using this demo
            
            1. The layout planner works best with descriptive prompts
            2. You can specify keywords to ensure they appear in the layout
            3. Increase temperature for more diverse layouts
            4. The context image is optional and just for visualization
            """)
        
        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.TabItem("Layout Visualization"):
                    layout_output = gr.Image(label="Text Layout Visualization", type="pil")
                
                with gr.TabItem("Context Image"):
                    context_image_output = gr.Image(label="Context Image (Optional)", type="pil")
                
                with gr.TabItem("Layout Information"):
                    layout_elements_output = gr.JSON(label="Layout Elements")
                
                with gr.TabItem("Raw Layout Output"):
                    layout_text_output = gr.Textbox(label="Raw Layout Planner Output", lines=10)
    
    # Examples
    gr.Examples(
        examples=[
            ["A new year greeting card of happy 2024, surrounded by balloons", "", 512, 512, 0.7, True],
            ["A beautiful city skyline stamp of Shanghai", "", 512, 512, 0.7, True],
            ["The words 'KFC VIVO50' are inscribed upon the wall in a neon light effect", "KFC/VIVO50", 512, 512, 0.7, True],
            ["A logo of superman", "", 512, 512, 0.7, True],
            ["A pencil sketch of a tree with the title nothing to tree here", "nothing/tree/here", 512, 512, 0.7, True],
            ["Delicate greeting card of happy birthday to xyz", "happy/birthday/xyz", 768, 512, 1.0, True],
            ["Book cover of good morning baby", "good/morning/baby", 512, 768, 0.7, True],
        ],
        inputs=[prompt_input, keywords_input, width_input, height_input, temperature_input, show_image_input]
    )
    
    # Function to process the request
    def process_ui_request(prompt, keywords, width, height, temperature, show_image):
        try:
            width = int(width)
            height = int(height)
            
            layout_elements, layout_text, layout_image, context_image, layout_data = model.process_request(
                prompt, 
                keywords, 
                width, 
                height, 
                temperature,
                show_image
            )
            
            if show_image and context_image is not None:
                return layout_image, context_image, layout_data, layout_text
            else:
                return layout_image, None, layout_data, layout_text
        
        except Exception as e:
            error_message = f"Error: {str(e)}"
            print(error_message)
            return None, None, {"error": error_message}, error_message
    
    # Connect the button to the processing function
    generate_button.click(
        fn=process_ui_request,
        inputs=[prompt_input, keywords_input, width_input, height_input, temperature_input, show_image_input],
        outputs=[layout_output, context_image_output, layout_elements_output, layout_text_output]
    )
    
    gr.Markdown("""
    ## About TextDiffuser-2
    
    TextDiffuser-2 is a system that uses language models for text rendering in images. The layout planner component is responsible for determining where text should be positioned in the generated image.
    
    This demo focuses only on the layout planning aspect, allowing you to generate and export layout information that can be used in your own HTML5 Canvas UI editor.
    
    For the full TextDiffuser-2 implementation, please visit the [official repository](https://github.com/microsoft/unilm/tree/master/textdiffuser-2).
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch()