mirror of
https://github.com/gradio-app/gradio.git
synced 2025-01-12 10:34:32 +08:00
23 lines
575 B
Python
23 lines
575 B
Python
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=4 # Set the batch size based on your CPU/GPU memory
|
|
).queue()
|
|
|
|
if __name__ == "__main__":
|
|
demo.launch()
|