Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py - TextDiffuser-2 implementation for Hugging Face Spaces
+# app.py - TextDiffuser-2 implementation with focus on layout planning
 import os
 import torch
 import gradio as gr
@@ -7,357 +7,535 @@ 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}")
 
-class SimpleTextDiffuser:
+# Define global storage for user interactions
+global_dict = {}
+
+class TextDiffuserLayoutPlanner:
     """
-    Simple implementation of TextDiffuser-2 concept for Hugging Face Spaces
+    Implementation focused on the layout planning aspect of TextDiffuser-2
     """
     def __init__(self):
-        # Load language model for layout generation
-        # Using a small model for efficiency
-        self.tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
-        self.language_model = AutoModelForCausalLM.from_pretrained("distilgpt2")
-        self.language_model.to(device)
+        # Load the layout planner model
+        self.layout_model_path = "JingyeChen22/textdiffuser2_layout_planner"
         
-        # Only load the diffusion model if we have a GPU
-        self.diffusion_model = None
-        if torch.cuda.is_available():
-            self.diffusion_model = StableDiffusionPipeline.from_pretrained(
-                "runwayml/stable-diffusion-v1-5",
-                torch_dtype=torch.float16
+        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
             )
-            self.diffusion_model.to(device)
+            
+            # 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
         
