Spaces:
Runtime error
Runtime error
File size: 2,592 Bytes
9f1dfb4 0bf6f40 9f1dfb4 0bf6f40 9f1dfb4 0bf6f40 9f1dfb4 |
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 |
import torch
import gradio as gr
from diffusers import DiffusionPipeline
from PIL import Image, ImageDraw, ImageFont
# Initialize the model globally for faster inference
pipe = None
def load_model():
global pipe
if pipe is None:
pipe = DiffusionPipeline.from_pretrained("segmind/tiny-sd", torch_dtype=torch.float32)
pipe = pipe.to("cpu")
return pipe
def generate_image_with_text(prompt, caption_text):
"""
Generate an image based on a text prompt and add simple text with border at the bottom
"""
# Ensure model is loaded
pipe = load_model()
# Generate image
image = pipe(prompt, num_inference_steps=20, height=384, width=384).images[0]
# Add space for text at bottom
width, height = image.size
new_img = Image.new("RGB", (width, height + 40), (0, 0, 0))
new_img.paste(image, (0, 0))
# Add text with border
draw = ImageDraw.Draw(new_img)
# Use default font
font = ImageFont.load_default()
# Calculate text position to center it
text_width = draw.textlength(caption_text, font=font)
text_x = (width - text_width) / 2
text_y = height + 10
# Draw text border (offset in 4 directions)
offset = 1 # Border thickness
for dx, dy in [(-offset, -offset), (-offset, offset), (offset, -offset), (offset, offset),
(0, -offset), (0, offset), (-offset, 0), (offset, 0)]:
draw.text((text_x + dx, text_y + dy), caption_text, fill=(0, 0, 0), font=font)
# Draw main text in white
draw.text((text_x, text_y), caption_text, fill=(255, 255, 255), font=font)
return new_img
# Create Gradio interface
title = "Text-to-Image Generator with Caption"
description = """
Generate an image from a text prompt and add a caption at the bottom.
The model used is lightweight and runs on CPU.
"""
demo = gr.Interface(
fn=generate_image_with_text,
inputs=[
gr.Textbox(label="Image Prompt", placeholder="Describe the image you want to generate..."),
gr.Textbox(label="Caption Text", placeholder="Text to display at the bottom of the image")
],
outputs=gr.Image(type="pil", label="Generated Image"),
title=title,
description=description,
examples=[
["A serene mountain landscape at sunset", "Beautiful Sunset View"],
["A cute cat playing with yarn", "Playful Kitten"],
["Abstract colorful shapes", "Modern Art"]
],
cache_examples=False,
)
# Load the model when the app starts
load_model()
# Launch the app
if __name__ == "__main__":
demo.launch() |