| import torch | |
| from diffusers import DiffusionPipeline | |
| import gradio as gr | |
| generator = DiffusionPipeline.from_pretrained("CompVis/ldm-text2im-large-256") | |
| # move to GPU if available | |
| if torch.cuda.is_available(): | |
| generator = generator.to("cuda") | |
| def generate(prompts): | |
| images = generator(list(prompts)).images | |
| return [images] | |
| demo = gr.Interface( | |
| generate, | |
| "textbox", | |
| "image", | |
| batch=True, | |
| max_batch_size=2, # Set the batch size based on your CPU/GPU memory | |
| ).queue() | |
| if __name__ == "__main__": | |
| demo.launch(share=True) | |