gradio/demo/diffusers_with_batching/run.ipynb
aliabid94 9b42ba8f10
Update guides esp plots (#8907)
* changes

* changes

* revert changes

* changes

* add changeset

* notebooks script

* changes

* changes

---------

Co-authored-by: Ali Abid <aliabid94@gmail.com>
Co-authored-by: gradio-pr-bot <gradio-pr-bot@users.noreply.github.com>
Co-authored-by: Ali Abdalla <ali.si3luwa@gmail.com>
2024-07-29 22:08:51 -07:00

1.2 KiB

Gradio Demo: diffusers_with_batching

In [ ]:
!pip install -q gradio torch transformers diffusers
In [ ]:
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
)

if __name__ == "__main__":
    demo.launch()