segformer fp16

app.py

@@ -27,7 +27,7 @@ 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(
     "SimianLuo/LCM_Dreamshaper_v7",
@@ -54,7 +54,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)
+    inputs = seg_feature_extractor(images=img, return_tensors="pt").to(preferred_device).to(preferred_dtype)
     outputs = seg_model(**inputs)
     logits = outputs.logits[0]
     mask = Image.fromarray((ban_cars_mask[ torch.argmax(logits, dim=0).cpu().numpy() ]) * 255)