up

app.py

@@ -1,53 +1,11 @@
-from diffusers import (
-    StableDiffusionPipeline,
-    StableDiffusionImg2ImgPipeline,
-    DPMSolverMultistepScheduler,
-)
+from diffusers import DiffusionPipeline
 import gradio as gr
 import torch
-from PIL import Image
 import time
 import psutil
-import random
-from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
 
 
 start_time = time.time()
-current_steps = 25
-
-SAFETY_CHECKER = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=torch.float16)
-
-
-class Model:
-    def __init__(self, name, path=""):
-        self.name = name
-        self.path = path
-
-        if path != "":
-            self.pipe_t2i = StableDiffusionPipeline.from_pretrained(
-                path, torch_dtype=torch.float16, safety_checker=SAFETY_CHECKER
-            )
-            self.pipe_t2i.scheduler = DPMSolverMultistepScheduler.from_config(
-                self.pipe_t2i.scheduler.config
-            )
-            self.pipe_i2i = StableDiffusionImg2ImgPipeline(**self.pipe_t2i.components)
-        else:
-            self.pipe_t2i = None
-            self.pipe_i2i = None
-
-
-models = [
-    Model("Protogen v2.2 (Anime)", "darkstorm2150/Protogen_v2.2_Official_Release"),
-    Model("Protogen x3.4 (Photorealism)", "darkstorm2150/Protogen_x3.4_Official_Release"),
-    Model("Protogen x5.3 (Photorealism)", "darkstorm2150/Protogen_x5.3_Official_Release"),
-    Model("Protogen x5.8 Rebuilt (Scifi+Anime)", "darkstorm2150/Protogen_x5.8_Official_Release"),
-    Model("Protogen Dragon (RPG Model)", "darkstorm2150/Protogen_Dragon_Official_Release"),
-    Model("Protogen Nova", "darkstorm2150/Protogen_Nova_Official_Release"),
-    Model("Protogen Eclipse", "darkstorm2150/Protogen_Eclipse_Official_Release"),
-    Model("Protogen Infinity", "darkstorm2150/Protogen_Infinity_Official_Release"),
-]
-
-MODELS = {m.name: m for m in models}
 
 device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
 
@@ -62,263 +20,80 @@ def error_str(error, title="Error"):
 
 
 def inference(
-    model_name,
+    repo_id,
+    pr,
     prompt,
-    guidance,
-    steps,
-    n_images=1,
-    width=512,
-    height=512,
-    seed=0,
-    img=None,
-    strength=0.5,
-    neg_prompt="",
 ):
 
     print(psutil.virtual_memory())  # print memory usage
 
-    if seed == 0:
-        seed = random.randint(0, 2147483647)
-
-    generator = torch.Generator("cuda").manual_seed(seed)
-
-    try:
-        if img is not None:
-            return (
-                img_to_img(
-                    model_name,
-                    prompt,
-                    n_images,
-                    neg_prompt,
-                    img,
-                    strength,
-                    guidance,
-                    steps,
-                    width,
-                    height,
-                    generator,
-                    seed,
-                ),
-                f"Done. Seed: {seed}",
-            )
-        else:
-            return (
-                txt_to_img(
-                    model_name,
-                    prompt,
-                    n_images,
-                    neg_prompt,
-                    guidance,
-                    steps,
-                    width,
-                    height,
-                    generator,
-                    seed,
-                ),
-                f"Done. Seed: {seed}",
-            )
-    except Exception as e:
-        return None, error_str(e)
-
-
-def txt_to_img(
-    model_name,
-    prompt,
-    n_images,
-    neg_prompt,
-    guidance,
-    steps,
-    width,
-    height,
-    generator,
-    seed,
-):
-    pipe = MODELS[model_name].pipe_t2i
-
-    if torch.cuda.is_available():
-        pipe = pipe.to("cuda")
-        pipe.enable_xformers_memory_efficient_attention()
-
-    result = pipe(
-        prompt,
-        negative_prompt=neg_prompt,
-        num_images_per_prompt=n_images,
-        num_inference_steps=int(steps),
-        guidance_scale=guidance,
-        width=width,
-        height=height,
-        generator=generator,
-    )
-
-    pipe.to("cpu")
-
-    return replace_nsfw_images(result)
-
-
-def img_to_img(
-    model_name,
-    prompt,
-    n_images,
-    neg_prompt,
-    img,
-    strength,
-    guidance,
-    steps,
-    width,
-    height,
-    generator,
-    seed,
-):
-    pipe = MODELS[model_name].pipe_i2i
-
-    if torch.cuda.is_available():
-        pipe = pipe.to("cuda")
-        pipe.enable_xformers_memory_efficient_attention()
+    seed = 0
+    torch_device = "cuda" if "GPU" in device else "cpu"
 
-    ratio = min(height / img.height, width / img.width)
-    img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)
+    generator = torch.Generator(torch_device).manual_seed(seed)
 
-    result = pipe(
-        prompt,
-        negative_prompt=neg_prompt,
-        num_images_per_prompt=n_images,
-        image=img,
-        num_inference_steps=int(steps),
-        strength=strength,
-        guidance_scale=guidance,
-        generator=generator,
-    )
-
-    pipe.to("cpu")
-
-    return replace_nsfw_images(result)
+    dtype = torch.float16 if torch_device == "cuda" else torch.float32
 
