colab working script

app.py

@@ -4,46 +4,131 @@ from PIL import Image, ImageDraw, ImageFont
 import requests
 from io import BytesIO
 import gradio as gr
+import gc
+import textwrap
 
-# log gpu availabilitu
+# log gpu availability
 print(f"Is CUDA available: {torch.cuda.is_available()}")
 print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 
 
 def image_to_template(generated_image, logo, button_text, punchline, theme_color):
-    # Resize logo if needed
-    logo = logo.resize((100, 100))  # Example size, adjust as needed
-
-    # Create a blank canvas with extra space for logo, punchline, and button
-    canvas_width = max(generated_image.width, logo.width) * 2
-    canvas_height = generated_image.height + logo.height + 100
-    canvas = Image.new('RGB', (canvas_width, canvas_height), 'white')
-    
-    # Paste the logo and the generated image onto the canvas
-    canvas.paste(logo, (10, 10))  # Adjust position as needed
-    canvas.paste(generated_image, (0, logo.height + 20))
-
-    # Add punchline and button
-    draw = ImageDraw.Draw(canvas)
-    font = ImageFont.load_default()  # Or use a custom font
-    text_color = theme_color
-
-    # Draw punchline
-    draw.text((10, logo.height + generated_image.height + 30), punchline, fill=text_color, font=font)
-
-    # Draw button
-    button_position = (10, logo.height + generated_image.height + 60)  # Adjust as needed
-    draw.rectangle([button_position, (canvas_width - 10, canvas_height - 10)], outline=theme_color, fill=text_color)
-    draw.text(button_position, button_text, font=font)
-
-    return canvas
-
-def generate_template(initial_image, logo, prompt, button_text, punchline, image_color, theme_color):
+    template_width = 540
+    button_font_size = 10
+    punchline_font_size = 30
+    decoration_height = 10
+    margin = 20
+    # wrap punchline text
+    punchline = textwrap.wrap(punchline, width=35)
+    n_of_lines_punchline = len(punchline)
+
+    generated_image = generated_image.convert("RGBA")
+    logo = logo.convert("RGBA")
+
+    # image shape
+    image_width = template_width // 2
+    image_height = image_width * generated_image.height // generated_image.width
+    image_shape = (image_width, image_height)
+
+    # logo shape
+    logo_width = image_width // 3
+    logo_height = logo_width * logo.height // logo.width
+    logo_shape = (logo_width, logo_height)
+
+    # Define fonts
+    button_font = ImageFont.truetype("Montserrat-Bold.ttf", button_font_size)
+    punchline_font = ImageFont.truetype("Montserrat-Bold.ttf", punchline_font_size)
+
+    # button shape
+    button_width = template_width // 3
+    button_height = button_font_size * 3
+
+    # template height calculation
+    template_height = (
+        image_height
+        + logo_height
+        + button_height
+        + n_of_lines_punchline * punchline_font_size
+        + (5 * margin)
+        + (2 * decoration_height)
+    )
+
+    # Calculate positions for the centered layout
+    logo_pos = ((template_width - logo_width) // 2, margin + decoration_height)
+    image_pos = (
+        (template_width - image_width) // 2,
+        logo_pos[1] + logo_height + margin,
+    )
+
+    # Decoration positions
+    top_decoration_pos = [
+        margin,
+        -decoration_height // 2,
+        template_width - margin,
+        decoration_height // 2,
+    ]
+    bottom_decoration_pos = [
+        margin,
+        template_height - decoration_height // 2,
+        template_width - margin,
+        template_height + decoration_height // 2,
+    ]
+
+    # Generate Components
+    generated_image.thumbnail(image_shape, Image.ANTIALIAS)
+    logo.thumbnail(logo_shape, Image.ANTIALIAS)
+    background = Image.new("RGBA", (template_width, template_height), "WHITE")
+    # round the corners of generated image
+    mask = Image.new("L", generated_image.size, 0)
+    draw = ImageDraw.Draw(mask)
+    draw.rounded_rectangle((0, 0) + generated_image.size, 20, fill=255)
+    generated_image.putalpha(mask)
+    # Paste the logo and the generated image onto the background
+    background.paste(logo, logo_pos, logo)
+    background.paste(generated_image, image_pos, generated_image)
+    # Draw the decorations, punchline, and button
+    draw = ImageDraw.Draw(background)
+    # Decorations on top and bottom
+    draw.rounded_rectangle(bottom_decoration_pos, radius=20, fill=theme_color)
+    draw.rounded_rectangle(top_decoration_pos, radius=20, fill=theme_color)
+    # Punchline text
+    text_heights = []
+    for line in punchline:
+        text_width, text_height = draw.textsize(line, font=punchline_font)
+        punchline_pos = (
+            (template_width - text_width) // 2,
+            image_pos[1] + generated_image.height + margin + sum(text_heights),
+        )
+        draw.text(punchline_pos, line, fill=theme_color, font=punchline_font)
+        text_heights.append(text_height)
+
+    # Button with rounded corners
+    button_text_width, button_text_height = draw.textsize(button_text, font=button_font)
+    button_shape = [
+        ((template_width - button_width) // 2, punchline_pos[1] + text_height + margin),
+        (
+            (template_width + button_width) // 2,
+            punchline_pos[1] + text_height + margin + button_height,
+        ),
+    ]
+    draw.rounded_rectangle(button_shape, radius=20, fill=theme_color)
+    # Button text
+    button_text_pos = (
+        (template_width - button_text_width) // 2,
+        button_shape[0][1] + (button_height - button_text_height) // 2,
+    )
+    draw.text(button_text_pos, button_text, fill="white", font=button_font)
+
+    return background
+
+
+def generate_template(
+    initial_image, logo, prompt, button_text, punchline, image_color, theme_color
+):
     pipeline = AutoPipelineForImage2Image.from_pretrained(
-        "./models/kandinsky-2-2-decoder", 
-        torch_dtype=torch.float16, 
-        variant="fp16", 
-        use_safetensors=True
+        "./models/kandinsky-2-2-decoder",
+        torch_dtype=torch.float16,
+        use_safetensors=True,
     )
 
     # pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)
@@ -53,34 +138,53 @@ def generate_template(initial_image, logo, prompt, button_text, punchline, image
     negative_prompt = "low quality, bad quality, blurry, unprofessional"
 
     generated_image = pipeline(
-        prompt=prompt, 
+        prompt=prompt,
         negative_prompt=negative_prompt,
-        image=initial_image,  
-        height=256, 
-        width=256).images[0]
+        image=initial_image,
+        height=256,
+        width=256,
+    ).images[0]
 
-    template_image = image_to_template(generated_image, logo, button_text, punchline, theme_color)
+    template_image = image_to_template(
+        generated_image, logo, button_text, punchline, theme_color
+    )
+
+    # free cpu and gpu memory
+    del pipeline
+    gc.collect()
+    torch.cuda.empty_cache()
 
     return template_image
 
+
 # Set up Gradio interface
 iface = gr.Interface(
     fn=generate_template,
-    inputs=[gr.inputs.Image(type="pil", label="Initial Image"), 
-            gr.inputs.Image(type="pil", label="Logo"), 
-            gr.inputs.Textbox(label="Prompt"), 
-            gr.inputs.Textbox(label="Button Text"),
-            gr.inputs.Textbox(label="Punchline"),
-            gr.inputs.ColorPicker(label="Image Color"),
-            gr.inputs.ColorPicker(label="Theme Color")],
-    outputs=[gr.outputs.Image(type="pil")],
+    inputs=[
+        gr.Image(type="pil", label="Initial Image"),
+        gr.Image(type="pil", label="Logo"),
+        gr.Textbox(label="Prompt"),
+        gr.Textbox(label="Button Text"),
+        gr.Textbox(label="Punchline"),
+        gr.ColorPicker(label="Image Color"),
+        gr.ColorPicker(label="Theme Color"),
+    ],
+    outputs=[gr.Image(type="pil")],
     title="Ad Template Generation Using Diffusion Models Demo",
     description="Generate ad template based on your inputs using a trained model.",
     concurrency_limit=2,
-    # examples=[
-    #     []
-    # ]
+    examples=[
+        [
+            "/path/to/example_initial_image1.jpg",  # Initial Image
+            "/path/to/example_logo1.png",  # Logo
+            "A scenic mountain landscape",  # Prompt
+            "Discover More",  # Button Text
+            "Escape into Nature",  # Punchline
+            "#00FF00",  # Image Color
+            "#0000FF",  # Theme Color
+        ]
+    ],
 )
 
 # Run the interface
-iface.launch()
+iface.launch(debug=True)

requirements.txt

@@ -3,6 +3,6 @@ transformers
 --extra-index-url https://download.pytorch.org/whl/cu113
 torch
 gradio
-PIL
+pillow
 requests
 safetensors