Update withanyone/flux/model.py

withanyone/flux/model.py

@@ -144,7 +144,8 @@ class SiglipEmbedding(nn.Module):
         super().__init__()
         self.model = SiglipModel.from_pretrained(siglip_path).vision_model.to(torch.bfloat16)
         self.processor = AutoProcessor.from_pretrained(siglip_path)
-        self.model.to(torch.cuda.current_device())
+        # self.model.to(torch.cuda.current_device())
+        self.model.to("cuda" if torch.cuda.is_available() else "cpu", dtype=torch.bfloat16)
         
         # BiRefNet matting setup
         self.use_matting = use_matting