reset pytorch matrix multiplication precision for rmbg

app.py

@@ -10,21 +10,6 @@ from sam2.sam2_image_predictor import SAM2ImagePredictor
 import numpy as np
 from simple_lama_inpainting import SimpleLama
 
-torch.set_float32_matmul_precision(["high", "highest"][0])
-
-birefnet = AutoModelForImageSegmentation.from_pretrained(
-    "ZhengPeng7/BiRefNet", trust_remote_code=True
-)
-birefnet.to("cuda")
-
-
-transform_image = transforms.Compose(
-    [
-        transforms.Resize((1024, 1024)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-    ]
-)
 
 pipe = FluxFillPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16
@@ -113,6 +98,18 @@ def rmbg(image=None, url=None):
         image = url
     image = load_img(image).convert("RGB")
     image_size = image.size
+    torch.set_float32_matmul_precision(["high", "highest"][0])
+    birefnet = AutoModelForImageSegmentation.from_pretrained(
+        "ZhengPeng7/BiRefNet", trust_remote_code=True
+    )
+    birefnet.to("cuda")
+    transform_image = transforms.Compose(
+        [
+            transforms.Resize((1024, 1024)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
     input_images = transform_image(image).unsqueeze(0).to("cuda")
     # Prediction
     with torch.no_grad():
@@ -121,6 +118,7 @@ def rmbg(image=None, url=None):
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(image_size)
     image.putalpha(mask)
+    torch.set_float32_matmul_precision(["high", "highest"][1])
     return image