File size: 613 Bytes
629bb6c
9fe573e
 
629bb6c
9fe573e
629bb6c
9fe573e
 
 
4bb0520
 
9fe573e
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import gradio as gr
from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("dreamlike-art/dreamlike-photoreal-2.0")

# move to GPU if available
if torch.cuda.is_available():
    pipeline = pipeline.to("cuda")
else:
    print('No CUDA, using CPU')

def generate(prompts):
  images = pipeline(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()