Update ldm/modules/encoders/modules.py

ldm/modules/encoders/modules.py

@@ -61,7 +61,7 @@ class FrozenT5Embedder(AbstractEncoder):
         super().__init__()
         self.tokenizer = T5Tokenizer.from_pretrained(version)
         self.transformer = T5EncoderModel.from_pretrained(version)
-        self.device = device
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.max_length = max_length   # TODO: typical value?
         if freeze:
             self.freeze()
@@ -98,7 +98,7 @@ class FrozenCLIPEmbedder(AbstractEncoder):
         assert layer in self.LAYERS
         self.tokenizer = CLIPTokenizer.from_pretrained(version)
         self.transformer = CLIPTextModel.from_pretrained(version)
-        self.device = device
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.max_length = max_length
         if freeze:
             self.freeze()
@@ -148,7 +148,7 @@ class FrozenOpenCLIPEmbedder(AbstractEncoder):
         del model.visual
         self.model = model
 
-        self.device = device
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.max_length = max_length
         if freeze:
             self.freeze()
@@ -194,7 +194,7 @@ class FrozenOpenCLIPEmbedder(AbstractEncoder):
 
 
 class FrozenCLIPT5Encoder(AbstractEncoder):
-    def __init__(self, clip_version="openai/clip-vit-large-patch14", t5_version="google/t5-v1_1-xl", device="cuda",
+    def __init__(self, clip_version="openai/clip-vit-large-patch14", t5_version="google/t5-v1_1-xl", device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'),
                  clip_max_length=77, t5_max_length=77):
         super().__init__()
         self.clip_encoder = FrozenCLIPEmbedder(clip_version, device, max_length=clip_max_length)