title

README.md

@@ -1,5 +1,5 @@
 ---
-title: Unofficial SDXL Turbo Real Time Text To Image
+title: Unofficial SDXL Turbo Img2Img Txt2Img
 emoji: 🏆
 colorFrom: yellow
 colorTo: purple

app.py

@@ -1,4 +1,4 @@
-from diffusers import DiffusionPipeline
+from diffusers import AutoPipelineForImage2Image, AutoPipelineForText2Image
 import torch
 import os
 
@@ -35,31 +35,79 @@ if mps_available:
     torch_dtype = torch.float32
 
 if SAFETY_CHECKER == "True":
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", revision="pr/4")
+    i2i_pipe = AutoPipelineForImage2Image.from_pretrained(
+        "stabilityai/sdxl-turbo",
+        torch_dtype=torch_dtype,
+        variant="fp16" if torch_dtype == torch.float16 else "fp32",
+    )
+    t2i_pipe = AutoPipelineForText2Image.from_pretrained(
+        "stabilityai/sdxl-turbo",
+        torch_dtype=torch_dtype,
+        variant="fp16" if torch_dtype == torch.float16 else "fp32",
+    )
 else:
-    pipe = DiffusionPipeline.from_pretrained(
-        "stabilityai/sdxl-turbo", revision="pr/4", safety_checker=None
+    i2i_pipe = AutoPipelineForImage2Image.from_pretrained(
+        "stabilityai/sdxl-turbo",
+        safety_checker=None,
+        torch_dtype=torch_dtype,
+        variant="fp16" if torch_dtype == torch.float16 else "fp32",
+    )
+    t2i_pipe = AutoPipelineForText2Image.from_pretrained(
+        "stabilityai/sdxl-turbo",
+        safety_checker=None,
+        torch_dtype=torch_dtype,
+        variant="fp16" if torch_dtype == torch.float16 else "fp32",
     )
 
 
-pipe.to(device=torch_device, dtype=torch_dtype).to(device)
-pipe.unet.to(memory_format=torch.channels_last)
-pipe.set_progress_bar_config(disable=True)
-
-
-def predict(prompt, steps, seed=1231231):
-    generator = torch.manual_seed(seed)
-    last_time = time.time()
-    results = pipe(
-        prompt=prompt,
-        generator=generator,
-        num_inference_steps=steps,
-        guidance_scale=0.0,
-        width=512,
-        height=512,
-        # original_inference_steps=params.lcm_steps,
-        output_type="pil",
-    )
+t2i_pipe.to(device=torch_device, dtype=torch_dtype).to(device)
+t2i_pipe.set_progress_bar_config(disable=True)
+i2i_pipe.to(device=torch_device, dtype=torch_dtype).to(device)
+i2i_pipe.set_progress_bar_config(disable=True)
+
+
+def pad_image(image):
+    w, h = image.size
+    if w == h:
+        return image
+    elif w > h:
+        new_image = Image.new(image.mode, (w, w), (0, 0, 0))
+        new_image.paste(image, (0, (w - h) // 2))
+        return new_image
+    else:
+        new_image = Image.new(image.mode, (h, h), (0, 0, 0))
+        new_image.paste(image, ((h - w) // 2, 0))
+        return new_image
+
+
+async def predict(init_image, prompt, strength, steps, seed=1231231):
+    if init_image is not None:
+        init_image = pad_image(init_image).convert("RGB").resize((512, 512))
+        generator = torch.manual_seed(seed)
+        last_time = time.time()
+        results = i2i_pipe(
+            prompt=prompt,
+            image=init_image,
+            generator=generator,
+            num_inference_steps=steps,
+            guidance_scale=0.0,
+            strength=strength,
+            width=512,
+            height=512,
+            output_type="pil",
+        )
+    else:
+        generator = torch.manual_seed(seed)
+        last_time = time.time()
+        results = t2i_pipe(
+            prompt=prompt,
+            generator=generator,
+            num_inference_steps=steps,
+            guidance_scale=0.0,
+            width=512,
+            height=512,
+            output_type="pil",
+        )
     print(f"Pipe took {time.time() - last_time} seconds")
     nsfw_content_detected = (
         results.nsfw_content_detected[0]
@@ -75,7 +123,7 @@ def predict(prompt, steps, seed=1231231):
 css = """
 #container{
     margin: 0 auto;
-    max-width: 40rem;
+    max-width: 80rem;
 }
 #intro{
     max-width: 100%;
@@ -84,9 +132,10 @@ css = """
 }
 """
 with gr.Blocks(css=css) as demo:
+    init_image_state = gr.State()
     with gr.Column(elem_id="container"):
         gr.Markdown(
-            """# SDXL Turbo - Text To Image
+            """# SDXL Turbo Image to Image/Text to Image
             ## Unofficial Demo
             SDXL Turbo model can generate high quality images in a single pass read more on [stability.ai post](https://stability.ai/news/stability-ai-sdxl-turbo).  
             **Model**: https://huggingface.co/stabilityai/sdxl-turbo
@@ -94,18 +143,40 @@ with gr.Blocks(css=css) as demo:
             elem_id="intro",
         )
         with gr.Row():
-            with gr.Row():
-                prompt = gr.Textbox(
-                    placeholder="Insert your prompt here:", scale=5, container=False
+            prompt = gr.Textbox(
+                placeholder="Insert your prompt here:",
+                scale=5,
+                container=False,
+            )
+            generate_bt = gr.Button("Generate", scale=1)
+        with gr.Row():
+            with gr.Column():
+                image_input = gr.Image(
+                    sources=["upload", "webcam", "clipboard"],
+                    label="Webcam",
+                    type="pil",
                 )
-                generate_bt = gr.Button("Generate", scale=1)
+            with gr.Column():
+                image = gr.Image(type="filepath")
+                with gr.Accordion("Advanced options", open=False):
+                    strength = gr.Slider(
+                        label="Strength",
+                        value=0.7,
+                        minimum=0.0,
+                        maximum=1.0,
+                        step=0.001,
+                    )
+                    steps = gr.Slider(
+                        label="Steps", value=2, minimum=1, maximum=10, step=1
+                    )
+                    seed = gr.Slider(
+                        randomize=True,
+                        minimum=0,
+                        maximum=12013012031030,
+                        label="Seed",
+                        step=1,
+                    )
 
-        image = gr.Image(type="filepath")
-        with gr.Accordion("Advanced options", open=False):
-            steps = gr.Slider(label="Steps", value=2, minimum=1, maximum=10, step=1)
-            seed = gr.Slider(
-                randomize=True, minimum=0, maximum=12013012031030, label="Seed", step=1
-            )
         with gr.Accordion("Run with diffusers"):
             gr.Markdown(
                 """## Running SDXL Turbo with `diffusers`
@@ -116,7 +187,7 @@ with gr.Blocks(css=css) as demo:
             from diffusers import DiffusionPipeline
 
             pipe = DiffusionPipeline.from_pretrained(
-                "stabilityai/sdxl-turbo", revision="refs/pr/4"
+                "stabilityai/sdxl-turbo"
             ).to("cuda")
             results = pipe(
                 prompt="A cinematic shot of a baby racoon wearing an intricate italian priest robe",
@@ -129,11 +200,19 @@ with gr.Blocks(css=css) as demo:
             """
             )
 
-        inputs = [prompt, steps, seed]
+        inputs = [image_input, prompt, strength, steps, seed]
         generate_bt.click(fn=predict, inputs=inputs, outputs=image, show_progress=False)
         prompt.input(fn=predict, inputs=inputs, outputs=image, show_progress=False)
         steps.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
         seed.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
+        strength.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
+        image_input.change(
+            fn=lambda x: x,
+            inputs=image_input,
+            outputs=init_image_state,
+            show_progress=False,
+            queue=False,
+        )
 
 demo.queue()
 demo.launch()