use lcm for inpainting instead of inpainting model with lcm lora and deep cache etc, parameterize num inference steps and random seed

app.py

@@ -14,6 +14,9 @@ from datetime import datetime
 # but segformer does not work on mps lolololol
 preferred_device = "cuda" if torch.cuda.is_available() else "cpu"
 preferred_dtype = torch.float16 if preferred_device == 'cuda' else torch.float32
+inpaint_preferred_device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
+torch.backends.cuda.matmul.allow_tf32 = True
+preferred_backend = "aot_eager" if inpaint_preferred_device == "mps" else ("tensorrt" if inpaint_preferred_device == "cuda" else "inductor")
 
 seg_model_img_size = 768
 seg_model_size = 0
@@ -21,30 +24,20 @@ seg_model_size = 0
 seg_feature_extractor = SegformerFeatureExtractor.from_pretrained(f"nvidia/segformer-b{seg_model_size}-finetuned-cityscapes-{seg_model_img_size}-{seg_model_img_size}")
 seg_model = SegformerForSemanticSegmentation.from_pretrained(
     f"nvidia/segformer-b{seg_model_size}-finetuned-cityscapes-{seg_model_img_size}-{seg_model_img_size}"
-).to(preferred_device) #.to(preferred_dtype)
+).to(preferred_device).to(preferred_dtype)
 
 inpainting_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
-    "runwayml/stable-diffusion-inpainting",
-    variant="fp16",
+    "SimianLuo/LCM_Dreamshaper_v7",
     torch_dtype=preferred_dtype,
     safety_checker=None,
-).to(preferred_device)
+).to(inpaint_preferred_device)
 
-from DeepCache import DeepCacheSDHelper
-helper = DeepCacheSDHelper(pipe=inpainting_pipeline)
-helper.set_params(cache_interval=3, cache_branch_id=0)
-helper.enable()
-
-# if preferred_device == "cuda":
-#     inpainting_pipeline.unet = torch.compile(inpainting_pipeline.unet)
-#     inpainting_pipeline.vae = torch.compile(inpainting_pipeline.vae)
-
-# inpainting_pipeline.scheduler = LCMScheduler.from_config(inpainting_pipeline.scheduler.config)
-# inpainting_pipeline.load_lora_weights("latent-consistency/lcm-lora-sdv1-5", torch_dtype=preferred_dtype)
-# inpainting_pipeline.fuse_lora()
+inpainting_pipeline.unet = torch.compile(inpainting_pipeline.unet, backend=preferred_backend)
+inpainting_pipeline.vae = torch.compile(inpainting_pipeline.vae, backend=preferred_backend)
+seg_model = torch.compile(seg_model, backend=preferred_backend)
 
 seg_working_size = (seg_model_img_size, seg_model_img_size)
-repaint_working_size = (512, 512)
+repaint_working_size = (768, 768)
 
 default_inpainting_prompt = "award-winning photo of a leafy pedestrian mall full of people, with multiracial genderqueer joggers and bicyclists and wheelchair users talking and laughing"
 
@@ -63,11 +56,11 @@ def get_seg_mask(img):
     outputs = seg_model(**inputs)
     logits = outputs.logits[0]
     mask = Image.fromarray((ban_cars_mask[ torch.argmax(logits, dim=0).cpu().numpy() ]) * 255)
-    blurred_widened_mask = ImageEnhance.Contrast(mask.filter(ImageFilter.GaussianBlur(5))).enhance(9000)
+    blurred_widened_mask = ImageEnhance.Contrast(mask.filter(ImageFilter.GaussianBlur(2))).enhance(9000)
     return blurred_widened_mask
 
 
-def app(img, prompt):
+def app(img, prompt, num_inference_steps, seed):
     start_time = datetime.now().timestamp()
     old_size = Image.fromarray(img).size
     img = np.array(Image.fromarray(img).resize(seg_working_size))
@@ -79,10 +72,11 @@ def app(img, prompt):
         prompt=prompt,
         image=Image.fromarray(img).resize(repaint_working_size),
         mask_image=(mask).resize(repaint_working_size),
-        strength=0.95,
-        num_inference_steps=16,
+        strength=1,
+        num_inference_steps=num_inference_steps,
         height=repaint_working_size[0],
         width=repaint_working_size[1],
+        generator=torch.manual_seed(int(seed)),
     ).images[0]
     #overlay_img.save("overlay_raw.jpg")
     end_time = datetime.now().timestamp()
@@ -94,12 +88,13 @@ def app(img, prompt):
     #overlay_img.save("overlay_with_text.jpg")
     return overlay_img
 
-### kick the tires before we start
+# warmup, for compiling and then for timing
 
-for i in tqdm(range(2)):
-    app(np.array(Image.fromarray(np.zeros((1024,1024,3), dtype=np.uint8))), default_inpainting_prompt).save("zeros_inpainting_oneshot.jpg")
+for i in range(2):
+    for j in tqdm(range(3 ** i)):
+        app(np.array(Image.fromarray(np.zeros((1024,1024,3), dtype=np.uint8))), default_inpainting_prompt, 4, 42).save("zeros_inpainting_oneshot.jpg")
 
 #ideally:
 #iface = gr.Interface(app, gr.Image(sources=["webcam"], streaming=True), "image", live=True)
-iface = gr.Interface(app, [gr.Image(), gr.Textbox(value=default_inpainting_prompt)], "image")
+iface = gr.Interface(app, [gr.Image(), gr.Textbox(value=default_inpainting_prompt), gr.Number(minimum=1, maximum=8, value=4), gr.Number(value=42)], "image")
 iface.launch()

requirements.txt

@@ -4,4 +4,3 @@ torch==2.2.1
 accelerate
 peft
 optimum
-DeepCache