# 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()