Update trellis/pipelines/trellis_text_to_3d.py

trellis/pipelines/trellis_text_to_3d.py

@@ -9,7 +9,6 @@ from .base import Pipeline
 from . import samplers
 from ..modules import sparse as sp
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 class TrellisTextTo3DPipeline(Pipeline):
     """
@@ -70,9 +69,10 @@ class TrellisTextTo3DPipeline(Pipeline):
         Initialize the text conditioning model.
         """
         # load model
-        model = CLIPTextModel.from_pretrained(name).to(self.device)
-        tokenizer = AutoTokenizer.from_pretrained(name).to(self.device)
+        model = CLIPTextModel.from_pretrained(name)
+        tokenizer = AutoTokenizer.from_pretrained(name)
         model.eval()
+        model = model.cuda()
         self.text_cond_model = {
             'model': model,
             'tokenizer': tokenizer,
@@ -86,7 +86,7 @@ class TrellisTextTo3DPipeline(Pipeline):
         """
         assert isinstance(text, list) and all(isinstance(t, str) for t in text), "text must be a list of strings"
         encoding = self.text_cond_model['tokenizer'](text, max_length=77, padding='max_length', truncation=True, return_tensors='pt')
-        tokens = encoding['input_ids'].to(self.device)
+        tokens = encoding['input_ids'].cuda()
         embeddings = self.text_cond_model['model'](input_ids=tokens).last_hidden_state
         
         return embeddings
@@ -225,4 +225,4 @@ class TrellisTextTo3DPipeline(Pipeline):
         torch.manual_seed(seed)
         coords = self.sample_sparse_structure(cond, num_samples, sparse_structure_sampler_params)
         slat = self.sample_slat(cond, coords, slat_sampler_params)
-        return self.decode_slat(slat, formats)
+        return self.decode_slat(slat, formats)