added reference image to test diffusion

app/pipeline.py

@@ -13,13 +13,16 @@ try:
 except:
     pass
 
-import psutil
 from pydantic import BaseModel, Field
 from PIL import Image
+import psutil
 import math
 import time
 import os
 
+from dotenv import load_dotenv
+load_dotenv()
+
 taesd_model = "madebyollin/taesd"
 controlnet_model = "thibaud/controlnet-sd21-canny-diffusers"
 base_model = "stabilityai/sd-turbo"
@@ -168,7 +171,7 @@ class Pipeline:
             taesd_model, torch_dtype=torch_dtype, use_safetensors=True
         ).to(device)
 
-        if os.getenv("TORCH_COMPILE", False):
+        if bool(os.getenv("TORCH_COMPILE")):
             self.pipe.unet = torch.compile(
                 self.pipe.unet, mode="reduce-overhead", fullgraph=True
             )
@@ -181,3 +184,49 @@ class Pipeline:
                 image=[Image.new("RGB", (768, 768))],
                 control_image=[Image.new("RGB", (768, 768))],
         )
+
+    def predict(self, params: "Pipeline.InputParams", image) -> Image.Image:
+        generator = torch.manual_seed(params.seed)
+        prompt_embeds = self.pipe.compel_proc(params.prompt)
+        control_image = self.canny_torch(
+            image, params.canny_low_threshold, params.canny_high_threshold
+        )
+        steps = params.steps
+        strength = params.strength
+        if int(steps * strength) < 1:
+            steps = math.ceil(1 / max(0.10, strength))
+        last_time = time.time()
+        results = self.pipe(
+            image=image,
+            control_image=control_image,
+            prompt_embeds=prompt_embeds,
+            generator=generator,
+            strength=strength,
+            num_inference_steps=steps,
+            guidance_scale=params.guidance_scale,
+            width=params.width,
+            height=params.height,
+            output_type="pil",
+            controlnet_conditioning_scale=params.controlnet_scale,
+            control_guidance_start=params.controlnet_start,
+            control_guidance_end=params.controlnet_end,
+        )
+        print(f"Time taken: {time.time() - last_time}")
+
+        nsfw_content_detected = (
+            results.nsfw_content_detected[0]
+            if "nsfw_content_detected" in results
+            else False
+        )
+        if nsfw_content_detected:
+            return None
+        result_image = results.images[0]
+
+        if os.getenv("CONTROL_NET_OVERLAY"):
+            # paste control_image on top of result_image
+            w0, h0 = (200, 200)
+            control_image = control_image.resize((w0, h0))
+            w1, h1 = result_image.size
+            result_image.paste(control_image, (w1 - w0, h1 - h0))
+
+        return result_image

app/pipeline_test.py

@@ -1,9 +1,12 @@
 from pipeline import Pipeline
 from device import device, torch_dtype
+from diffusers.utils import load_image
 
 def main():
     p = Pipeline(device, torch_dtype)
-    print(p.InputParams.schema())
+    params = Pipeline.InputParams()
+    image = load_image("https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png")
+    p.predict(params, image).show()
 
 if __name__ == "__main__":
     main()

requirements.txt

@@ -10,6 +10,7 @@ pillow
 pydantic
 utils
 psutil
+dotenv
 
 transformers==4.35.2
 torch==2.1.1