diffusers/check_pr

1from diffusers import DiffusionPipeline2import gradio as gr3import torch4import time5import psutil6 7 8start_time = time.time()9 10device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"11 12 13def error_str(error, title="Error"):14    return (15        f"""#### {title}16            {error}"""17        if error18        else ""19    )20 21 22def inference(23    repo_id,24    discuss_nr,25    prompt,26):27 28    print(psutil.virtual_memory())  # print memory usage29 30    seed = 031    torch_device = "cuda" if "GPU" in device else "cpu"32 33    generator = torch.Generator(torch_device).manual_seed(seed)34 35    dtype = torch.float16 if torch_device == "cuda" else torch.float3236 37    try:38        revision = f"refs/pr/{discuss_nr}" if (discuss_nr != "" or discuss_nr is None) else None39        pipe = DiffusionPipeline.from_pretrained(repo_id, revision=revision, torch_dtype=dtype)40        pipe.to(torch_device)41 42        return pipe(prompt, generator=generator, num_inference_steps=25).images, f"Done. Seed: {seed}"43    except Exception as e:44        url = f"https://huggingface.co/{repo_id}/discussions/{discuss_nr}"45        message = f"There is a problem with your diffusers weights of the PR: {url}. Error message: \n"46        return None, error_str(message + e)47 48 49with gr.Blocks(css="style.css") as demo:50    gr.HTML(51        f"""52            <div class="diffusion">53              <p>54               Space to test whether `diffusers` PRs work.55              </p>56              <p>57               Running on <b>{device}</b>58              </p>59            </div>60        """61    )62    with gr.Row():63 64        with gr.Column(scale=55):65            with gr.Group():66                repo_id = gr.Textbox(67                    label="Repo id on Hub",68                    placeholder="Path to model, e.g. CompVis/stable-diffusion-v1-4 for https://huggingface.co/CompVis/stable-diffusion-v1-4",69                )70                discuss_nr = gr.Textbox(71                    label="Discussion number",72                    placeholder="Number of the discussion that should be checked, e.g. 171 for https://huggingface.co/CompVis/stable-diffusion-v1-4/discussions/171",73                )74                prompt = gr.Textbox(75                    label="Prompt",76                    default="An astronaut riding a horse on Mars.",77                    placeholder="Enter prompt.",78                )79                gallery = gr.Gallery(80                    label="Generated images", show_label=False, elem_id="gallery"81                ).style(grid=[2], height="auto")82 83            error_output = gr.Markdown()84 85            generate = gr.Button(value="Generate").style(86                rounded=(False, True, True, False)87            )88 89    inputs = [90        repo_id,91        discuss_nr,92        prompt,93    ]94    outputs = [gallery, error_output]95    prompt.submit(inference, inputs=inputs, outputs=outputs)96    generate.click(inference, inputs=inputs, outputs=outputs)97 98print(f"Space built in {time.time() - start_time:.2f} seconds")99 100demo.queue(concurrency_count=1)101demo.launch()102