T5 offload w/ bnb

Mentioned in #22

README.md

@@ -140,7 +140,7 @@ If you have a second GPU, selecting "cuda:1" as the device will allow you to use
 
 Loaded in bnb4bit mode, it only takes around 6GB VRAM, making it work with 12GB cards. The only drawback is that it'll constantly stay in VRAM since BitsAndBytes doesn't allow moving the weights to the system RAM temporarily. Switching to a different workflow *should* still release the VRAM as expected. Pascal cards (1080ti, P40) seem to struggle with 4bit. Select "cpu" if you encounter issues.
 
-On windows, you may need a newer version of bitsandbytes for 4bit. Try `python -m pip install bitsandbytes --prefer-binary --extra-index-url=https://jllllll.github.io/bitsandbytes-windows-webui`
+On windows, you may need a newer version of bitsandbytes for 4bit. Try `python -m pip install bitsandbytes`
 
 > [!IMPORTANT]  
 > You may also need to upgrade transformers and install spiece for the tokenizer. `pip install -r requirements.txt`

T5/nodes.py

@@ -53,7 +53,7 @@ class T5v11Loader:
 
 	def load_model(self, t5v11_name, t5v11_ver, path_type, device, dtype):
 		if "bnb" in dtype:
-			assert device == "gpu", "BitsAndBytes only works on CUDA! Set device to 'gpu'."
+			assert device == "gpu" or device.startswith("cuda"), "BitsAndBytes only works on CUDA! Set device to 'gpu'."
 		dtype = string_to_dtype(dtype, "text_encoder")
 		if device == "cpu":
 			assert dtype in [None, torch.float32], f"Can't use dtype '{dtype}' with CPU! Set dtype to 'default'."

T5/t5v11.py

@@ -33,9 +33,9 @@ class T5v11Model(torch.nn.Module):
             else:
                 if dtype: model_args["torch_dtype"] = dtype
                 self.bnb = False
-                # second GPU offload hack part 2
-                if device.startswith("cuda"):
-                    model_args["device_map"] = device
+            # second GPU offload hack part 2
+            if device.startswith("cuda"):
+                model_args["device_map"] = device
             print(f"Loading T5 from '{textmodel_path}'")
             self.transformer = T5EncoderModel.from_pretrained(textmodel_path, **model_args)
         else: