File size: 628 Bytes
0e7cefd 7f5f38c 0e7cefd 7f5f38c 0e7cefd e88d83b d1e4372 e88d83b 0e7cefd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | import torch
from diffusers import DiffusionPipeline # type: ignore
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 # type: ignore
return [images]
demo = gr.Interface(generate,
"textbox",
"image",
batch=True,
max_batch_size=4, # Set the batch size based on your CPU/GPU memory
api_name="predict"
)
if __name__ == "__main__":
demo.launch()
|