XGBooster/StableDiffusionBasic

1import gradio as gr2import torch3from diffusers import StableDiffusionPipeline4import os5 6def get_token() -> str:7  return os.environ.get("HUGGING_FACE_TOKEN") 8 9def save_images(images: list) -> list:10 11  output_files_names = []12  for id, image in enumerate(images):13    filename = f"output{id}.png"14    image.save(filename)15    output_files_names.append(filename)16    17  return output_files_names18 19 20 21def create_img(prompt :str, number_output_requested: int) -> list:22  AUTH_TOKEN = get_token()23  generator = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", 24                                                        revision="fp16", 25                                                        torch_dtype=torch.float16, 26                                                        use_auth_token=AUTH_TOKEN)27  generator.to("cuda")28  prompt = [prompt] * number_output_requested29  with torch.autocast("cuda"):30    images = generator(prompt).images31    output_paths = save_images(images)32  return output_paths33 34diffusers_app = gr.Interface(35        fn=create_img,36        inputs =37        [38          gr.Textbox(label="Write your prompt below", placeholder = "A squirrel bench pressing 200 kg"),39          gr.Slider(value=1, minimum=1, maximum=8, step=1, label="Number of pictures to generate")40        ],41        outputs = gr.Gallery(label="Generated Images").style(grid=[2]),42        title="Text to Image with Stable Diffusion",43        description="This is a basic app to generate pictures with Stable Diffusion."44) 45 46diffusers_app.launch()