T5 offload w/ bnb

Mentioned in #22

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: