File size: 3,169 Bytes
2740cfb 68a1920 2740cfb 68a1920 2740cfb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import torch
from diffusers import AutoPipelineForImage2Image
from PIL import Image, ImageDraw, ImageFont
import requests
from io import BytesIO
import gradio as gr
# log gpu availabilitu
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):
pipeline = AutoPipelineForImage2Image.from_pretrained(
"./models/kandinsky-2-2-decoder",
torch_dtype=torch.float16,
variant="fp16",
use_safetensors=True
)
# pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)
pipeline.enable_model_cpu_offload()
prompt = f"{prompt}, include the color {image_color}"
negative_prompt = "low quality, bad quality, blurry, unprofessional"
generated_image = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
image=initial_image,
height=256,
width=256).images[0]
template_image = image_to_template(generated_image, logo, button_text, punchline, theme_color)
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")],
title="Ad Template Generation Using Diffusion Models Demo",
description="Generate ad template based on your inputs using a trained model.",
concurrency_limit=2,
# examples=[
# []
# ]
)
# Run the interface
iface.launch()
|