-        print("Models initialized")
+        # 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, image_size=(512, 512), num_text_elements=3):
-        """Generate text layout based on prompt"""
-        width, height = image_size
-        
-        # Format the prompt for layout generation
-        layout_prompt = f"""
-        Create a layout for an image with:
-        - Description: {prompt}
-        - Image size: {width}x{height}
-        - Number of text elements: {num_text_elements}
-        
-        Generate text content and positions:
+    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
         
-        # Generate layout using LM
-        input_ids = self.tokenizer.encode(layout_prompt, return_tensors="pt").to(device)
-        with torch.no_grad():
-            output = self.language_model.generate(
-                input_ids,
-                max_length=input_ids.shape[1] + 150,
-                temperature=0.7,
-                num_return_sequences=1,
-                pad_token_id=self.tokenizer.eos_token_id
-            )
-        
-        layout_text = self.tokenizer.decode(output[0], skip_special_tokens=True)
-        
-        # Parse the generated layout (simplified)
-        # In a real implementation, this would be more sophisticated
-        text_elements = []
-        
-        # Simple fallback: generate random layout
-        import random
-        
-        # Create a title element
-        title = prompt.split()[:5]
-        title = " ".join(title) + "..."
-        title_x = width // 4
-        title_y = height // 4
-        text_elements.append({
-            "text": title,
-            "position": (title_x, title_y),
-            "size": 24,
-            "color": (0, 0, 0),
-            "type": "title"
-        })
-        
-        # Create additional text elements
-        sample_texts = [
-            "Premium Quality",
-            "Best Value",
-            "Limited Edition",
-            "New Collection",
-            "Special Offer",
-            "Coming Soon",
-            "Best Seller",
-            "Top Choice",
-            "Featured Product",
-            "Exclusive Deal"
-        ]
-        
-        for i in range(1, num_text_elements):
-            x = random.randint(width // 8, width * 3 // 4)
-            y = random.randint(height // 3, height * 3 // 4)
-            text = sample_texts[i % len(sample_texts)]
-            color = (
-                random.randint(0, 200),
-                random.randint(0, 200),
-                random.randint(0, 200)
-            )
+        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)
             
-            text_elements.append({
-                "text": text,
-                "position": (x, y),
-                "size": 18,
-                "color": color,
-                "type": f"element_{i}"
-            })
-        
-        return text_elements, layout_text
-    
-    def generate_image(self, prompt, image_size=(512, 512)):
-        """Generate base image using diffusion model or placeholder"""
+        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
         
-        if self.diffusion_model and torch.cuda.is_available():
-            # Generate image using diffusion model
-            image = self.diffusion_model(
-                prompt=prompt,
-                height=height,
-                width=width,
-                num_inference_steps=30
-            ).images[0]
+        # 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:
-            # Create a placeholder gradient image
-            image = Image.new("RGB", image_size, (240, 240, 240))
-            
-            # Add a colored gradient background
-            for y in range(height):
-                for x in range(width):
-                    r = int(240 - 100 * (y / height))
-                    g = int(240 - 50 * (x / width))
-                    b = int(240 - 75 * ((x + y) / (width + height)))
-                    image.putpixel((x, y), (r, g, b))
+            # 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 image
+        return layout_elements, layout_text, layout_image
     
-    def render_text(self, image, text_elements):
-        """Render text elements onto the image"""
-        image_with_text = image.copy()
-        draw = ImageDraw.Draw(image_with_text)
+    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 element in text_elements:
+        for line in lines:
+            line = line.strip()
+            if not line or '###' in line or '.com' in line:
+                continue
+            
             try:
-                font_size = element["size"]
+                # Parse the line to extract text and position
+                parts = line.split()
+                if len(parts) < 5:  # Need at least text and 4 coordinates
+                    continue
                 
-                # Try to load a font, fall back to default if not available
-                try:
-                    font = ImageFont.truetype("DejaVuSans.ttf", font_size)
-                except IOError:
-                    try:
-                        font = ImageFont.truetype("Arial.ttf", font_size)
-                    except IOError:
-                        font = ImageFont.load_default()
+                # Last 4 parts should be coordinates, everything else is text
+                coords = parts[-1]
+                text = ' '.join(parts[:-1])
                 
-                # Draw text with background for better visibility
-                text = element["text"]
-                position = element["position"]
-                color = element["color"]
-                
-                # Get text size to create background
-                bbox = draw.textbbox(position, text, font=font)
-                text_width = bbox[2] - bbox[0]
-                text_height = bbox[3] - bbox[1]
-                
-                # Draw semi-transparent background
-                padding = 5
-                background_box = [
-                    position[0] - padding,
-                    position[1] - padding,
-                    position[0] + text_width + padding,
-                    position[1] + text_height + padding
-                ]
-                draw.rectangle(background_box, fill=(255, 255, 255, 200))
-                
-                # Draw text
-                draw.text(position, text, fill=color, font=font)
+                # 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 rendering text: {e}")
+                print(f"Error parsing layout line: {e}")
                 continue
         
-        return image_with_text
+        return layout_elements
     
-    def visualize_layout(self, text_elements, image_size=(512, 512)):
-        """Create a visualization of the text layout"""
+    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
-        image = Image.new("RGB", image_size, (255, 255, 255))
+        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
-        for x in range(0, width, 50):
-            draw.line([(x, 0), (x, height)], fill=(230, 230, 230))
-        for y in range(0, height, 50):
-            draw.line([(0, y), (width, y)], fill=(230, 230, 230))
+        # 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 element in text_elements:
-            position = element["position"]
+        for i, element in enumerate(layout_elements):
+            box = element.get("box", (0, 0, 0, 0))
             text = element["text"]
-            element_type = element.get("type", "unknown")
-            
-            # Draw position marker
-            circle_radius = 5
-            circle_bbox = [
-                position[0] - circle_radius,
-                position[1] - circle_radius,
-                position[0] + circle_radius,
-                position[1] + circle_radius
-            ]
-            draw.ellipse(circle_bbox, fill=(255, 0, 0))
+            
+            # Draw bounding box
+            draw.rectangle(box, outline=(255, 0, 0), width=2)
             
             # Draw text label
-            try:
-                font = ImageFont.truetype("DejaVuSans.ttf", 12)
-            except IOError:
-                font = ImageFont.load_default()
+            draw.text(
+                (box[0] + 5, box[1] - 20), 
+                f"{i+1}: {text}", 
+                font=font_small, 
+                fill=(0, 0, 0)
+            )
             
-            # Draw text preview and position info
-            info_text = f"{text} ({element_type})"
-            pos_text = f"Position: ({position[0]}, {position[1]})"
-            draw.text((position[0] + 10, position[1]), info_text, fill=(0, 0, 0), font=font)
-            draw.text((position[0] + 10, position[1] + 15), pos_text, fill=(0, 0, 255), font=font)
+            # 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_text_image(self, prompt, width=512, height=512, num_text_elements=3):
-        """Generate an image with rendered text based on prompt"""
-        # Validate inputs
-        width = max(256, min(1024, width))
-        height = max(256, min(1024, height))
-        num_text_elements = max(1, min(5, num_text_elements))
+    def generate_context_image(self, prompt, image_size=(512, 512)):
+        """
+        Generate a context image based on the prompt
         
