app.py

# app.py - TextDiffuser-2 implementation for Hugging Face Spaces
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

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

class SimpleTextDiffuser:
    """
    Simple implementation of TextDiffuser-2 concept for Hugging Face Spaces
    """
    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)
        
        # 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
            )
            self.diffusion_model.to(device)
        
        print("Models initialized")
    
    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:
        """
        
        # 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)
            )
            
            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"""
        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]
        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))
        
        return 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)
        
        for element in text_elements:
            try:
                font_size = element["size"]
                
                # 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()
                
                # 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)
                
            except Exception as e:
                print(f"Error rendering text: {e}")
                continue
        
        return image_with_text
    
    def visualize_layout(self, text_elements, image_size=(512, 512)):
        """Create a visualization of the text layout"""
        width, height = image_size
        image = Image.new("RGB", image_size, (255, 255, 255))
        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 text elements
        for element in text_elements:
            position = element["position"]
            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 text label
            try:
                font = ImageFont.truetype("DejaVuSans.ttf", 12)
            except IOError:
                font = ImageFont.load_default()
            
            # 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)
        
        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))
        
        image_size = (width, height)
        
        # Step 1: Generate text layout
        text_elements, layout_text = self.generate_layout(prompt, image_size, num_text_elements)
        
        # Step 2: Generate base image
        base_image = self.generate_image(prompt, image_size)
        
        # Step 3: Render text onto the image
        image_with_text = self.render_text(base_image, text_elements)
        
        # Step 4: Create layout visualization
        layout_visualization = self.visualize_layout(text_elements, image_size)
        
        # Step 5: Format layout information for display
        layout_info = {
            "prompt": prompt,
            "image_size": image_size,
            "num_text_elements": num_text_elements,
            "text_elements": text_elements,
            "layout_generation_prompt": layout_text
        }
        
        formatted_layout = json.dumps(layout_info, indent=2)
        
        return image_with_text, layout_visualization, formatted_layout

# 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

# Create the Gradio app
with gr.Blocks(title="TextDiffuser-2 Demo") as demo:
    gr.Markdown("""
    # TextDiffuser-2 Demo
    
    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.
    
    Generate images with text by providing a descriptive prompt below.
    """)
    
    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
            )
            
            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
            )
            
            submit_button = gr.Button("Generate Image", variant="primary")
        
        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")
                
                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:
    """)
    
    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],
        ],
        inputs=[prompt_input, width_input, height_input, num_elements_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]
    )
    
    gr.Markdown("""
    ## About
    
    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.
    
    Running on: """ + str(device))

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