| 71 | |
| 72 | |
| 73 | class FreeSplatterRunner: |
| 74 | def __init__(self, device): |
| 75 | self.device = device |
| 76 | |
| 77 | # background remover |
| 78 | self.rembg = AutoModelForImageSegmentation.from_pretrained( |
| 79 | "briaai/RMBG-2.0", |
| 80 | trust_remote_code=True, |
| 81 | cache_dir='ckpts/', |
| 82 | ).to(device) |
| 83 | self.rembg.eval() |
| 84 | |
| 85 | # diffusion models |
| 86 | pipeline = DiffusionPipeline.from_pretrained( |
| 87 | "sudo-ai/zero123plus-v1.1", |
| 88 | custom_pipeline="sudo-ai/zero123plus-pipeline", |
| 89 | torch_dtype=torch.float16, |
| 90 | cache_dir="ckpts/", |
| 91 | ) |
| 92 | pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config( |
| 93 | pipeline.scheduler.config, timestep_spacing='trailing' |
| 94 | ) |
| 95 | self.zero123plus_v11 = pipeline.to(device) |
| 96 | |
| 97 | pipeline = DiffusionPipeline.from_pretrained( |
| 98 | "sudo-ai/zero123plus-v1.2", |
| 99 | custom_pipeline="sudo-ai/zero123plus-pipeline", |
| 100 | torch_dtype=torch.float16, |
| 101 | cache_dir="ckpts/", |
| 102 | ) |
| 103 | pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config( |
| 104 | pipeline.scheduler.config, timestep_spacing='trailing' |
| 105 | ) |
| 106 | self.zero123plus_v12 = pipeline.to(device) |
| 107 | |
| 108 | pipeline = HunYuan3D_MVD_Std_Pipeline.from_pretrained( |
| 109 | './ckpts/Hunyuan3D-1/mvd_std', |
| 110 | torch_dtype=torch.float16, |
| 111 | use_safetensors=True, |
| 112 | ) |
| 113 | self.hunyuan3d_mvd_std = pipeline.to(device) |
| 114 | |
| 115 | # freesplatter |
| 116 | config_file = 'configs/freesplatter-object.yaml' |
| 117 | ckpt_path = hf_hub_download('TencentARC/FreeSplatter', repo_type='model', filename='freesplatter-object.safetensors', local_dir='./ckpts/FreeSplatter') |
| 118 | model = instantiate_from_config(OmegaConf.load(config_file).model) |
| 119 | state_dict = {} |
| 120 | with safe_open(ckpt_path, framework="pt", device="cpu") as f: |
| 121 | for key in f.keys(): |
| 122 | state_dict[key] = f.get_tensor(key) |
| 123 | model.load_state_dict(state_dict, strict=True) |
| 124 | self.freesplatter = model.eval().to(device) |
| 125 | |
| 126 | config_file = 'configs/freesplatter-object-2dgs.yaml' |
| 127 | ckpt_path = hf_hub_download('TencentARC/FreeSplatter', repo_type='model', filename='freesplatter-object-2dgs.safetensors', local_dir='./ckpts/FreeSplatter') |
| 128 | model = instantiate_from_config(OmegaConf.load(config_file).model) |
| 129 | state_dict = {} |
| 130 | with safe_open(ckpt_path, framework="pt", device="cpu") as f: |