fp32 for segformer

app.py

@@ -13,7 +13,7 @@ from datetime import datetime
 # 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.float32 # torch.float16 if preferred_device == 'cuda' else torch.float32
+preferred_dtype = torch.float16 if preferred_device == 'cuda' else torch.float32
 
 seg_model_img_size = 768
 seg_model_size = 0
@@ -21,7 +21,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(
     "runwayml/stable-diffusion-inpainting",