Stable-X commited on
Commit
0109587
·
verified ·
1 Parent(s): 4813e61

Update trellis/pipelines/trellis_image_to_3d.py

Browse files
trellis/pipelines/trellis_image_to_3d.py CHANGED
@@ -338,7 +338,7 @@ class TrellisImageTo3DPipeline(Pipeline):
338
  transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
339
  ])
340
 
341
- input_images = transform_image(image).unsqueeze(0).to(self.device)
342
 
343
  with torch.no_grad():
344
  preds = self.birefnet_model(input_images)[-1].sigmoid().cpu()
@@ -818,7 +818,7 @@ class TrellisVGGTTo3DPipeline(TrellisImageTo3DPipeline):
818
  new_pipeline.birefnet_model = AutoModelForImageSegmentation.from_pretrained(
819
  'ZhengPeng7/BiRefNet',
820
  trust_remote_code=True
821
- ).to(new_pipeline.device)
822
  new_pipeline.birefnet_model.eval()
823
 
824
  new_pipeline.sparse_structure_sampler = getattr(samplers, args['sparse_structure_sampler']['name'])(**args['sparse_structure_sampler']['args'])
 
338
  transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
339
  ])
340
 
341
+ input_images = transform_image(image).unsqueeze(0).cpu()
342
 
343
  with torch.no_grad():
344
  preds = self.birefnet_model(input_images)[-1].sigmoid().cpu()
 
818
  new_pipeline.birefnet_model = AutoModelForImageSegmentation.from_pretrained(
819
  'ZhengPeng7/BiRefNet',
820
  trust_remote_code=True
821
+ ).cpu()
822
  new_pipeline.birefnet_model.eval()
823
 
824
  new_pipeline.sparse_structure_sampler = getattr(samplers, args['sparse_structure_sampler']['name'])(**args['sparse_structure_sampler']['args'])