-        image_size = (width, height)
+        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")
         
-        # Step 1: Generate text layout
-        text_elements, layout_text = self.generate_layout(prompt, image_size, num_text_elements)
+        # Create a placeholder gradient image
+        width, height = image_size
+        image = Image.new("RGB", image_size, (240, 240, 240))
         
-        # Step 2: Generate base image
-        base_image = self.generate_image(prompt, image_size)
+        # 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))
         
-        # Step 3: Render text onto the image
-        image_with_text = self.render_text(base_image, text_elements)
+        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
         
-        # Step 4: Create layout visualization
-        layout_visualization = self.visualize_layout(text_elements, image_size)
+        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)
         
-        # Step 5: Format layout information for display
-        layout_info = {
+        # 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,
-            "num_text_elements": num_text_elements,
-            "text_elements": text_elements,
-            "layout_generation_prompt": layout_text
+            "text_elements": layout_elements,
         }
         
-        formatted_layout = json.dumps(layout_info, indent=2)
-        
-        return image_with_text, layout_visualization, formatted_layout
+        return layout_elements, layout_text, layout_image, context_image, layout_data
 
 # Initialize the model
-model = SimpleTextDiffuser()
-
-# Define the Gradio interface
-def process_request(prompt, width, height, num_text_elements):
-    try:
-        width = int(width)
-        height = int(height)
-        num_text_elements = int(num_text_elements)
-        
-        image, layout, layout_info = model.generate_text_image(
-            prompt,
-            width=width,
-            height=height,
-            num_text_elements=num_text_elements
-        )
-        
-        return image, layout, layout_info
-    except Exception as e:
-        error_message = f"Error: {str(e)}"
-        print(error_message)
-        return None, None, error_message
+model = TextDiffuserLayoutPlanner()
 
-# Create the Gradio app
-with gr.Blocks(title="TextDiffuser-2 Demo") as demo:
+# Create the Gradio interface
+with gr.Blocks(title="TextDiffuser-2 Layout Planner") as demo:
     gr.Markdown("""
-    # TextDiffuser-2 Demo
+    # TextDiffuser-2 Layout Planner
     
-    This demo implements the concepts from the paper "[TextDiffuser-2: Unleashing the Power of Language Models for Text Rendering](https://arxiv.org/abs/2311.16465)" by Jingye Chen et al.
+    This application focuses on the layout planning aspect of TextDiffuser-2. It allows you to:
     
-    Generate images with text by providing a descriptive prompt below.
+    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 modern business poster with company name and tagline",
-                lines=3
+                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)
             
-            num_elements_input = gr.Slider(
-                label="Number of Text Elements",
-                minimum=1,
-                maximum=5,
-                value=3,
-                step=1
+            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."
             )
             
-            submit_button = gr.Button("Generate Image", variant="primary")
+            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("Generated Image"):
-                    image_output = gr.Image(label="Image with Text")
-                
                 with gr.TabItem("Layout Visualization"):
-                    layout_output = gr.Image(label="Text Layout")
+                    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_info_output = gr.Code(language="json", label="Layout Data")
-    
-    gr.Markdown("""
-    ## Example Prompts
-    
-    Try these prompts or create your own:
-    """)
+                    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
+    gr.Examples(
         examples=[
-            ["A movie poster for a sci-fi thriller", 512, 768, 3],
-            ["A motivational quote on a sunset background", 768, 512, 2],
-            ["A coffee shop menu with prices", 512, 512, 4],
-            ["A modern business card design", 512, 384, 3],
+            ["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, width_input, height_input, num_elements_input]
+        inputs=[prompt_input, keywords_input, width_input, height_input, temperature_input, show_image_input]
     )
     
-    submit_button.click(
-        fn=process_request,
-        inputs=[prompt_input, width_input, height_input, num_elements_input],
-        outputs=[image_output, layout_output, layout_info_output]
+    # 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
+    ## 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 is a simplified implementation for demonstration purposes. The full approach described in the paper involves deeper integration of language models with the diffusion process.
+    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.
     
-    Running on: """ + str(device))
+    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__":