+    try:
+        pipe = DiffusionPipeline.from_pretrained(repo_id, revision=pr, torch_dtype=dtype)
+        pipe.to(torch_device)
 
-def replace_nsfw_images(results):
-    for i in range(len(results.images)):
-        if results.nsfw_content_detected[i]:
-            results.images[i] = Image.open("nsfw.png")
-    return results.images
+        return pipe(prompt, generator=generator, num_inference_steps=25).images
+    except Exception as e:
+        url = f"https://huggingface.co/{repo_id}/discussions/{pr.split('/')[-1]}"
+        message = f"There is a problem with your diffusers weights of the PR: {url}. Error message: \n"
+        return None, error_str(message + e)
 
 
 with gr.Blocks(css="style.css") as demo:
+    gr.HTML(
+        f"""
+            <div class="diffusion">
+              <p>
+               Space to test whether `diffusers` PRs work.
+              </p>
+              <p>
+               Running on <b>{device}</b>
+              </p>
+            </div>
+        """
+    )
     with gr.Row():
 
         with gr.Column(scale=55):
             with gr.Group():
-                prompt = gr.Textbox(
+                repo_id = gr.Textbox(
                     label="Repo id on Hub",
                     placeholder="Path to model, e.g. CompVis/stable-diffusion-v1-4",
                 )
-                with gr.Box(visible=False) as custom_model_group:
-                    custom_model_path = gr.Textbox(
-                        label="Custom model path",
-                        placeholder="Path to model, e.g. darkstorm2150/Protogen_x3.4_Official_Release",
-                        interactive=True,
-                    )
-                    gr.HTML(
-                        "<div><font size='2'>Custom models have to be downloaded first, so give it some time.</font></div>"
-                    )
-
-                with gr.Row():
-                    prompt = gr.Textbox(
-                        label="Prompt",
-                        show_label=False,
-                        max_lines=2,
-                        placeholder="Enter prompt.",
-                    ).style(container=False)
-                    generate = gr.Button(value="Generate").style(
-                        rounded=(False, True, True, False)
-                    )
-
-                # image_out = gr.Image(height=512)
+                pr = gr.Textbox(
+                    label="PR branch",
+                    placeholder="PR branch that should be checked, e.g. refs/pr/171",
+                )
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    default="An astronaut riding a horse on Mars.",
+                    placeholder="Enter prompt.",
+                )
                 gallery = gr.Gallery(
                     label="Generated images", show_label=False, elem_id="gallery"
                 ).style(grid=[2], height="auto")
 
-            state_info = gr.Textbox(label="State", show_label=False, max_lines=2).style(
-                container=False
-            )
             error_output = gr.Markdown()
 
-        with gr.Column(scale=45):
-            with gr.Tab("Options"):
-                with gr.Group():
-                    neg_prompt = gr.Textbox(
-                        label="Negative prompt",
-                        placeholder="What to exclude from the image",
-                    )
-
-                    n_images = gr.Slider(
-                        label="Images", value=1, minimum=1, maximum=4, step=1
-                    )
-
-                    with gr.Row():
-                        guidance = gr.Slider(
-                            label="Guidance scale", value=7.5, maximum=15
-                        )
-                        steps = gr.Slider(
-                            label="Steps",
-                            value=current_steps,
-                            minimum=2,
-                            maximum=75,
-                            step=1,
-                        )
-
-                    with gr.Row():
-                        width = gr.Slider(
-                            label="Width", value=512, minimum=64, maximum=1024, step=8
-                        )
-                        height = gr.Slider(
-                            label="Height", value=512, minimum=64, maximum=1024, step=8
-                        )
-
-                    seed = gr.Slider(
-                        0, 2147483647, label="Seed (0 = random)", value=0, step=1
-                    )
-
-            with gr.Tab("Image to image"):
-                with gr.Group():
-                    image = gr.Image(
-                        label="Image", height=256, tool="editor", type="pil"
-                    )
-                    strength = gr.Slider(
-                        label="Transformation strength",
-                        minimum=0,
-                        maximum=1,
-                        step=0.01,
-                        value=0.5,
-                    )
+            generate = gr.Button(value="Generate").style(
+                rounded=(False, True, True, False)
+            )
 
     inputs = [
-        model_name,
+        repo_id,
+        pr,
         prompt,
-        guidance,
-        steps,
-        n_images,
-        width,
-        height,
-        seed,
-        image,
-        strength,
-        neg_prompt,
     ]
     outputs = [gallery, error_output]
     prompt.submit(inference, inputs=inputs, outputs=outputs)
     generate.click(inference, inputs=inputs, outputs=outputs)
 
-    gr.HTML(
-        """
-    <div style="border-top: 1px solid #303030;">
-      <br>
-      <p>Models by <a href="https://huggingface.co/darkstorm2150">@darkstorm2150</a> and others. ❤️</p>
-      <p>This space uses the <a href="https://github.com/LuChengTHU/dpm-solver">DPM-Solver++</a> sampler by <a href="https://arxiv.org/abs/2206.00927">Cheng Lu, et al.</a>.</p>
-      <p>Space by: Darkstorm (Victor Espinoza)<br>
-      <a href="https://www.instagram.com/officialvictorespinoza/">Instagram</a>
-    </div>
-    """
-    )
-
 print(f"Space built in {time.time() - start_time:.2f} seconds")
 
 demo.queue(concurrency_count=1)