Spaces:
Runtime error
Runtime error
| from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline | |
| import torch | |
| import gradio as gr | |
| import random | |
| import socket | |
| import re | |
| import cpuinfo | |
| import gpustat | |
| import pathlib | |
| def print_system_info(): | |
| # print CPU information | |
| cpu_info = cpuinfo.get_cpu_info() | |
| cpu_str = f"CPU:\n Model: {cpu_info['brand_raw']}\n Cores: {cpu_info['count']}\n" | |
| # print GPU information | |
| gpu_stats = gpustat.GPUStatCollection.new_query() | |
| gpu_str = "GPU:\n" | |
| for i, gpu in enumerate(gpu_stats.gpus): | |
| gpu_str += f" GPU {i}: {gpu.name}\n" | |
| return cpu_str + gpu_str | |
| #Preload model/llm | |
| pipe = StableDiffusionXLPipeline.from_pretrained( | |
| "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16",add_watermarker=False | |
| ).to("cuda") | |
| def to_filename(text, max_length=200): | |
| # replace any non-alphanumeric characters with an empty string | |
| filename = re.sub(r'\W+', '', text) | |
| # replace any spaces with an underscore | |
| filename = re.sub(r'\s+', '_', filename) | |
| # truncate the filename to the desired length | |
| filename = filename[:max_length] | |
| return filename | |
| #Description in Markdown | |
| description = """ | |
| This is a demo ⚡ of Stable Diffusion XL, a new diffusion model that can generate high quality images. [Stable Diffusion XL](https://huggingface.co./docs/diffusers/using-diffusers/sdxl) | |
| It is being presented 📦 using Gradio, a library for quickly creating UIs for machine learning models. [Gradio](https://www.gradio.app/docs/interface) | |
| """ | |
| article = """ | |
| Coded 🧾 by [Matthew Rogers](https://matthewrogers.org) | |
| Running on: | |
| """ | |
| # get system information | |
| system_info = print_system_info() | |
| # append system information to article | |
| article += system_info | |
| # create a directory to save the generated images | |
| saved_dir = pathlib.Path("saved") | |
| saved_dir.mkdir(parents=True, exist_ok=True) | |
| # Function to generate the output and the image | |
| def generateImage(prompt, prompt_2,num_inference_steps=50,seed=0): | |
| #create image, use default 50 for num_inference_steps | |
| # set a random seed if seed is not defined | |
| if seed == 0 or seed == None: | |
| seed = random.randint(0, 2**32 - 1) | |
| print("seed:",seed) | |
| # set the seed for the PyTorch random number generator | |
| torch.manual_seed(seed) | |
| if num_inference_steps == 0: | |
| num_inference_steps = 50 | |
| image = pipe(prompt=prompt,prompt_2=prompt_2, width=1024, height=1024, num_inference_steps=num_inference_steps, ).images[0] | |
| used_seed = torch.initial_seed() | |
| print("used_seed:",used_seed) | |
| # generate a filename based on the input prompts and the number of inference steps | |
| filename = to_filename(f"{prompt}_{prompt_2}_{num_inference_steps}_{used_seed}", max_length=200) | |
| # add the file extension | |
| filename = f"saved/{filename}.png" | |
| image.save( filename, "PNG") | |
| #return values | |
| return image, used_seed, num_inference_steps, (sorted(saved_dir.glob("*.png"), key=lambda f: f.stat().st_ctime, reverse=True)) | |
| gui = gr.Interface( | |
| fn=generateImage, | |
| inputs=[gr.Textbox(lines=2,label="Prompt"), gr.Textbox(lines=2,label="Prompt2"), gr.Slider(0, 150, value=80, label="Steps"),gr.Number(label="Seed")], | |
| outputs=[gr.Image(label="Generated Image"), gr.Number(label="Seed"), gr.Number(label="Steps"), gr.Gallery(label="Gallery")], | |
| title="Stable Diffusion XL on Gradio", | |
| description=description, | |
| article=article, | |
| allow_flagging="never", | |
| ) | |
| gui.launch() |