Update withanyone/flux/model.py

withanyone/flux/model.py

@@ -145,13 +145,14 @@ class SiglipEmbedding(nn.Module):
         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("cuda" if torch.cuda.is_available() else "cpu", dtype=torch.bfloat16)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
         
         # BiRefNet matting setup
         self.use_matting = use_matting
         if self.use_matting:
             self.birefnet = AutoModelForImageSegmentation.from_pretrained(
-                'briaai/RMBG-2.0', trust_remote_code=True).to(torch.cuda.current_device(), dtype=torch.bfloat16)
+                'briaai/RMBG-2.0', trust_remote_code=True).to(self.device, dtype=torch.bfloat16)
             # Apply half precision to the entire model after loading
             self.matting_transform = transforms.Compose([
                 # transforms.Resize((512, 512)),
@@ -165,7 +166,7 @@ class SiglipEmbedding(nn.Module):
             return image
             
         # Convert to input format and move to GPU
-        input_image = self.matting_transform(image).unsqueeze(0).to(torch.cuda.current_device(), dtype=torch.bfloat16)
+        input_image = self.matting_transform(image).unsqueeze(0).to(self.device, dtype=torch.bfloat16)
 
         # Generate prediction
         with torch.no_grad(), autocast(dtype=torch.bfloat16):
@@ -205,7 +206,7 @@ class SiglipEmbedding(nn.Module):
                     
                 pixel_values = self.processor(images=processed_images, return_tensors="pt").pixel_values
                 # device
-                pixel_values = pixel_values.to(torch.cuda.current_device(), dtype=torch.bfloat16)
+                pixel_values = pixel_values.to(self.device, dtype=torch.bfloat16)
                 last_hidden_state = self.model(pixel_values).last_hidden_state # 2, 256 768
                 # pooled_output = self.model(pixel_values).pooler_output # 2, 768
                 siglip_embedding.append(last_hidden_state)
@@ -217,14 +218,14 @@ class SiglipEmbedding(nn.Module):
                 for _ in range(4 - batch_size):
                     pixel_values = self.processor(images=processed_images, return_tensors="pt").pixel_values
                     # device
-                    pixel_values = pixel_values.to(torch.cuda.current_device(), dtype=torch.bfloat16)
+                    pixel_values = pixel_values.to(self.device, dtype=torch.bfloat16)
                     last_hidden_state = self.model(pixel_values).last_hidden_state
 
         elif isinstance(refimage, torch.Tensor):
             # refimage is a tensor of shape (batch_size, num_of_person, 3, H, W)
             batch_size, num_of_person, C, H, W = refimage.shape
             refimage = refimage.view(batch_size * num_of_person, C, H, W)
-            refimage = refimage.to(torch.cuda.current_device(), dtype=torch.bfloat16)
+            refimage = refimage.to(self.device, dtype=torch.bfloat16)
             last_hidden_state = self.model(refimage).last_hidden_state
             siglip_embedding = last_hidden_state.view(batch_size, num_of_person, 256, 768)