Spaces:
Runtime error
Runtime error
| 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() |