File size: 628 Bytes
145747e
f682c12
145747e
 
 
 
 
 
 
 
f682c12
145747e
 
91a656b
 
 
 
1115b57
 
91a656b
145747e
 
 
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()