inpainting

app.py

@@ -11,13 +11,19 @@ import numpy as np
 from PIL import Image
 from datetime import datetime
 
-preferred_dtype = torch.float32
 preferred_device = "cuda" if torch.cuda.is_available() else "cpu"
+preferred_dtype = torch.float32 if preferred_device == 'cpu' else torch.float16
 
 def label_func(fn): return path/"labels"/f"{fn.stem}_P{fn.suffix}"
 
 segmodel = load_learner("camvid-512.pkl")
 
+inpainting_pipeline = AutoPipelineForInpainting(
+    model="runwayml/stable-diffusion-inpainting",
+    revision="fp16",
+    torch_dtype=preferred_dtype,
+).to(preferred_device)
+
 working_size = (512, 512)
 
 default_inpainting_prompt = "watercolor of a leafy pedestrian mall at golden hour with multiracial genderqueer joggers and bicyclists and wheelchair users talking and laughing"
@@ -49,10 +55,20 @@ def app(img, prompt):
     img = np.array(Image.fromarray(img).resize(working_size))
     mask = ban_cars_mask[get_seg_mask(img)]
     mask = mask * 255
-    overlay_img = Image.fromarray(np.stack([img[:, :, 0], mask, img[:,:,2]], axis=-1))
+    mask_time = datetime.now().timestamp()
+    overlay_img = inpainting_pipeline(
+        prompt=prompt,
+        image=Image.fromarray(img),
+        mask=Image.fromarray(mask),
+        strength=0.95,
+        num_inference_steps=13,
+    ).images[0]
     end_time = datetime.now().timestamp()
     draw = ImageDraw.Draw(overlay_img)
-    draw.text((50, 10), f"Old size: {old_size}\nDuration: {int(1000 * (end_time - start_time))}ms\n<{prompt}>", fill=(255, 255, 255))
+    # replace spaces with newlines after many words to line break prompt
+    prompt = " ".join([prompt.split(" ")[i] if (i+1) % 5 else prompt.split(" ")[i] + "\n" for i in range(len(prompt.split(" ")))])
+
+    draw.text((50, 10), f"Old size: {old_size}\nTotal duration: {int(1000 * (end_time - start_time))}ms\nSegmentation {int(1000 * (mask_time - start_time))}ms / inpainting {int(1000 * (end_time - mask_time))} \n<{prompt}>", fill=(255, 255, 255))
     return overlay_img
 
 #ideally: