| 209 | |
| 210 | |
| 211 | class DiffLocksInference(): |
| 212 | def __init__(self, path_ckpt_strandcodec, path_config_difflocks, path_ckpt_difflocks, path_ckpt_rgb2material=None, cfg_val=1.0, nr_iters_denoise=100, nr_chunks_decode=50): |
| 213 | super(DiffLocksInference, self).__init__() |
| 214 | |
| 215 | self.nr_chunks_decode_strands=nr_chunks_decode |
| 216 | self.nr_iters_denoise=nr_iters_denoise |
| 217 | self.cfg_val=cfg_val |
| 218 | |
| 219 | |
| 220 | self.mediapipe_img=Mediapipe(VisionRunningMode.IMAGE) |
| 221 | |
| 222 | #create dinov2 |
| 223 | image_size=770 #nearest images size that divides cleanly by patch size 14 |
| 224 | self.dinov2_latents_preprocessor = T.Compose([ |
| 225 | T.Resize(image_size, interpolation=T.InterpolationMode.BICUBIC), |
| 226 | T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), |
| 227 | ]) |
| 228 | self.dinov2_latents_model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14_reg') |
| 229 | self.dinov2_latents_model.cuda() |
| 230 | |
| 231 | #difflocks strand codec |
| 232 | self.strand_codec = StrandCodec(do_vae=False, |
| 233 | decode_type="dir", |
| 234 | scale_init=30.0, |
| 235 | nr_verts_per_strand=256, nr_values_to_decode=255, |
| 236 | dim_per_value_decoded=3).cuda() |
| 237 | self.strand_codec.load_state_dict(torch.load(path_ckpt_strandcodec)) |
| 238 | self.strand_codec.eval() |
| 239 | |
| 240 | |
| 241 | #difflocks diffusion |
| 242 | config = K.config.load_config(path_config_difflocks) |
| 243 | self.model_config = config['model'] |
| 244 | inner_model_ema = K.config.make_model(config).cuda() |
| 245 | inner_model_ema.eval() |
| 246 | model_ema = K.config.make_denoiser_wrapper(config)(inner_model_ema) |
| 247 | model_ema.eval() |
| 248 | #IMPORTANT set the dropout rate here for the condition to whatever you need, to make it either conditional or unconditional |
| 249 | # model_ema.inner_model.condition_dropout_rate=0.0 |
| 250 | # if state_path.exists() : |
| 251 | # state = json.load(open(state_path)) |
| 252 | # ckpt_path = state['latest_checkpoint'] |
| 253 | print(f'Resuming from {path_ckpt_difflocks}...') |
| 254 | ckpt = torch.load(path_ckpt_difflocks, map_location='cpu') |
| 255 | model_ema.inner_model.load_state_dict(ckpt['model_ema']) |
| 256 | del ckpt |
| 257 | self.model_ema=model_ema.cuda() |
| 258 | |
| 259 | #rgb2material |
| 260 | if path_ckpt_rgb2material is not None: |
| 261 | self.rgb2material = RGB2MaterialModel( |
| 262 | input_dim=1024, |
| 263 | out_dim=11, |
| 264 | hidden_dim=64).cuda() |
| 265 | self.rgb2material.load_state_dict(torch.load(path_ckpt_rgb2material)) |
| 266 | self.rgb2material.eval() |
| 267 | else: |
| 268 | self.rgb2material=None |