seg preferred device vs inpaint, and compile everything

app.py

@@ -12,14 +12,15 @@ from datetime import datetime
 # ideally:
 # preferred_device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
 # 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
+seg_preferred_device = "cuda" if torch.cuda.is_available() else "cpu"
+seg_preferred_dtype = torch.float16 if seg_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")
+inpaint_preferred_dtype = torch.float32 if inpaint_preferred_device == 'cpu' else torch.float16
 torch.backends.cuda.matmul.allow_tf32 = True
 
 print(f"backends: {torch._dynamo.list_backends()}")
 
-preferred_backend = "aot_eager" if inpaint_preferred_device == "mps" else ("onnxrt" if inpaint_preferred_device == "cuda" else "inductor")
+preferred_backend = "aot_eager" if inpaint_preferred_device == "mps" else "inductor"
 
 seg_model_img_size = 768
 seg_model_size = 0
@@ -27,17 +28,18 @@ 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(seg_preferred_device).to(seg_preferred_dtype)
 
 inpainting_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
     "SimianLuo/LCM_Dreamshaper_v7",
-    torch_dtype=preferred_dtype,
+    torch_dtype=inpaint_preferred_dtype,
     safety_checker=None,
 ).to(inpaint_preferred_device)
 
-#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)
+inpainting_pipeline.text_encoder = torch.compile(inpainting_pipeline.text_encoder, backend=preferred_backend)
+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)
 
@@ -54,7 +56,7 @@ ban_cars_mask = np.array(ban_cars_mask, dtype=np.uint8)
 
 
 def get_seg_mask(img):
-    inputs = seg_feature_extractor(images=img, return_tensors="pt").to(preferred_device).to(preferred_dtype)
+    inputs = seg_feature_extractor(images=img, return_tensors="pt").to(seg_preferred_device).to(seg_preferred_dtype)
     outputs = seg_model(**inputs)
     logits = outputs.logits[0]
     mask = Image.fromarray((ban_cars_mask[ torch.argmax(logits, dim=0).cpu().numpy() ]) * 255)