DiffFacto: Controllable Part-Based 3D Point Cloud Generation with Cross Diffusion, ICCV 2023
# model settings
cimle=True
cimle_cache_interval=50
model = dict(
    type='AnchorDiffAE',
    encoder=dict(
        type='PartEncoderForTransformerDecoder',
        encoder=dict(
            type='PointNetV2',
            zdim=256,
            point_dim=3,
            per_part_mlp=True,
        ),
        part_aligner=dict(
            type="PartAlignerTransformer",
            in_channels = 256,
            out_channels=6,
            n_class=4,
            d_head=32,
            depth=5,
            n_heads=8,
            dropout=0.,
            use_checkpoint=False,
            use_linear=True,
            class_cond=True,
            single_attn=True,
            add_class_cond=True,
            cimle=True,
            noise_scale=50,
            cond_noise_type=0
        ),
        n_class=4,
        kl_weight=0,
        fit_loss_type=4,
        fit_loss_weight=1.0,
        use_flow=True,
        latent_flow_depth=14, 
        latent_flow_hidden_dim=256,
        include_z=False,
        include_part_code=True,
        include_params=True,
        use_gt_params=False,
        kl_weight_annealing=False,
        gen=True,
        prior_var=1.0,
        # selective_noise_sampling=True
        
    ),
    diffusion=dict(
        type='AnchoredDiffusion',
        net = dict(
            type='TransformerNet',
            in_channels=3,
            out_channels=3,
            n_heads=8,
            d_head=16,
            depth=5,
            dropout=0.2,
            context_dim=256 + 6,
            n_class=4,
            class_cond=True,
            use_linear=True,
            cat_params_to_x=True,
            use_checkpoint=False,
            single_attn=True,
            cat_class_to_x=True,
        ),
        beta_1=1e-4,
        beta_T=.02,
        k=1.0,
        res=False,
        mode='linear',
        use_beta=False,
        rescale_timesteps=False,
        model_mean_type="epsilon",
        learn_variance=True,
        loss_type='mse',
        include_anchors=False,
        
        classifier_weight=1.,
        guidance=False,
        ddim_sampling=False,
        ddim_nsteps=25,
        ddim_discretize='quad',
        ddim_eta=1.
    ),
    sampler = dict(type='Uniform'),
    num_anchors=4,
    num_timesteps=100,
    npoints = 2048,
    
    gen=True,
    cimle=True,
    cimle_sample_num=1,
    ret_traj = True,
    ret_interval = 10,
    forward_sample=False,
    drift_anchors=False,
    interpolate=False,
    save_weights=False,
    save_dir="/mnt/disk3/wang/diffusion/anchorDIff/work_dirs/pn_aware_attn_both_feat_guidance_200T_no_augment/checkpoints"
)

dataset = dict(
    train=dict(
        type="ShapeNetSegPart",
        batch_size = 128,
        split='trainval',
        root='/orion/u/w4756677/datasets/diffFacto_data',
        npoints=2048,
        scale_mode='shape_unit',
        part_scale_mode='shape_canonical',
        eval_mode='gen',
        drop_last=False,
        clip=False,
        num_workers=4,
        class_choice='Airplane',
    ),
    val=dict(
        type="ShapeNetSegPart",
        batch_size= 128,
        split='test',
        root='/orion/u/w4756677/datasets/diffFacto_data',
        npoints=2048,
        shuffle=False,
        scale_mode='shape_unit',
        part_scale_mode='shape_canonical',
        eval_mode='gen',
        drop_last=True,
        num_workers=0,
        class_choice='Airplane',
        clip=False,
        save_only=True
    ),
)

optimizer = dict(type='Adamax', lr=0.002, weight_decay=0.)

scheduler = dict(
    type='LinearLR',
    start_lr=2e-3,
    end_lr = 1e-4,
    start_epoch=4000,
    end_epoch=8000,
)

logger = dict(
    type="RunLogger")

# when we the trained model from cshuan, image is rgb
resume_path="pretrained/airplane.pth"
save_num_batch = 1
max_epoch = 4000
eval_interval = 250
checkpoint_interval = 250
log_interval = 150
max_norm=10
# model_only=True
train_aligner=True
# eval_both=True
configs/gen_airplane.py
diffFacto/diffFacto
# model settings
cimle=True
cimle_cache_interval=50
model = dict(
    type='AnchorDiffAE',
    encoder=dict(
        type='PartEncoderForTransformerDecoder',
        encoder=dict(
            type='PointNetV2',
            zdim=256,
            point_dim=3,
            per_part_mlp=True,
        ),
        part_aligner=dict(
            type="PartAlignerTransformer",
            in_channels = 256,
            out_channels=6,
            n_class=4,
            d_head=32,
            depth=5,
            n_heads=8,
            dropout=0.,
            use_checkpoint=False,
            use_linear=True,
            class_cond=True,
            single_attn=True,
            add_class_cond=True,
            cimle=True,
            noise_scale=50,
            cond_noise_type=0,
            
        ),
        n_class=4,
        kl_weight=0,
        fit_loss_type=4,
        fit_loss_weight=1.0,
        use_flow=True,
        latent_flow_depth=14, 
        latent_flow_hidden_dim=256,
        include_z=False,
        include_part_code=True,
        include_params=True,
        use_gt_params=False,
        kl_weight_annealing=False,
        gen=True,
        prior_var=1.0,
        # selective_noise_sampling=True
    ),
    diffusion=dict(
        type='AnchoredDiffusion',
        net = dict(
            type='TransformerNet',
            in_channels=3,
            out_channels=3,
            n_heads=8,
            d_head=16,
            depth=5,
            dropout=0.2,
            context_dim=256 + 6,
            n_class=4,
            class_cond=True,
            use_linear=True,
            cat_params_to_x=True,
            use_checkpoint=False,
            single_attn=True,
            cat_class_to_x=True,
        ),
        beta_1=1e-4,
        beta_T=.02,
        k=1.0,
        res=False,
        mode='linear',
        use_beta=False,
        rescale_timesteps=False,
        model_mean_type="epsilon",
        learn_variance=True,
        loss_type='mse',
        include_anchors=False,
        
        classifier_weight=1.,
        guidance=False,
        ddim_sampling=False,
        ddim_nsteps=25,
        ddim_discretize='quad',
        ddim_eta=1.
    ),
    sampler = dict(type='Uniform'),
    num_anchors=4,
    num_timesteps=100,
    npoints = 8192,
    
    gen=True,
    cimle=True,
    cimle_sample_num=1,
    ret_traj = False,
    ret_interval = 10,
    forward_sample=False,
    drift_anchors=False,
    interpolate=False,
    save_weights=False,
)

dataset = dict(
    train=dict(
        type="ShapeNetSegPart",
        batch_size = 128,
        split='trainval',
        root='/orion/u/w4756677/datasets/diffFacto_data',
        npoints=2048,
        scale_mode='shape_unit',
        part_scale_mode='shape_canonical',
        eval_mode='gen',
        drop_last=False,
        num_workers=4,
        class_choice='Car',
    ),
    val=dict(
        type="ShapeNetSegPart",
        batch_size= 128,
        split='test',
        root='/orion/u/w4756677/datasets/diffFacto_data',
        npoints=2048,
        shuffle=False,
        scale_mode='shape_unit',
        part_scale_mode='shape_canonical',
        eval_mode='gen',
        drop_last=False,
        num_workers=0,
        class_choice='Car',
        save_only=True
    ),
)

optimizer = dict(type='Adam', lr=0.002, weight_decay=0.)

scheduler = dict(
    type='LinearLR',
    start_lr=2e-3,
    end_lr = 1e-4,
    start_epoch=4000,
    end_epoch=8000,
)

logger = dict(
    type="RunLogger")

# when we the trained model from cshuan, image is rgb
resume_path="pretrained/car.pth"
save_num_batch = 1
max_epoch = 4000
eval_interval = 250
checkpoint_interval = 250
log_interval = 50
max_norm=10
# model_only=True
train_aligner=True
# eval_both=True
configs/gen_car.py
diffFacto/diffFacto
# model settings
cimle=True
cimle_cache_interval=50
model = dict(
    type='AnchorDiffAE',
    encoder=dict(
        type='PartEncoderForTransformerDecoder',
        encoder=dict(
            type='PointNetV2',
            zdim=256,
            point_dim=3,
            per_part_mlp=True,
        ),
        part_aligner=dict(
            type="PartAlignerTransformer",
            in_channels = 256,
            out_channels=6,
            n_class=4,
            d_head=32,
            depth=5,
            n_heads=8,
            dropout=0.,
            use_checkpoint=False,
            use_linear=True,
            class_cond=True,
            single_attn=True,
            add_class_cond=True,
            cimle=True,
            noise_scale=100,
            cond_noise_type=0
        ),
        n_class=4,
        kl_weight=0,
        fit_loss_type=4,
        fit_loss_weight=1.0,
        use_flow=True,
        latent_flow_depth=14, 
        latent_flow_hidden_dim=256,
        include_z=False,
        include_part_code=True,
        include_params=True,
        use_gt_params=False,
        kl_weight_annealing=False,
        gen=True,
        prior_var=1.0,
        # selective_noise_sampling=True
    ),
    diffusion=dict(
        type='AnchoredDiffusion',
        net = dict(
            type='TransformerNet',
            in_channels=3,
            out_channels=3,
            n_heads=8,
            d_head=16,
            depth=5,
            dropout=0.2,
            context_dim=256 + 6,
            n_class=4,
            class_cond=True,
            use_linear=True,
            cat_params_to_x=True,
            use_checkpoint=False,
            single_attn=True,
            cat_class_to_x=True,
        ),
        beta_1=1e-4,
        beta_T=.02,
        k=1.0,
        res=False,
        mode='linear',
        use_beta=False,
        rescale_timesteps=False,
        model_mean_type="epsilon",
        learn_variance=True,
        loss_type='mse',
        include_anchors=False,
        
        classifier_weight=1.,
        guidance=False,
        ddim_sampling=False,
        ddim_nsteps=25,
        ddim_discretize='quad',
        ddim_eta=1.
    ),
    sampler = dict(type='Uniform'),
    num_anchors=4,
    num_timesteps=100,
    npoints = 2048,
    
    gen=True,
    cimle=True,
    cimle_sample_num=1,
    ret_traj = True,
    ret_interval = 10,
    forward_sample=False,
    drift_anchors=False,
    interpolate=False,
    save_weights=False,
)

dataset = dict(
    train=dict(
        type="ShapeNetSegPart",
        batch_size = 128,
        split='trainval',
        root='/orion/u/w4756677/datasets/diffFacto_data',
        npoints=2048,
        scale_mode='shape_unit',
        part_scale_mode='shape_canonical',
        eval_mode='gen',
        drop_last=False,
        clip=False,
        num_workers=4,
        class_choice='Chair',
    ),
    val=dict(
        type="ShapeNetSegPart",
        batch_size= 128,
        split='test',
        root='/orion/u/w4756677/datasets/diffFacto_data',
        npoints=2048,
        shuffle=False,
        scale_mode='shape_unit',
        part_scale_mode='shape_canonical',
        eval_mode='gen_part',
        drop_last=False,
        clip=False,
        num_workers=0,
        class_choice='Chair',
        save_only=True
    ),
)

optimizer = dict(type='Adam', lr=0.002, weight_decay=0.)

scheduler = dict(
    type='LinearLR',
    start_lr=2e-3,
    end_lr = 1e-4,
    start_epoch=4000,
    end_epoch=8000,
)

logger = dict(
    type="RunLogger")

# when we the trained model from cshuan, image is rgb
save_num_batch = 1
max_epoch = 4000
eval_interval = 250
checkpoint_interval = 250
log_interval = 50
max_norm=10
# model_only=True
train_aligner=True
resume_path="pretrained/chair.pth"
configs/gen_chair.py
diffFacto/diffFacto
# model settings
cimle=True
cimle_cache_interval=50
model = dict(
    type='AnchorDiffAE',
    encoder=dict(
        type='PartEncoderForTransformerDecoder',
        encoder=dict(
            type='PointNetV2',
            zdim=256,
            point_dim=3,
            per_part_mlp=True,
        ),
        part_aligner=dict(
            type="PartAlignerTransformer",
            in_channels = 256,
            out_channels=6,
            n_class=4,
            d_head=32,
            depth=5,
            n_heads=8,
            dropout=0.,
            use_checkpoint=False,
            use_linear=True,
            class_cond=True,
            single_attn=True,
            add_class_cond=True,
            cimle=True,
            noise_scale=10,
            cond_noise_type=0
        ),
        n_class=4,
        kl_weight=0,
        fit_loss_type=4,
        fit_loss_weight=1.0,
        use_flow=True,
        latent_flow_depth=14, 
        latent_flow_hidden_dim=256,
        include_z=False,
        include_part_code=True,
        include_params=True,
        use_gt_params=False,
        kl_weight_annealing=False,
        gen=True,
        prior_var=1.0
    ),
    diffusion=dict(
        type='AnchoredDiffusion',
        net = dict(
            type='TransformerNet',
            in_channels=3,
            out_channels=3,
            n_heads=8,
            d_head=16,
            depth=5,
            dropout=0.2,
            context_dim=256 + 6,
            n_class=4,
            class_cond=True,
            use_linear=True,
            cat_params_to_x=True,
            use_checkpoint=False,
            single_attn=True,
            cat_class_to_x=True,
        ),
        beta_1=1e-4,
        beta_T=.02,
        k=1.0,
        res=False,
        mode='linear',
        use_beta=False,
        rescale_timesteps=False,
        model_mean_type="epsilon",
        learn_variance=True,
        loss_type='mse',
        include_anchors=False,
        
        classifier_weight=1.,
        guidance=False,
        ddim_sampling=False,
        ddim_nsteps=25,
        ddim_discretize='quad',
        ddim_eta=1.
    ),
    sampler = dict(type='Uniform'),
    num_anchors=4,
    num_timesteps=100,
    npoints = 2048,
    
    gen=True,
    cimle=True,
    cimle_sample_num=1,
    ret_traj = True,
    ret_interval = 10,
    forward_sample=False,
    drift_anchors=False,
    interpolate=False,
    save_weights=False,
)

dataset = dict(
    train=dict(
        type="ShapeNetSegPart",
        batch_size = 128,
        split='trainval',
        root='/orion/u/w4756677/datasets/diffFacto_data',
        npoints=2048,
        scale_mode='shape_unit',
        part_scale_mode='shape_canonical',
        eval_mode='gen',
        drop_last=False,
        clip=False,
        num_workers=4,
        class_choice='Lamp',
    ),
    val=dict(
        type="ShapeNetSegPart",
        batch_size= 128,
        split='test',
        root='/orion/u/w4756677/datasets/diffFacto_data',
        npoints=2048,
        shuffle=False,
        scale_mode='shape_unit',
        part_scale_mode='shape_canonical',
        eval_mode='gen',
        drop_last=False,
        clip=False,
        num_workers=0,
        class_choice='Lamp',
        save_only=True
    ),
)

optimizer = dict(type='Adam', lr=0.002, weight_decay=0.)

scheduler = dict(
    type='LinearLR',
    start_lr=2e-3,
    end_lr = 1e-4,
    start_epoch=4000,
    end_epoch=8000,
)

logger = dict(
    type="RunLogger")

# when we the trained model from cshuan, image is rgb
save_num_batch = 1
max_epoch = 4000
eval_interval = 250
checkpoint_interval = 250
log_interval = 50
max_norm=10
# model_only=True
train_aligner=True
resume_path="pretrained/lamp.pth"
configs/gen_lamp.py
diffFacto/diffFacto
# model settings
model = dict(
    type='AnchorDiffAE',
    encoder=dict(
        type='PartEncoderForTransformerDecoder',
        encoder=dict(
            type='PointNetV2',
            zdim=256,
            point_dim=3,
            per_part_mlp=True,
        ),
        n_class=4,
        kl_weight=5e-4,
        fit_loss_type=4,
        fit_loss_weight=1.0,
        use_flow=True,
        latent_flow_depth=14, 
        latent_flow_hidden_dim=256,
        include_z=False,
        include_part_code=True,
        include_params=True,
        use_gt_params=True,
        kl_weight_annealing=False,
        min_kl_weight=1e-7,
        kl_weight_annealing_end_epoch=4000,
        gen=True,
        prior_var=1.0
    ),
    diffusion=dict(
        type='AnchoredDiffusion',
        net = dict(
            type='TransformerNet',
            in_channels=3,
            out_channels=3,
            n_heads=8,
            d_head=16,
            depth=5,
            dropout=0.2,
            context_dim=256 + 6,
            n_class=4,
            class_cond=True,
            use_linear=True,
            cat_params_to_x=True,
            use_checkpoint=False,
            single_attn=True,
            cat_class_to_x=True,
        ),
        beta_1=1e-4,
        beta_T=.02,
        k=1.0,
        res=False,
        mode='linear',
        use_beta=False,
        rescale_timesteps=False,
        model_mean_type="epsilon",
        learn_variance=True,
        loss_type='mse',
        include_anchors=False,
        
        classifier_weight=1.,
        guidance=False,
        ddim_sampling=False,
        ddim_nsteps=25,
        ddim_discretize='quad',
        ddim_eta=1.
    ),
    sampler = dict(type='Uniform'),
    num_anchors=4,
    num_timesteps=100,
    npoints = 2048,
    
    ret_traj = False,
    ret_interval = 10,
    forward_sample=False,
    drift_anchors=False,
    interpolate=False,
    save_weights=False,
    save_dir="/mnt/disk3/wang/diffusion/anchorDIff/work_dirs/pn_aware_attn_both_feat_guidance_200T_no_augment/checkpoints"
)

dataset = dict(
    train=dict(
        type="ShapeNetSegPart",
        batch_size = 128,
        split='trainval',
        root='/orion/u/w4756677/datasets/diffFacto_data',
        npoints=2048,
        scale_mode='shape_unit',
        part_scale_mode='shape_canonical',
        eval_mode='ae',
        clip=False,
        num_workers=4,
        class_choice='Chair',
    ),
    val=dict(
        type="ShapeNetSegPart",
        batch_size= 64,
        split='test',
        root='/orion/u/w4756677/datasets/diffFacto_data',
        npoints=2048,
        shuffle=False,
        scale_mode='shape_unit',
        part_scale_mode='shape_canonical',
        clip=False,
        eval_mode='ae',
        num_workers=0,
        class_choice='Chair',
        save_only=True
    ),
)

optimizer = dict(type='Adam', lr=0.002, weight_decay=0.)

scheduler = dict(
    type='LinearLR',
    start_lr=2e-3,
    end_lr = 1e-4,
    start_epoch=4000,
    end_epoch=8000,
)


logger = dict(
    type="RunLogger")

# when we the trained model from cshuan, image is rgb
save_num_batch = 1
max_epoch = 8000
eval_interval = 500
checkpoint_interval = 500
log_interval = 50
max_norm=10
# eval_both=True
configs/train_chair_stage1.py
diffFacto/diffFacto
# model settings
cimle=True
cimle_cache_interval=50
model = dict(
    type='AnchorDiffAE',
    encoder=dict(
        type='PartEncoderForTransformerDecoder',
        encoder=dict(
            type='PointNetV2',
            zdim=256,
            point_dim=3,
            per_part_mlp=True,
        ),
        part_aligner=dict(
            type="PartAlignerTransformer",
            in_channels = 256,
            out_channels=6,
            n_class=4,
            d_head=32,
            depth=5,
            n_heads=8,
            dropout=0.,
            use_checkpoint=False,
            use_linear=True,
            class_cond=True,
            single_attn=True,
            add_class_cond=True,
            cimle=True,
            noise_scale=50,
            cond_noise_type=0
        ),
        n_class=4,
        kl_weight=0,
        fit_loss_type=4,
        fit_loss_weight=1.0,
        use_flow=True,
        latent_flow_depth=14, 
        latent_flow_hidden_dim=256,
        include_z=False,
        include_part_code=True,
        include_params=True,
        use_gt_params=False,
        kl_weight_annealing=False,
        gen=True,
        prior_var=1.0
    ),
    diffusion=dict(
        type='AnchoredDiffusion',
        net = dict(
            type='TransformerNet',
            in_channels=3,
            out_channels=3,
            n_heads=8,
            d_head=16,
            depth=5,
            dropout=0.2,
            context_dim=256 + 6,
            n_class=4,
            class_cond=True,
            use_linear=True,
            cat_params_to_x=True,
            use_checkpoint=False,
            single_attn=True,
            cat_class_to_x=True,
        ),
        beta_1=1e-4,
        beta_T=.02,
        k=1.0,
        res=False,
        mode='linear',
        use_beta=False,
        rescale_timesteps=False,
        model_mean_type="epsilon",
        learn_variance=True,
        loss_type='mse',
        include_anchors=False,
        
        classifier_weight=1.,
        guidance=False,
        ddim_sampling=False,
        ddim_nsteps=25,
        ddim_discretize='quad',
        ddim_eta=1.
    ),
    sampler = dict(type='Uniform'),
    num_anchors=4,
    num_timesteps=100,
    npoints = 2048,
    
    gen=True,
    cimle=True,
    # cimle_sample_num=1,
    ret_traj = True,
    ret_interval = 1,
    forward_sample=False,
    drift_anchors=False,
    interpolate=False,
    save_weights=False,
    save_dir="/mnt/disk3/wang/diffusion/anchorDIff/work_dirs/pn_aware_attn_both_feat_guidance_200T_no_augment/checkpoints"
)

dataset = dict(
    train=dict(
        type="ShapeNetSegPart",
        batch_size = 128,
        split='trainval',
        root='/orion/u/w4756677/datasets/diffFacto_data',
        npoints=2048,
        scale_mode='shape_unit',
        part_scale_mode='shape_canonical',
        eval_mode='gen',
        drop_last=False,
        clip=False,
        num_workers=4,
        class_choice='Chair',
    ),
    val=dict(
        type="ShapeNetSegPart",
        batch_size= 32,
        split='test',
        root='/orion/u/w4756677/datasets/diffFacto_data',
        npoints=2048,
        shuffle=False,
        scale_mode='shape_unit',
        part_scale_mode='shape_canonical',
        eval_mode='gen',
        drop_last=False,
        clip=False,
        num_workers=0,
        class_choice='Chair',
        save_only=True
    ),
)

optimizer = dict(type='Adam', lr=0.002, weight_decay=0.)

scheduler = dict(
    type='LinearLR',
    start_lr=2e-3,
    end_lr = 1e-4,
    start_epoch=4000,
    end_epoch=8000,
)

logger = dict(
    type="RunLogger")

# when we the trained model from cshuan, image is rgb
save_num_batch = 1000
max_epoch = 4000
eval_interval = 250
checkpoint_interval = 250
log_interval = 50
max_norm=10
model_only=True
train_aligner=True
resume_path="work_dirs/anchordiff_exp624_debugged/checkpoints/ckpt_8000.pth"
configs/train_chair_stage2.py
diffFacto/diffFacto
from . import optimizers
from . import runner
from . import samplers
from . import models
from . import datasets
from . import metrics

python/difffacto/__init__.py
diffFacto/diffFacto
# reference maskrcnn_benchmark default
from collections import OrderedDict
from difffacto.utils.misc import check_file
import os
import yaml
import copy
from importlib import import_module
import sys
import inspect

__all__ = ["get_cfg","init_cfg","save_cfg","print_cfg"]
BASE_KEY = "_base_"
COVER_KEY = "_cover_"
IGNORE_KEY = "_ignore_"

class Config(OrderedDict):
    def __init__(self, *args):
        super().__init__()
        if len(args)==1:
            self.load_from_file(args[0])
        else:
            assert(len(args)==0)

    def __getattr__(self, name):
        if name in self:
            return self[name]
        return None

    def __setattr__(self, name, value):
        self[name] = value

    # Load config file to a Dict without considering BASE_KEY
    @staticmethod
    def _load_dict_from_file_no_base(filename):
        if check_file(filename,ext=[".yaml"]):
            with open(filename, "r") as f:
                cfg = yaml.safe_load(f.read())
        elif check_file(filename,ext=[".py"]):
            f_dir = os.path.dirname(filename)
            f_name = os.path.basename(filename)
            module_name = f_name[:-3]

            # temp_module_name = osp.splitext(temp_config_name)[0]
            sys.path.insert(0, f_dir)
            # Config._validate_py_syntax(filename)
            mod = import_module(module_name)
            sys.path.pop(0)
            cfg = {
                name: value
                for name, value in mod.__dict__.items()
                if not name.startswith('__')
            }
            # delete imported module
            del sys.modules[module_name]
        else:
            assert(False), "unsupported config type."
        return cfg

    # Load config file to a Dict with considering BASE_KEY
    @staticmethod
    def _load_dict_from_file(filename):
        cfg = Config._load_dict_from_file_no_base(filename)
        cfg_dir = os.path.dirname(filename)
        if (BASE_KEY in cfg):
            if (isinstance(cfg[BASE_KEY], list)):
                base_filenames = cfg[BASE_KEY]
            else:
                assert(isinstance(cfg[BASE_KEY], str))
                base_filenames = [cfg[BASE_KEY]]
            cfg_base = {}
            for bfn in base_filenames:
                Config.merge_dict_b2a(cfg_base, Config._load_dict_from_file(os.path.join(cfg_dir, bfn)))
            cfg.pop(BASE_KEY)
            Config.merge_dict_b2a(cfg_base, cfg)
            cfg = cfg_base
        return cfg

    @staticmethod
    def merge_dict_b2a(a, b):
        def clear_cover_key(a):
            if (not isinstance(a, dict)):
                return a
            out = copy.deepcopy(a)
            if (COVER_KEY in out):
                out.pop(COVER_KEY)
            for k, v in out.items():
                out[k] = clear_cover_key(v)
            return out

        assert(isinstance(a, dict))
        assert(isinstance(b, dict))
        if (COVER_KEY in b):
            a.clear()
            temp = clear_cover_key(copy.deepcopy(b))
            a.update(temp)
            return
        for k, v in b.items():
            if ((not(k in a)) or (isinstance(v, dict) and v.get(COVER_KEY, False)) or (not isinstance(v, dict)) or (not isinstance(a[k], dict))):
                a[k] = clear_cover_key(copy.deepcopy(v))
            else:
                Config.merge_dict_b2a(a[k], v)

    def load_from_file(self, filename):
        cfg = Config._load_dict_from_file(filename)
        self.clear()
        self.update(self.dfs(cfg))
        if (self.name is None):
            self.name = os.path.splitext(os.path.basename(filename))[0]
        if self.work_dir is None:
            self.work_dir = f"work_dirs/{self.name}"
    
    def dfs(self, cfg_other):
        if isinstance(cfg_other,dict):
            now_cfg = Config()
            for k,d in cfg_other.items():
                if (inspect.ismodule(d)):
                    continue
                now_cfg[k]=self.dfs(d)
        elif isinstance(cfg_other,list):
            now_cfg = [self.dfs(d) for d in cfg_other if (not inspect.ismodule(d))]
        else:
            now_cfg = copy.deepcopy(cfg_other)
        return now_cfg
    
    def dump(self):
        """convert Config to dict"""
        now = dict()
        for k,d in self.items():
            if isinstance(d,Config):
                d = d.dump()
            if isinstance(d,list):
                d = [dd.dump() if isinstance(dd,Config) else dd for dd in d]
            now[k]=d
        return now

'''
TODO: _cover_ between sibling nodes 

Configs forms a tree like structure through '_base_', while '_cover_' will only work on child nodes, and will not work between sibling nodes.

Maybe we need a cover works on sibling nodes, like:
    _base_: ['a.yaml', '@b.yaml']
means each attribute with a depth of 1 in b.yaml will have the attribute '_cover' before merge into a.yaml, and '@@b.yaml' means a depth of 2, etc.
'''

_cfg = Config()

def init_cfg(filename):
    print("Loading config from: ",filename)
    _cfg.load_from_file(filename)

def get_cfg():
    return _cfg

def update_cfg(**kwargs):
    _cfg.update(kwargs)

def save_cfg(save_file):
    with open(save_file,"w") as f:
        f.write(yaml.dump(_cfg.dump()))

def save_args(save_file, args):
    with open(save_file,"w") as f:
        f.write(yaml.dump(args.__dict__))



def print_cfg():
    data  =  yaml.dump(_cfg.dump())
    # TODO: data keys are not sorted
    print(data) 


if __name__ == '__main__':
    path = '/mnt/disk1/wang/anchorDIff/scripts/exp1.py'
    cfg = init_cfg(path)
    print_cfg()
python/difffacto/config/config.py
diffFacto/diffFacto
from .shapenet55 import ShapeNet
from .shapenet_seg import ShapeNetSeg, ShapeNetSegPart, _ShapeNetSegParts, _ShapeNetSeg
from .shapenet_sup_seg import ShapeNetSegSuperSegment
from .shapenet_sup_seg_partglot import ShapeNetSegSuperSegmentParglot
# from .partglot import PartglotDataset
from .partnet_dataset import Partnet

python/difffacto/datasets/__init__.py
diffFacto/diffFacto
import os
import torch
import torch.nn.functional as F
from difffacto.utils.registry import DATASETS
from difffacto.utils.misc import fps
from torch.utils.data import Dataset
import numpy as np
from tqdm.auto import tqdm
from .dataset_utils import pc_norm, shapenet_part_normal_cat_to_id, build_loss_funcs, DataLoaderWrapper, DataLoaderWrapperOne
from .evaluation_utils import EMD_CD, compute_all_metrics
import json
import pickle
import random

@DATASETS.register_module()
class CustomDataset(Dataset):

    def __init__(
        self,
        data_dir = '../datasets/shapenetcore_partanno_segmentation_benchmark_v0_normal',
        npoints=2500, 
        scale_mode='shape_unit',
        part_scale_mode='shape_canonical',
        eval_mode='ae',
        clip=True,
        n_class=4
    ):
        self.scale_mode=scale_mode 
        self.part_scale_mode=part_scale_mode 
        self.npoints = npoints # Sampling points
        self.root = data_dir # File root path
        self.eval_mode=eval_mode
        self.n_class=n_class
        self.clip=clip

        d = pickle.load(open(self.root, 'rb'))
        self.data = d['pred']
        self.label =d['pred_seg_mask']

    def __getitem__(self, index):
        # resample
        point_set = self.data[index] # Sample by index
        label = self.label[index] # Sample by index
        
        shifts = np.zeros((self.n_class, 3))
        scales = np.ones((self.n_class, 3))
        present = torch.zeros(self.n_class)

        choice = np.random.choice(point_set.shape[0], self.npoints, replace=True)
        point_set = point_set[choice]
        label = label[choice]
        point_set, shift, scale = pc_norm(point_set, self.scale_mode, to_torch=False) 
        out = np.zeros_like(point_set)
        for i in range(self.n_class):
            idx = label == i
            if idx.sum() >= 10: # if less than 10 points we ignore
                part_point_set = point_set[idx]
                part_point_set_mean = part_point_set.mean(0)
                part_point_set_std = part_point_set.std(0)
                if np.any(part_point_set_std == 0.0):
                    present[i] = 0
                    part_point_set_std = np.ones_like(part_point_set_std)
                else:
                    present[i] = 1
                print(self.part_scale_mode)
                shifted_chosen, pshift, pscale = pc_norm(part_point_set, self.part_scale_mode, to_torch=False, clip=self.clip)
                print(pshift, pscale)
                shifts[i] = pshift[0]
                scales[i] = pscale[0]
                out[idx] = shifted_chosen
            elif np.any(idx):
                # we change its label to their nearest neighbors
                part_point = point_set[idx] # n, 3
                rest_point = point_set[~idx]
                rest_seg = label[~idx]
                dist = (part_point[:, None] - rest_point[None]) ** 2
                label_idx = dist.sum(-1).argmin(1)
                label[idx] = np.take_along_axis(rest_seg, label_idx, axis=0)
        



        
        return {
            'input':out, 
            'seg_mask':label,
            'shift':shift,
            'scale':scale, 
            'id':index,
            } # pointset is point cloud data, cls has 16 categories, and seg is a small category corresponding to different points in the data

    def __len__(self):
        return self.data.shape[0]

    

    def evaluate(self, results, save_num_batch, device):
        save_dict = dict()
        preds = []
        refs = []
        for idx, pred_dict in enumerate(tqdm(results)):
            shift = pred_dict['shift']
            scale = pred_dict['scale']
            del pred_dict['shift']
            del pred_dict['scale']
            pred = pred_dict['pred']
            ref = pred_dict['input_ref']
            if pred.shape[1] > 2048:
                pred = fps(pred.to('cuda').contiguous(), 2048).cpu()
            if ref.shape[1] > 2048:
                ref = fps(ref.to('cuda').contiguous(), 2048).cpu()
            if self.eval_mode == 'ae':
                pred = pred * scale + shift
                ref = ref * scale + shift
            else:
                pred_max = pred.max(1)[0].reshape(-1, 1, 3) # (1, 3)
                pred_min = pred.min(1)[0].reshape(-1, 1, 3) # (1, 3)
                pred_shift = ((pred_min + pred_max) / 2).reshape(-1, 1, 3)
                pred_scale = (pred_max - pred_min).max(-1)[0].reshape(-1, 1, 1) / 2
                pred = (pred - pred_shift) / pred_scale
                ref_max = ref.max(1)[0].reshape(-1, 1, 3) # (1, 3)
                ref_min = ref.min(1)[0].reshape(-1, 1, 3) # (1, 3)
                ref_shift = ((ref_min + ref_max) / 2).reshape(-1, 1, 3)
                ref_scale = (ref_max - ref_min).max(-1)[0].reshape(-1, 1, 1) / 2
                ref = (ref - ref_shift) / ref_scale
                if self.using_whole_chair_only:
                    present = pred_dict['present'][:, :3].sum(1) == 3 
                    pred = pred[present]
                    ref = ref[present]
            preds.append(pred)
            refs.append(ref)
            pred_dict = {k: v * scale + shift if v.shape[-1] == 3 else v for k, v in pred_dict.items()}    
                     
            if idx < save_num_batch:
                for k, v in pred_dict.items():
                    if k not in save_dict.keys():
                        save_dict[k] = []
                    save_dict[k].append(v.detach().cpu().numpy())
        preds = torch.cat(preds, dim=0).cuda()
        refs = torch.cat(refs, dim=0).cuda()
        if self.eval_mode == 'ae':
            metrics = EMD_CD(preds, refs, 32)
        elif self.eval_mode == 'gen':
            # metrics = dict(l=torch.zeros(1))
            print(preds.shape, refs.shape)
            metrics = compute_all_metrics(preds, refs, 32)
        ssave_dict = dict()
        for k, v in save_dict.items():
            ssave_dict[k] = np.concatenate(v, axis=0)
        return ssave_dict, metrics

python/difffacto/datasets/custom.py
diffFacto/diffFacto
import numpy as np
import torch
from torch.utils.data import DataLoader
from difffacto.utils.misc import separate_point_cloud
from difffacto.utils.registry import build_from_cfg, METRICS
import random
from plyfile import  PlyData

def load_ply_withNormals(file_name ):
    
    ply_data = PlyData.read(file_name)
    vertices = ply_data['vertex']
    
    points  = np.vstack([vertices['x'],  vertices['y'],  vertices['z']]).T

    return points

class IterWrapper:
    def __init__(self, loader):
        self.loader = loader
        self.loader_iter = iter(loader.loader)
    def __next__(self):
        d = self.loader_iter.next()
        partials, _ = separate_point_cloud(d['ref'].cuda(), self.loader.npoints, self.loader.crop_range)
        d['partial'] = partials 
        return d

class DataLoaderWrapper:
    def __init__(self, train_dataset, batch_size, num_workers, drop_last, crop_range, npoints, sampler=None):
        self.loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=num_workers, drop_last=drop_last, sampler=sampler, worker_init_fn=worker_init_fn)
        self.batch_size = self.loader.batch_size
        self.crop_range=crop_range
        self.npoints=npoints
    def __iter__(self):
        return IterWrapper(self)

    def __len__(self):
        return len(self.loader)

    def evaluate(self, results, save_num_batch, device):
        return self.loader.dataset.evaluate(results, save_num_batch, device)


class DataLoaderWrapperOne(DataLoader):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


    def evaluate(self, results, save_num_batch, device):
        return self.dataset.evaluate(results, save_num_batch, device)

def worker_init_fn(worker_id):
    np.random.seed(np.random.get_state()[1][0] + worker_id)

def pc_norm(pc, scale_mode, to_torch=True, stats=None, clip=True):
        if scale_mode == 'global_unit':
            shift = stats['mean'].reshape(1, 3)
            scale = stats['std'].reshape(1, 1)
        elif scale_mode == 'shape_unit':
            shift = pc.mean(0).reshape(1, 3)
            scale = pc.flatten().std().reshape(1, 1)
        elif scale_mode == 'shape_canonical':
            shift = pc.mean(0).reshape(1, 3)
            scale = pc.std(0).reshape(1, 3)
            if clip:
                scale = scale.clip(1e-2, 1)
            if np.any(scale == 0.0):
                scale[0, scale[0] == 0.0] = 1.0
        elif scale_mode == 'shape_canonical_bbox':
            pc_max = pc.max(0).reshape(1,3) # (1, 3)
            pc_min = pc.min(0).reshape(1,3) # (1, 3)
            shift = ((pc_min + pc_max) / 2).reshape(1, 3)
            scale = (pc_max - pc_min).reshape(1, 3) / 2
            if clip:
                scale = scale.clip(1e-2, 1)
            if np.any(scale == 0.0):
                scale[0, scale[0] == 0.0] = 1.0
        elif scale_mode == 'shape_half':
            shift = pc.mean(0).reshape(1, 3)
            scale = pc.flatten().std().reshape(1, 1) / (0.5)
        elif scale_mode == 'shape_34':
            shift = pc.mean(0).reshape(1, 3)
            scale = pc.flatten().std().reshape(1, 1) / (0.75)
        elif scale_mode == 'shape_bbox':
            pc_max = pc.max(0).reshape(1,3) # (1, 3)
            pc_min = pc.min(0).reshape(1,3) # (1, 3)
            shift = ((pc_min + pc_max) / 2).reshape(1, 3)
            scale = (pc_max - pc_min).max().reshape(1, 1) / 2
        else:
            shift = np.zeros([1, 3])
            scale = np.ones([1, 1])
        pc = (pc - shift) / scale
        if to_torch:
            pc, shift, scale = map(lambda t: torch.from_numpy(t).float(), [pc, shift, scale])
        return pc, shift, scale

def build_loss_funcs(loss_dict):
    loss = dict()
    if isinstance(loss_dict, dict):
        loss[loss_dict['type']] = build_from_cfg(loss_dict, METRICS)
    elif isinstance(loss_dict, list):
        for val_l in loss_dict:
            loss[val_l['type']] = build_from_cfg(val_l, METRICS)
    else:
        loss = loss_dict
    return loss

def augment(input, seg_mask, n_class, vertical_only=True, shift_only=False):
    out = torch.zeros_like(input)
    scales = []
    shifts = []
    for i in range(n_class):
        part = input[seg_mask == i]
        rand_scale = torch.rand(1,3) / 2 + 0.7 if not shift_only else torch.ones(1,3)
        rand_shift = torch.rand(1, 3) - 0.5
        if vertical_only:
            rand_shift[:, [0,2]] = 0.
        scales.append(rand_scale)
        shifts.append(rand_shift)
        part = (part + rand_shift) * rand_scale
        out[seg_mask == i] += part
    scales = torch.cat(scales, dim=0)
    shifts = torch.cat(shifts, dim=0)
    return out, scales, shifts

shapenet55_id_to_cat={
    "02691156": "airplane", "02747177": "trash bin", 
    "02773838": "bag", "02801938": "basket", 
    "02808440": "bathtub", "02818832": "bed", 
    "02828884": "bench", "02843684": "birdhouse", 
    "02871439": "bookshelf", "02876657": "bottle", 
    "02880940": "bowl", "02924116": "bus", 
    "02933112": "cabinet", "02942699": "camera", 
    "02946921": "can", "02954340": "cap", 
    "02958343": "car", "02992529": "cellphone", 
    "03001627": "chair", "03046257": "clock", 
    "03085013": "keyboard", "03207941": "dishwasher", 
    "03211117": "display", "03261776": "earphone", 
    "03325088": "faucet", "03337140": "file cabinet",
    "03467517": "guitar", "03513137": "helmet",
    "03593526": "jar", "03624134": "knife", 
    "03636649": "lamp", "03642806": "laptop", 
    "03691459": "loudspeaker", "03710193": "mailbox", 
    "03759954": "microphone", "03761084": "microwaves", 
    "03790512": "motorbike", "03797390": "mug", 
    "03928116": "piano", "03938244": "pillow", 
    "03948459": "pistol", "03991062": "flowerpot", 
    "04004475": "printer", "04074963": "remote", 
    "04090263": "rifle", "04099429": "rocket", 
    "04225987": "skateboard", "04256520": "sofa", 
    "04330267": "stove", "04379243": "table", 
    "04401088": "telephone", "04460130": "tower", 
    "04468005": "train", "04530566": "watercraft", 
    "04554684": "washer"
    }


shapenet_part_normal_cat_to_id={
    'Airplane': '02691156', 'Bag': '02773838', 
    'Cap': '02954340', 'Car': '02958343', 
    'Chair': '03001627', 'Earphone': '03261776', 
    'Guitar': '03467517', 'Knife': '03624134', 
    'Lamp': '03636649', 'Laptop': '03642806', 
    'Motorbike': '03790512', 'Mug': '03797390', 
    'Pistol': '03948459', 'Rocket': '04099429', 
    'Skateboard': '04225987', 'Table': '04379243'
    }


shapenet_chair_part_distribution={
    '1110': 0.7209302325581395, 
    '1111': 0.2630199803471995, 
    '1101': 0.009498853586636095, 
    '1001': 0.00032754667540124465, 
    '1100': 0.002947920078611202, 
    '0111': 0.0013101867016049786, 
    '0110': 0.0016377333770062235, 
    '1011': 0.00032754667540124465
}
python/difffacto/datasets/dataset_utils.py
diffFacto/diffFacto
import torch
import numpy as np
import torch.nn.functional as F
import warnings
from scipy.stats import entropy
from sklearn.neighbors import NearestNeighbors
from numpy.linalg import norm
from tqdm.auto import tqdm
# Import CUDA version of approximate EMD,
# from https://github.com/zekunhao1995/pcgan-pytorch/
from difffacto.metrics.chamfer_dist import ChamferDistanceL2_split
from difffacto.metrics.emd.emd_module import EMD
from difffacto.utils.misc import fps
from functools import partial
from .iou import get_3d_box, box3d_iou

def distChamferCUDA(x, y):
    func = ChamferDistanceL2_split(reduce=False)
    return func(x, y)



def part_chamfer(n_class, A, B, accelerated=False):
    dist = []
    for i in range(n_class):
        _A = A.get(i, None)
        _B = B.get(i, None)
        if (_A is not None) != (_B is not None):
            return torch.ones(1) * float('inf')
        if _A is None:
            continue
        _A_points = torch.rand([512, 3]) * (_A[1] - _A[0]) + _A[0]
        _B_points = torch.rand([512, 3]) * (_B[1] - _B[0]) + _B[0]
        if accelerated:
            _dl, _dr = distChamferCUDA(_A_points[None].cuda(), _B_points[None].cuda())
        else:
            _dl, _dr = distChamfer(_A_points[None], _B_points[None])
        dist.append(_dl.mean() + _dr.mean())
    dist = torch.tensor(dist).mean()
    return dist 

def part_l2(n_class, A, B, accelerated=False):
    dist = []
    for i in range(n_class):
        _A = A.get(i, None)
        _B = B.get(i, None)
        if (_A is not None) != (_B is not None):
            return torch.ones(1) * float('inf')
        if _A is None:
            continue
        A_max = _A[1][0]
        A_min = _A[0][0]
        B_max = _B[1][0]
        B_min = _B[0][0]
        scale_A = (A_max - A_min) / 2.0 
        scale_B = (B_max - B_min) / 2.0
        shift_A = (A_max + A_min) / 2.0
        shift_B = (B_max + B_min) / 2.0 
        d = F.mse_loss(torch.cat([scale_A, shift_A], dim=0), torch.cat([scale_B, shift_B], dim=0))
        dist.append(d)
    dist = torch.tensor(dist).mean()
    return dist 

def part_miou(n_class, A, B, accelerated=True):
    dist = []
    for i in range(n_class):
        _A = A.get(i, None)
        _B = B.get(i, None)
        if (_A is not None) != (_B is not None):
            return torch.ones(1) * float('inf')
        if _A is None:
            continue
        A_max = _A[1][0].numpy()
        A_min = _A[0][0].numpy()
        B_max = _B[1][0].numpy()
        B_min = _B[0][0].numpy()
        A_box  = get_3d_box(A_max - A_min, 0, (A_max + A_min) / 2.0) 
        B_box  = get_3d_box(B_max - B_min, 0, (B_max + B_min) / 2.0) 
        iou_3d, _ = box3d_iou(A_box, B_box)
        dist.append(iou_3d)
    dist = torch.tensor(dist).mean()
    return 1.0 - dist 

def emd_approx(sample, ref):
    B, N, N_ref = sample.size(0), sample.size(1), ref.size(1)
    assert N == N_ref, "Not sure what would EMD do in this case"
    emd = EMD(0.002, 10000, True)
    cost = emd(sample, ref)  # (B,)
    return cost


# Borrow from https://github.com/ThibaultGROUEIX/AtlasNet
def distChamfer(a, b):
    x, y = a, b
    bs, num_points, points_dim = x.size()
    xx = torch.bmm(x, x.transpose(2, 1))
    yy = torch.bmm(y, y.transpose(2, 1))
    zz = torch.bmm(x, y.transpose(2, 1))
    diag_ind = torch.arange(0, num_points).to(a).long()
    rx = xx[:, diag_ind, diag_ind].unsqueeze(1).expand_as(xx)
    ry = yy[:, diag_ind, diag_ind].unsqueeze(1).expand_as(yy)
    P = (rx.transpose(2, 1) + ry - 2 * zz)
    return P.min(1)[0], P.min(2)[0]


def EMD_CD(sample_pcs, ref_pcs, batch_size, accelerated_cd=False, reduced=True):
    N_sample = sample_pcs.shape[0]
    N_ref = ref_pcs.shape[0]
    assert N_sample == N_ref, "REF:%d SMP:%d" % (N_ref, N_sample)

    cd_lst = []
    emd_lst = []
    iterator = range(0, N_sample, batch_size)

    for b_start in tqdm(iterator, desc='EMD-CD'):
        b_end = min(N_sample, b_start + batch_size)
        sample_batch = sample_pcs[b_start:b_end]
        ref_batch = ref_pcs[b_start:b_end]

        if accelerated_cd:
            dl, dr = distChamferCUDA(sample_batch, ref_batch)
        else:
            dl, dr = distChamfer(sample_batch, ref_batch)
        cd_lst.append(dl.mean(dim=1) + dr.mean(dim=1))

        emd_batch = emd_approx(sample_batch, ref_batch)
        emd_lst.append(emd_batch)

    if reduced:
        cd = torch.cat(cd_lst).mean()
        emd = torch.cat(emd_lst).mean()
    else:
        cd = torch.cat(cd_lst)
        emd = torch.cat(emd_lst)

    results = {
        'MMD-CD': cd,
        'MMD-EMD': emd,
    }
    return results


def _pairwise_EMD_CD_(sample_pcs, ref_pcs, batch_size,
                      accelerated_cd=True, verbose=True, mask_sample=None, mask_ref=None):
    N_sample = sample_pcs.shape[0]
    N_ref = ref_pcs.shape[0]
    all_cd = []
    all_emd = []
    iterator = range(N_sample)
    if verbose:
        iterator = tqdm(iterator, desc='Pairwise EMD-CD')
    for sample_b_start in iterator:
        sample_batch = sample_pcs[sample_b_start]
        if mask_sample is not None:
            _mask_sample = mask_sample[sample_b_start]
        cd_lst = []
        emd_lst = []
        sub_iterator = range(0, N_ref, batch_size)
        if verbose:
            sub_iterator = tqdm(sub_iterator, leave=False)
        for ref_b_start in sub_iterator:
            ref_b_end = min(N_ref, ref_b_start + batch_size)
            ref_batch = ref_pcs[ref_b_start:ref_b_end]
            if mask_ref is not None:
                _mask_ref = mask_ref[ref_b_start:ref_b_end]

            batch_size_ref = ref_batch.size(0)
            point_dim = ref_batch.size(2)
            sample_batch_exp = sample_batch.view(1, -1, point_dim).expand(
                batch_size_ref, -1, -1)
            sample_batch_exp = sample_batch_exp.contiguous()

            if accelerated_cd:
                dl, dr = distChamferCUDA(sample_batch_exp, ref_batch)
                    
            else:
                dl, dr = distChamfer(sample_batch_exp, ref_batch)
            if mask_sample is not None:
                dl_mean = (dl * _mask_sample.unsqueeze(0)).sum(1) / _mask_sample.sum()
            else:
                dl_mean = dl.mean(1)
            if mask_ref is not None:
                dr_mean = (dr * _mask_ref).sum(1) / _mask_ref.sum(1)
            else:
                dr_mean = dr.mean(1)
            cd_lst.append((dl_mean + dr_mean).view(1, -1))

            emd_batch = emd_approx(sample_batch_exp, ref_batch)
            emd_lst.append(emd_batch.view(1, -1))

        cd_lst = torch.cat(cd_lst, dim=1)
        emd_lst = torch.cat(emd_lst, dim=1)
        all_cd.append(cd_lst)
        all_emd.append(emd_lst)
        # print(cd_lst.min(1)[0])

    all_cd = torch.cat(all_cd, dim=0)  # N_sample, N_ref
    all_emd = torch.cat(all_emd, dim=0)  # N_sample, N_ref

    return all_cd, all_emd


# Adapted from https://github.com/xuqiantong/
# GAN-Metrics/blob/master/framework/metric.py
def knn(Mxx, Mxy, Myy, k, sqrt=False, one_way=False):
    n0 = Mxx.size(0)
    n1 = Myy.size(0)
    print(n0, n1)
    label = torch.cat((torch.ones(n0), torch.zeros(n1))).to(Mxx)
    M = torch.cat([
        torch.cat((Mxx, Mxy), 1),
        torch.cat((Mxy.transpose(0, 1), Myy), 1)], 0)
    if sqrt:
        M = M.abs().sqrt()
    INFINITY = float('inf')
    val, idx = (M + torch.diag(INFINITY * torch.ones(n0 + n1).to(Mxx))).topk(
        k, 0, False)
    print(val.shape, idx.shape)
    print(val[0, :n0].mean(), val[0, n0:].mean())
    count = torch.zeros(n0 + n1).to(Mxx)
    for i in range(0, k):
        count = count + label.index_select(0, idx[i])
    pred = torch.ge(count, (float(k) / 2) * torch.ones(n0 + n1).to(Mxx)).float()
    if one_way:
        pred = pred[:n0]
        label = pred[:n0]

    s = {
        'tp': (pred * label).sum(),
        'fp': (pred * (1 - label)).sum(),
        'fn': ((1 - pred) * label).sum(),
        'tn': ((1 - pred) * (1 - label)).sum(),
    }

    s.update({
        'precision': s['tp'] / (s['tp'] + s['fp'] + 1e-10),
        'recall': s['tp'] / (s['tp'] + s['fn'] + 1e-10),
        'acc_t': s['tp'] / (s['tp'] + s['fn'] + 1e-10),
        'acc_f': s['tn'] / (s['tn'] + s['fp'] + 1e-10),
        'acc': torch.eq(label, pred).float().mean(),
    })
    return s


def lgan_mmd_cov(all_dist, thresh=1000):
    N_sample, N_ref = all_dist.size(0), all_dist.size(1)
    min_val_fromsmp, min_idx = torch.min(all_dist, dim=1)
    min_val, idx = torch.min(all_dist, dim=0)
    min_val, idxx = torch.sort(min_val)
    print("Min dist of test shape to generated sample")
    print(min_val, idx[idxx], idxx)
    min_val_fromsmp, idxxx = torch.sort(min_val_fromsmp)
    print("Min dist of generated sample to test shape")
    torch.set_printoptions(profile="full")
    print(min_val_fromsmp, min_idx[idxxx], idxxx)
    torch.set_printoptions(profile="default")
    sorted_idx = idx[idxx]
    outlier_mask = min_val > thresh 
    if torch.any(outlier_mask):
        sorted_idx[outlier_mask] = sorted_idx[0]
    mmd = min_val.mean()
    mmd_smp = min_val_fromsmp.mean()
    cov = float(sorted_idx.unique().view(-1).size(0)) / float(N_ref)
    print(sorted_idx.unique().view(-1))
    cov = torch.tensor(cov).to(all_dist)
    return {
        'lgan_mmd': mmd,
        'lgan_cov': cov,
        'lgan_mmd_smp': mmd_smp,
    }


def lgan_mmd_cov_match(all_dist):
    N_sample, N_ref = all_dist.size(0), all_dist.size(1)
    min_val_fromsmp, min_idx = torch.min(all_dist, dim=1)
    min_val, _ = torch.min(all_dist, dim=0)
    mmd = min_val.mean()
    mmd_smp = min_val_fromsmp.mean()
    cov = float(min_idx.unique().view(-1).size(0)) / float(N_ref)
    cov = torch.tensor(cov).to(all_dist)
    return {
        'lgan_mmd': mmd,
        'lgan_cov': cov,
        'lgan_mmd_smp': mmd_smp,
    }, min_idx.view(-1)

def compute_bbox_metric(preds, preds_mask, refs, refs_mask, batch_size, n_class=4, thresh=1.0, metric='chamfer', no_nn=False, cov_thresh=100):
    print("Bbox EMD CD")
    pred_size = preds.shape[0]
    ref_size = refs.shape[0]
    assert pred_size == preds_mask.shape[0]
    assert ref_size == refs_mask.shape[0]
    pred_params = []
    ref_params = []
    for i in range(pred_size):
        _pred = preds[i]
        _pred_seg_mask = preds_mask[i]
            
        _pred_max = _pred.max(0)[0].reshape(1,3) # (1, 3)
        _pred_min = _pred.min(0)[0].reshape(1,3) # (1, 3)
        _pred_shift = ((_pred_min + _pred_max) / 2).reshape(1, 3)
        _pred_scale = (_pred_max - _pred_min).max(-1)[0].reshape(1, 1) / 2
        _pred = (_pred - _pred_shift) / _pred_scale
        one_pred_params = dict()
        for j in range(n_class):
            idx = _pred_seg_mask == j 
            if torch.any(idx):
                points = _pred[idx]
                if points.shape[0] > 100:
                    part_max = torch.quantile(points, thresh, dim=0, keepdim=True).cpu()
                    part_min =  torch.quantile(points, 1 - thresh, dim=0, keepdim=True).cpu()
                    one_pred_params[j] = (part_min, part_max)
        pred_params.append(one_pred_params)
        
    for i in range(ref_size):
        _ref = refs[i]
        _ref_max = _ref.max(0)[0].reshape(1,3) # (1, 3)
        _ref_min = _ref.min(0)[0].reshape(1,3) # (1, 3)
        _ref_shift = ((_ref_min + _ref_max) / 2).reshape(1, 3)
        _ref_scale = (_ref_max - _ref_min).max(-1)[0].reshape(1, 1) / 2
        _ref = (_ref - _ref_shift) / _ref_scale
        _ref_seg_mask = refs_mask[i]
        one_ref_params = dict()
        for j in range(n_class):
            ref_idx = _ref_seg_mask == j
            if torch.any(ref_idx):
                ref_points = _ref[ref_idx]
                if ref_points.shape[0] > 100:
                    ref_part_max = torch.quantile(ref_points, thresh, dim=0, keepdim=True).cpu()
                    ref_part_min =  torch.quantile(ref_points, 1 - thresh, dim=0, keepdim=True).cpu()
                    one_ref_params[j] = (ref_part_min, ref_part_max)

        ref_params.append(one_ref_params)
    if metric == 'chamfer':
        metric = compute_all_metrics_cust_func(pred_params, ref_params, partial(part_chamfer, n_class), 'bbox_chamfer', accelerated_cd=True, no_nn=no_nn, thresh=cov_thresh)
    elif metric == 'iou':
        metric = compute_all_metrics_cust_func(pred_params, ref_params, partial(part_miou, n_class), 'bbox_iou', accelerated_cd=True, no_nn=no_nn, thresh=cov_thresh)
    elif metric == 'l2':
        metric = compute_all_metrics_cust_func(pred_params, ref_params, partial(part_l2, n_class), 'bbox_l2', accelerated_cd=True, no_nn=no_nn, thresh=cov_thresh)
    else:
        raise NotImplementedError
    metric = {f"bbox_{k}":v for k, v in metric.items()}
    
    return metric

def compute_all_metrics_cust_func(sample_pcs, ref_pcs, dist_func, dist_name, accelerated_cd=False, no_nn=False, thresh=1000):
    results = {}
    
    print("Pairwise EMD CD")
    M = len(sample_pcs)
    N = len(ref_pcs)
    rs_dists = torch.zeros([N, M])
    for i in tqdm(range(N)):
        for j in range(M):
            rs_dists[i, j] = dist_func(ref_pcs[i], sample_pcs[j], accelerated=accelerated_cd)
            
    res_cd = lgan_mmd_cov(rs_dists.t(), thresh=cov_thresh)
    results.update({
        f"{k}-{dist_name}": v for k, v in res_cd.items()
    })
    
    for k, v in results.items():
        print('[%s] %.8f' % (k, v.item()))
    if no_nn:
        return results
    rr_dists = torch.zeros([N, N])
    for i in tqdm(range(N)):
        for j in range(N):
            rr_dists[i, j] = dist_func(ref_pcs[i], ref_pcs[j], accelerated=accelerated_cd)
            
    ss_dists = torch.zeros([M, M])
    for i in tqdm(range(M)):
        for j in range(M):
            ss_dists[i, j] = dist_func(sample_pcs[i], sample_pcs[j], accelerated=accelerated_cd)
            
       
    # 1-NN results
    one_nn_cd_res = knn(rr_dists, rs_dists, ss_dists, 1, sqrt=False)
    results.update({
        f"1-NN-{dist_name}-{k}": v for k, v in one_nn_cd_res.items() if 'acc' in k
    })

    return results

def compute_snapping_metric(preds, preds_mask, cls='Chair'):
    print("Snapping Metric")
    if cls == 'Chair':
        connected = [(0, [1,2]), (1, [2]), (3, [0, 1])]
    elif cls == 'Airplane':
        connected = [(1, [0]), (2, [0]), (3, [0, 1])]
    dists = {p[0]:[] for p in connected}
    for k in tqdm(range(preds.shape[0])):
        pred = preds[k]
        pred_mask = preds_mask[k]
        min_d = []
        for i, js in connected:
            id_part_A = pred_mask == i 
            min_d = []
            for j in js:
                id_part_B = pred_mask == j 
                if torch.any(id_part_A) and torch.any(id_part_B):
                    point_A = pred[id_part_A]
                    point_B = pred[id_part_B]
                    dist = ((point_A[:, None] - point_B[None]) ** 2).sum(-1)
                    min_point_A_val, id_min_point_A = dist.min(1)[0].topk(50,dim=0, largest=False)
                    min_point_B_val, id_min_point_B = dist.min(0)[0].topk(50,dim=0, largest=False)
                    min_point_A = point_A[id_min_point_A]
                    min_point_B = point_B[id_min_point_B]
                    dl, dr = distChamfer(min_point_A[None], min_point_B[None])
                    min_d.append(dl.mean(1) + dr.mean(1))
            if len(min_d):
                min_d = torch.min(torch.tensor(min_d)).item()
                dists[i].append((k, min_d)) 
                # print(dist)
    d_array = torch.tensor([d[1] for d in dists[1]])
    d_ind_array = torch.tensor([d[0] for d in dists[1]])
    print(d_array.sort()[0], d_ind_array[d_array.sort()[1]])
    dists = {f"snapping_{cls}_{k}":torch.tensor(list(zip(*v))[1]).mean() for k, v in dists.items()}
    return dists
                    
            

def compute_part_metric(preds, preds_mask, refs, refs_mask, batch_size, n_class=4):
    print('part level Pairwise EMD CD')
    pred_size = preds.shape[0]
    ref_size = refs.shape[0]
    pred_part_pointcloud = [[] for _ in range(n_class)]
    pred_part_mask = [[] for _ in range(n_class)]
    ref_part_pointcloud = [[] for _ in range(n_class)]
    assert pred_size == preds_mask.shape[0]
    assert ref_size == refs_mask.shape[0]
    for i in tqdm(range(pred_size)):
        _pred = preds[i]
        _pred_seg_mask = preds_mask[i]
        for j in range(n_class):
            idx = _pred_seg_mask == j 
            if torch.any(idx):
                points = _pred[idx].unsqueeze(0)
                if points.shape[1] > 100:
                    mask = torch.ones([1, 512])
                    if points.shape[1] < 512:
                        mask[:, points.shape[1]:] = 0
                    while points.shape[1] < 512:
                        points = torch.cat([points, points], dim=1)
                    if points.shape[1] > 512:
                        points = points[:, :512]
                    pred_part_pointcloud[j].append(points)
                    assert (mask.shape == (1, 512))
                    pred_part_mask[j].append(mask)
    for i in tqdm(range(ref_size)):
        _ref = refs[i]
        _ref_seg_mask = refs_mask[i]
        for j in range(n_class):
            idx = _ref_seg_mask == j 
            if torch.any(idx):
                points = _ref[idx].unsqueeze(0)
                if points.shape[1] > 100:
                    while points.shape[1] < 512:
                        points = torch.cat([points, points], dim=1)
                    if points.shape[1] > 512:
                        points = points[:, :512]
                    ref_part_pointcloud[j].append(points)
    pred_part_pointcloud, ref_part_pointcloud, pred_part_mask = map(lambda l: [torch.cat(ll, dim=0) for ll in l], [pred_part_pointcloud, ref_part_pointcloud, pred_part_mask])
    weight = []
    for l, s, k in zip(pred_part_pointcloud, ref_part_pointcloud, pred_part_mask):
        print(s.shape, l.shape, k.shape)
        weight.append(s.shape[0])
    N = sum(weight)
    weight = [w / N for w in weight]
    print('Normalizing------------------')
    for i in range(4):
        _pred = pred_part_pointcloud[i]
        _pred_max = _pred.max(1)[0].reshape(-1, 1,3) # (1, 3)
        _pred_min = _pred.min(1)[0].reshape(-1, 1,3) # (1, 3)
        _pred_shift = ((_pred_min + _pred_max) / 2).reshape(-1, 1, 3)
        _pred_scale = (_pred_max - _pred_min).reshape(-1, 1, 3) / 2
        _pred = (_pred - _pred_shift) / _pred_scale
        pred_part_pointcloud[i] = _pred
        _ref = ref_part_pointcloud[i]
        _ref_max = _ref.max(1)[0].reshape(-1, 1,3) # (1, 3)
        _ref_min = _ref.min(1)[0].reshape(-1, 1,3) # (1, 3)
        _ref_shift = ((_ref_min + _ref_max) / 2).reshape(-1, 1, 3)
        _ref_scale = (_ref_max - _ref_min).reshape(-1, 1, 3) / 2
        _ref = (_ref - _ref_shift) / _ref_scale
        ref_part_pointcloud[i] = _ref
    metrics = []
    for i in range(n_class):
        print(f"Part {i} --------------------------------")
        _pred = pred_part_pointcloud[i]
        _ref = ref_part_pointcloud[i]
        _pred_mask = pred_part_mask[i].cuda()
        metric = compute_all_metrics(_pred, _ref, 32, mask=_pred_mask)
        metrics.append(metric)
    avg_m = {f"part_weighted_{k}":0 for k in metrics[0].keys()}
    for i, m in enumerate(metrics):
        for k, v in m.items():
            avg_m[f"part_weighted_{k}"] += v * weight[i]
    return avg_m

def compute_all_metrics(sample_pcs, ref_pcs, batch_size, accelerated_cd=True, one_way=False, mask=None):
    results = {}
    
    print("Pairwise EMD CD")
    M_rs_cd, M_rs_emd = _pairwise_EMD_CD_(
        ref_pcs, sample_pcs, batch_size, accelerated_cd=accelerated_cd, mask_ref=mask)
    print(M_rs_cd.min(1)[0].mean(), M_rs_cd.min(0)[0].mean())
    res_cd = lgan_mmd_cov(M_rs_cd.t())
    results.update({
        "%s-CD" % k: v for k, v in res_cd.items()
    })

    res_emd = lgan_mmd_cov(M_rs_emd.t())
    results.update({
        "%s-EMD" % k: v for k, v in res_emd.items()
    })

    for k, v in results.items():
        print('[%s] %.8f' % (k, v.item()))

    M_rr_cd, M_rr_emd = _pairwise_EMD_CD_(
        ref_pcs, ref_pcs, batch_size, accelerated_cd=accelerated_cd)
    
    if not one_way:
        M_ss_cd, M_ss_emd = _pairwise_EMD_CD_(
            sample_pcs, sample_pcs, batch_size, accelerated_cd=accelerated_cd, mask_ref=mask, mask_sample=mask)
    else:
        INFINITY = float('inf')
        M_ss_cd =  torch.zeros(M_rs_cd.shape[1], M_rs_cd.shape[1]).cuda() + INFINITY
        M_ss_emd = torch.zeros(M_rs_cd.shape[1], M_rs_cd.shape[1]) + INFINITY
        
    # 1-NN results
    one_nn_cd_res = knn(M_rr_cd, M_rs_cd, M_ss_cd, 1, sqrt=False, one_way=one_way)
    results.update({
        "1-NN-CD-%s" % k: v for k, v in one_nn_cd_res.items() if 'acc' in k
    })
    one_nn_emd_res = knn(M_rr_emd, M_rs_emd, M_ss_emd, 1, sqrt=False, one_way=one_way)
    results.update({
        "1-NN-EMD-%s" % k: v for k, v in one_nn_emd_res.items() if 'acc' in k
    })

    return results


#######################################################
# JSD : from https://github.com/optas/latent_3d_points
#######################################################
def unit_cube_grid_point_cloud(resolution, clip_sphere=False):
    """Returns the center coordinates of each cell of a 3D grid with
    resolution^3 cells, that is placed in the unit-cube. If clip_sphere it True
    it drops the "corner" cells that lie outside the unit-sphere.
    """
    grid = np.ndarray((resolution, resolution, resolution, 3), np.float32)
    spacing = 1.0 / float(resolution - 1)
    for i in range(resolution):
        for j in range(resolution):
            for k in range(resolution):
                grid[i, j, k, 0] = i * spacing - 0.5
                grid[i, j, k, 1] = j * spacing - 0.5
                grid[i, j, k, 2] = k * spacing - 0.5

    if clip_sphere:
        grid = grid.reshape(-1, 3)
        grid = grid[norm(grid, axis=1) <= 0.5]

    return grid, spacing


def jsd_between_point_cloud_sets(
        sample_pcs, ref_pcs, resolution=28):
    """Computes the JSD between two sets of point-clouds,
       as introduced in the paper
    ```Learning Representations And Generative Models For 3D Point Clouds```.
    Args:
        sample_pcs: (np.ndarray S1xR2x3) S1 point-clouds, each of R1 points.
        ref_pcs: (np.ndarray S2xR2x3) S2 point-clouds, each of R2 points.
        resolution: (int) grid-resolution. Affects granularity of measurements.
    """
    in_unit_sphere = True
    sample_grid_var = entropy_of_occupancy_grid(
        sample_pcs, resolution, in_unit_sphere)[1]
    ref_grid_var = entropy_of_occupancy_grid(
        ref_pcs, resolution, in_unit_sphere)[1]
    return jensen_shannon_divergence(sample_grid_var, ref_grid_var)


def entropy_of_occupancy_grid(
        pclouds, grid_resolution, in_sphere=False, verbose=False):
    """Given a collection of point-clouds, estimate the entropy of
    the random variables corresponding to occupancy-grid activation patterns.
    Inputs:
        pclouds: (numpy array) #point-clouds x points per point-cloud x 3
        grid_resolution (int) size of occupancy grid that will be used.
    """
    epsilon = 10e-4
    bound = 0.5 + epsilon
    if abs(np.max(pclouds)) > bound or abs(np.min(pclouds)) > bound:
        if verbose:
            warnings.warn('Point-clouds are not in unit cube.')

    if in_sphere and np.max(np.sqrt(np.sum(pclouds ** 2, axis=2))) > bound:
        if verbose:
            warnings.warn('Point-clouds are not in unit sphere.')

    grid_coordinates, _ = unit_cube_grid_point_cloud(grid_resolution, in_sphere)
    grid_coordinates = grid_coordinates.reshape(-1, 3)
    grid_counters = np.zeros(len(grid_coordinates))
    grid_bernoulli_rvars = np.zeros(len(grid_coordinates))
    nn = NearestNeighbors(n_neighbors=1).fit(grid_coordinates)

    for pc in tqdm(pclouds, desc='JSD'):
        _, indices = nn.kneighbors(pc)
        indices = np.squeeze(indices)
        for i in indices:
            grid_counters[i] += 1
        indices = np.unique(indices)
        for i in indices:
            grid_bernoulli_rvars[i] += 1

    acc_entropy = 0.0
    n = float(len(pclouds))
    for g in grid_bernoulli_rvars:
        if g > 0:
            p = float(g) / n
            acc_entropy += entropy([p, 1.0 - p])

    return acc_entropy / len(grid_counters), grid_counters


def jensen_shannon_divergence(P, Q):
    if np.any(P < 0) or np.any(Q < 0):
        raise ValueError('Negative values.')
    if len(P) != len(Q):
        raise ValueError('Non equal size.')

    P_ = P / np.sum(P)  # Ensure probabilities.
    Q_ = Q / np.sum(Q)

    e1 = entropy(P_, base=2)
    e2 = entropy(Q_, base=2)
    e_sum = entropy((P_ + Q_) / 2.0, base=2)
    res = e_sum - ((e1 + e2) / 2.0)

    res2 = _jsdiv(P_, Q_)

    if not np.allclose(res, res2, atol=10e-5, rtol=0):
        warnings.warn('Numerical values of two JSD methods don\'t agree.')

    return res


def _jsdiv(P, Q):
    """another way of computing JSD"""

    def _kldiv(A, B):
        a = A.copy()
        b = B.copy()
        idx = np.logical_and(a > 0, b > 0)
        a = a[idx]
        b = b[idx]
        return np.sum([v for v in a * np.log2(a / b)])

    P_ = P / np.sum(P)
    Q_ = Q / np.sum(Q)

    M = 0.5 * (P_ + Q_)

    return 0.5 * (_kldiv(P_, M) + _kldiv(Q_, M))


if __name__ == '__main__':
    a = torch.randn([16, 2048, 3]).cuda()
    b = torch.randn([16, 2048, 3]).cuda()
    print(EMD_CD(a, b, batch_size=8))
python/difffacto/datasets/evaluation_utils.py
diffFacto/diffFacto
import os
import torch
import torch.nn.functional as F
from difffacto.utils.registry import DATASETS
from torch.utils.data import Dataset
from difffacto.utils.misc import fps
import numpy as np
from tqdm.auto import tqdm
from .dataset_utils import pc_norm, shapenet_part_normal_cat_to_id, build_loss_funcs, DataLoaderWrapper, DataLoaderWrapperOne, load_ply_withNormals
from .evaluation_utils import EMD_CD, compute_all_metrics
import json


@DATASETS.register_module()
def Partnet(
    batch_size, 
    root, 
    n_part, 
    npoints, 
    part_npoints,
    split, 
    num_workers=0, 
    eval_mode='ae',
    scale_mode='shape_unit', 
    part_scale_mode=None,
    distributed=False, 
    shuffle=True, 
    class_choice='Chair',
    save_only=False,
    augment=False):
    assert not distributed
    dataset = _Partnet(root, npoints, part_npoints, split, scale_mode=scale_mode, class_choice=class_choice, save_only=save_only, eval_mode=eval_mode, augment=augment, n_part=n_part,part_scale_mode=part_scale_mode)

    return DataLoaderWrapperOne(dataset, batch_size=batch_size, num_workers=num_workers, drop_last=True, shuffle=shuffle), None




class _Partnet(Dataset):

    def __init__(
        self,
        root = '/orion/u/w4756677/datasets/coalace_data_processed/',
        npoints=2048,
        part_npoints=512,
        split='train', 
        n_part=4, 
        class_choice='Chair',
        normal_channel=False, 
        scale_mode='shape_bbox',
        part_scale_mode=None,
        eval_mode='ae',
        save_only=False,
        augment=False,
    ):
        self.eval_mode=eval_mode
        self.augment=augment
        assert not (self.augment and split == 'test')
        self.npoints = npoints # Sampling points per part
        self.root = root # File root path
        self.catfile = os.path.join(self.root, 'synsetoffset2category.txt') # Path corresponding to category and folder name
        self.cat = {}
        self.normal_channel = normal_channel # Whether to use rgb information
        self.scale_mode = scale_mode
        self.part_scale_mode=part_scale_mode
        self.part_npoints=part_npoints

        self.cat = shapenet_part_normal_cat_to_id #{'Airplane': '02691156', 'Bag': '02773838', 'Cap': '02954340', 'Car': '02958343', 'Chair': '03001627', 'Earphone': '03261776', 'Guitar': '03467517', 'Knife': '03624134', 'Lamp': '03636649', 'Laptop': '03642806', 'Motorbike': '03790512', 'Mug': '03797390', 'Pistol': '03948459', 'Rocket': '04099429', 'Skateboard': '04225987', 'Table': '04379243'}
        self.classes_original = dict(zip(self.cat, range(len(self.cat)))) #{'Airplane': 0, 'Bag': 1, 'Cap': 2, 'Car': 3, 'Chair': 4, 'Earphone': 5, 'Guitar': 6, 'Knife': 7, 'Lamp': 8, 'Laptop': 9, 'Motorbike': 10, 'Mug': 11, 'Pistol': 12, 'Rocket': 13, 'Skateboard': 14, 'Table': 15}

        self.cat = {class_choice:shapenet_part_normal_cat_to_id[class_choice]} if class_choice is not None else shapenet_part_normal_cat_to_id
        self.class_choice = class_choice
        self.save_only=save_only
        print(self.cat)
        self.noises = dict()


        self.meta = {} # Read the classified folder jason file and put their names in the list
        with open(os.path.join(self.root, 'train_test_split', 'shuffled_train_file_list.json'), 'r') as f:
            train_ids = set([str(d.split('/')[2]) for d in json.load(f)]) # '928c86eabc0be624c2bf2dc31ba1713' this is the first value
        with open(os.path.join(self.root, 'train_test_split', 'shuffled_val_file_list.json'), 'r') as f:
            val_ids = set([str(d.split('/')[2]) for d in json.load(f)])
        with open(os.path.join(self.root, 'train_test_split', 'shuffled_test_file_list.json'), 'r') as f:
            test_ids = set([str(d.split('/')[2]) for d in json.load(f)])
        for item in self.cat:
            self.meta[item] = []
            dir_point = os.path.join(self.root, self.cat[item]) # # Get the path corresponding to a folder, such as the first folder 02691156
            fns = sorted(os.listdir(dir_point))  # Get each txt file under the folder according to the path and put it in the list
            # print(fns[0][0:-4])
            densePoints = [f for f in fns if (f[-7:] == 'ori.ply')]
            labels = [f for f in fns if (f[-10:] == 'ori.labels')]
            assert len(densePoints) == len(labels)
            if split == 'trainval':
                labels = [fn for fn in labels if ((fn[0:-11] in train_ids) or (fn[0:-11] in val_ids))]
                densePoints = [fn for fn in densePoints if ((fn[0:-8] in train_ids) or (fn[0:-8] in val_ids))]
            elif split == 'train':
                labels = [fn for fn in labels if ((fn[0:-11] in train_ids))]
                densePoints = [fn for fn in densePoints if ((fn[0:-8] in train_ids))] # Judge whether the txt file in the folder is in the training TXT. If so, the txt file obtained in fns is the file to be trained in all txt files in this category, and put it into fns
            elif split == 'val':
                labels = [fn for fn in labels if ((fn[0:-11] in val_ids))]
                densePoints = [fn for fn in densePoints if ((fn[0:-8] in val_ids))]
            elif split == 'test':
                labels = [fn for fn in labels if ((fn[0:-11] in test_ids))]
                densePoints = [fn for fn in densePoints if ((fn[0:-8] in test_ids) )]
            else:
                print('Unknown split: %s. Exiting..' % (split))
                exit(-1)

            # print(os.path.basename(fns))
            for fn in densePoints:
                token = (os.path.splitext(os.path.basename(fn))[0])
                self.meta[item].append((os.path.join(dir_point, token + '.ply'), os.path.join(dir_point, token + '.labels')))  # Generate a dictionary and combine the category name and the training path as the training data in a large class
                #After the above code is executed, you can put all the data to be trained or verified into a dictionary. The key of the dictionary is the category to which the data belongs, such as aircraft. Value is the full path of its corresponding data
                #{aircraft: [path 1, path 2.....]}
        #####################################################################################################################################################
        self.datapath = []
        for item in self.cat: # self.cat is the dictionary corresponding to the category name and folder
            for fn in self.meta[item]:
                self.datapath.append((item, *fn)) # Generate tuples of labels and point cloud paths, and convert the dictionary in self.met into a tuple
        print(len(self.datapath))
        self.classes = {}
        for i in self.cat.keys():
            self.classes[i] = self.classes_original[i]
        ## self.classes corresponds the name and index of the category, such as aircraft < --- > 0
        # Mapping from category ('Chair') to a list of int [10,11,12,13] as segmentation labels

        # for cat in sorted(self.seg_classes.keys()):
        #     print(cat, self.seg_classes[cat])

        self.cache = {}  # from index to (point_set, cls, seg) tuple
        self.cache_size = 20000
        self.num_class = n_part

    def store_noise(self, all_noises, all_ids):
        self.noises = dict()
        for noises, ids in zip(all_noises, all_ids):
            assert noises.shape[0] == len(ids)
            for i, id in enumerate(ids):
                self.noises[id.item()] = noises[i].detach().cpu().numpy()

    def __getitem__(self, index):
        if index in self.cache: # The initial slef.cache is an empty dictionary, which is used to store the retrieved data and place it according to (point_set, cls, seg) while avoiding repeated sampling
            valid_point_set, cls, valid_seg, valid_seg_onehot, shift, scale = self.cache[index]
        else:
            fn = self.datapath[index] # The path to get the training data according to the index self.datepath is a tuple (class name, path)
            cat = self.datapath[index][0] # Get the class name
            cls = self.classes[cat] # Convert class name to index
            cls = np.array([cls]).astype(np.int32)
            point_set = load_ply_withNormals(fn[1]).astype(np.float32) # size 20488,7 read in this txt file, a total of 20488 points, each point xyz rgb + sub category label
            
            seg = np.loadtxt(fn[2]).astype(np.long) - 1 # Get the label of the sub category
            seg_onehot = F.one_hot(torch.from_numpy(seg), num_classes=self.num_class)
            
            # get rid of dummy points
            mask = np.abs(point_set).sum(1) != 0
            valid_point_set = point_set[mask]
            valid_seg = seg[mask]
            valid_seg_onehot = seg_onehot[mask]
            valid_point_set, shift, scale = pc_norm(valid_point_set, self.scale_mode, to_torch=False) 
            if len(self.cache) < self.cache_size:
                self.cache[index] = (valid_point_set, cls, valid_seg, valid_seg_onehot, shift, scale)
        
        noise = self.noises.get(index, np.zeros(1))#for cimle
        choice = np.random.choice(len(valid_seg), self.npoints, replace=False) # Random sampling of data in a category returns the index, allowing repeated sampling
        # resample
        ref_point_set = valid_point_set[choice, :] # Sample by index
        ref_seg = valid_seg[choice]
        ref_seg_onehot = valid_seg_onehot[choice]
        shifts = np.zeros((self.num_class, 3))
        scales = np.zeros((self.num_class, 3))
        present = torch.zeros(self.num_class)
        out = np.zeros((self.num_class, self.part_npoints, 3))
        out_mask = np.zeros((self.num_class, self.part_npoints))
        out_onehot = np.zeros((self.num_class, self.part_npoints, self.num_class))
        # print(seg.shape, seg_onehot.shape)
        for i in range(self.num_class):
            idx = valid_seg == i
            if idx.any(): # if less than 10 points we ignore
                present[i] = 1
                part_point_set = valid_point_set[idx]
                shifted_chosen, pshift, pscale = pc_norm(part_point_set, self.part_scale_mode, to_torch=False, clip=True) 
                if self.part_npoints < shifted_chosen.shape[0]:
                    choice = np.random.choice(shifted_chosen.shape[0], self.part_npoints, replace=False) # Random sampling of data in a category returns the index, allowing repeated sampling
                    shifted_chosen = shifted_chosen[choice]
                    mask_chosen = valid_seg[choice]
                    onehot_chosen = valid_seg_onehot[choice]
                    
                shifts[i] = part_point_set.mean(0)
                scales[i] = part_point_set.std(0).clip(1e-2, 1)
                out[i] = shifted_chosen
                out_mask[i] = mask_chosen
                out_onehot[i] = onehot_chosen
        ref_point_set, shift, scale = map(lambda t: torch.from_numpy(t).float(), [ref_point_set, shift, scale])
        out, shifts, scales = map(lambda t: torch.from_numpy(t).float(), [out, shifts, scales])
        if self.augment:
            rand_scale = torch.rand(1,3) / 2 + 0.7
            rand_shift = torch.rand(1, 3) - 0.5
            # (o + s) * scale
            ori_out = (ori_out + rand_shift) * rand_scale
            shift = shift  + scale * rand_shift
            scale = rand_scale * scale


        return {'part_scale': scales.transpose(0, 1), 
                'part_shift': shifts.transpose(0, 1), 
                'input':out, 
                "present": present,
                'attn_map': out_onehot, 
                'ref':ref_point_set,
                'ref_attn_map': ref_seg_onehot,
                'ref_seg_mask': ref_seg,
                'class':cls, 
                'seg_mask':out_mask, 
                'shift':shift, 
                'scale':scale, 
                'id':index,
                'noise':torch.from_numpy(noise)
                } # pointset is point cloud data, cls has 16 categories, and seg is a small category corresponding to different points in the data

    def __len__(self):
        return len(self.datapath)

    def evaluate(self, results, save_num_batch, device):
        save_dict = dict()
        preds = []
        refs = []
        for idx, pred_dict in enumerate(tqdm(results)):
            shift = pred_dict['shift']
            scale = pred_dict['scale']
            del pred_dict['shift']
            del pred_dict['scale']
            pred = pred_dict['pred']
            ref = pred_dict['input_ref']
            if pred.shape[1] > 2048:
                pred = fps(pred.to('cuda'), 2048)
            if ref.shape[1] > 2048:
                ref = fps(ref.to('cuda'), 2048)
            if self.eval_mode == 'ae':
                pred = pred * scale + shift
                ref = ref * scale + shift
            else:
                pred_max = pred.max(1)[0].reshape(-1, 1, 3) # (1, 3)
                pred_min = pred.min(1)[0].reshape(-1, 1, 3) # (1, 3)
                pred_shift = ((pred_min + pred_max) / 2).reshape(-1, 1, 3)
                pred_scale = (pred_max - pred_min).max(-1)[0].reshape(-1, 1, 1) / 2
                pred = (pred - pred_shift) / pred_scale
                ref_max = ref.max(1)[0].reshape(-1, 1, 3) # (1, 3)
                ref_min = ref.min(1)[0].reshape(-1, 1, 3) # (1, 3)
                ref_shift = ((ref_min + ref_max) / 2).reshape(-1, 1, 3)
                ref_scale = (ref_max - ref_min).max(-1)[0].reshape(-1, 1, 1) / 2
                ref = (ref - ref_shift) / ref_scale
                if self.using_whole_chair_only:
                    present = pred_dict['present'][:, :3].sum(1) == 3 
                    pred = pred[present]
                    ref = ref[present]
            preds.append(pred)
            refs.append(ref)
            B, N, C = ref.shape
            pred_dict = {k: v.reshape(B, N, C) * scale + shift if v.shape[-1] == 3 else v for k, v in pred_dict.items()}    
            pred_dict['pred'] = pred 
            pred_dict['input_ref'] = ref  
                     
            if idx < save_num_batch:
                for k, v in pred_dict.items():
                    if k not in save_dict.keys():
                        save_dict[k] = []
                    save_dict[k].append(v.detach().cpu().numpy())
        preds = torch.cat(preds, dim=0).cuda()
        refs = torch.cat(refs, dim=0).cuda()
        if self.eval_mode == 'ae':
            metrics = EMD_CD(preds, refs, 32)
        elif self.eval_mode == 'gen':
            # metrics = dict(l=torch.zeros(1))
            print(preds.shape, refs.shape)
            metrics = compute_all_metrics(preds, refs, 32)
        ssave_dict = dict()
        for k, v in save_dict.items():
            ssave_dict[k] = np.concatenate(v, axis=0)
        return ssave_dict, metrics

python/difffacto/datasets/partnet_dataset.py
diffFacto/diffFacto
import numpy as np
import torch
from torch.utils.data import Dataset
from difffacto.utils.registry import DATASETS 
from difffacto.utils.misc import *
import os
import pickle
from tqdm import tqdm
from pathlib import Path
from .dataset_utils import pc_norm, shapenet55_id_to_cat, build_loss_funcs, DataLoaderWrapper

@DATASETS.register_module()
def ShapeNet(batch_size, root, npoints, split, crop=[.25, .75], num_workers=0, scale_mode='shape_unit', mode='median', loss=None, distributed=False, shuffle=True, cats=['all']):
    mode={
            'easy': 1/4,
            'median' :1/2,
            'hard':3/4,
            'complete': 1.
        }[mode]
    if split == 'test':
        crop_range = int(mode * npoints)
    
    elif isinstance(crop, list):
        crop_range = [int(crop[0] * npoints), int(crop[1] * npoints)]
    elif isinstance(crop, float):
        crop_range = int(crop * npoints)
    else:
        raise NotImplementedError
    dataset = _ShapeNet(root, npoints, split, scale_mode=scale_mode, fine_loss=loss, cats=cats)
    if distributed:
        sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle = shuffle)
    else:
        sampler = None
    return DataLoaderWrapper(dataset, batch_size, num_workers, True, crop_range, npoints, sampler=sampler), sampler



    

    

class _ShapeNet(Dataset):
    def __init__(self, root, npoints, split, scale_mode='shape_unit', fine_loss=None, cats=['all']):
        if 'all' in cats:
            cats = list(shapenet55_id_to_cat.values())
        for cat in cats:
            if cat not in shapenet55_id_to_cat.values():
                assert False, "Unknown category"
        self.cats = cats
        self.data_root = root
        self.subset = split
        self.npoints = npoints
        self.scale_mode=scale_mode
        self.data_list_file = os.path.join(self.data_root, f'{self.subset}.txt')
        self.fine_loss_func = build_loss_funcs(fine_loss)
        self.parent_path = Path(root).parent.absolute()
        

        print(f'[DATASET] Open file {self.data_list_file}')
        with open(self.data_list_file, 'r') as f:
            lines = f.readlines()
        
        self.file_list = []
        for line in lines:
            line = line.strip()
            taxonomy_id = line.split('-')[0]
            if shapenet55_id_to_cat[taxonomy_id] in self.cats:
                model_id = line.split('-')[1].split('.')[0]
                self.file_list.append({
                    'taxonomy_id': taxonomy_id,
                    'model_id': model_id,
                    'file_path': os.path.join(self.parent_path, 'shapenet_pc', line)
                })
        print(f'[DATASET] {len(self.file_list)} instances were loaded')

    


    
        
    def __getitem__(self, idx):
        sample = self.file_list[idx]

        data = np.load(sample['file_path']).astype(np.float32)
        data = filter_points(data, self.npoints)
        data, shift, scale = pc_norm(data, self.scale_mode, to_torch=True)

        return {'taxonomy_id':sample['taxonomy_id'], 'model_id':sample['model_id'], 'shift':shift, 'scale': scale, 'pointcloud':data}

    def __len__(self):
        return len(self.file_list)
    
    def evaluate(self, results, save_num_batch, device):
        save_dict = dict()
        fine_metrics = {k:0. for k in self.fine_loss_func.keys()}
        for idx, pred_dict in enumerate(tqdm(results)):
            ref = pred_dict['ref'].to(device)
            pred = pred_dict['pred'].to(device)
            shift = pred_dict['shift'].to(device)
            scale = pred_dict['scale'].to(device)
            del pred_dict['shift']
            del pred_dict['scale']
            ref, pred = map(lambda p: p * scale + shift, [ref, pred])
            for k, v in self.fine_loss_func.items():
                fine_metrics[k] += v(ref, pred)
            if idx < save_num_batch:
                for k, v in pred_dict.items():
                    if k not in save_dict.keys():
                        save_dict[k] = []
                    save_dict[k].append(v.detach().cpu().numpy())
            
        fine_metrics = {'fine_' + str(k): v / len(results) for k, v in fine_metrics.items()}
        save_dict = {k: np.concatenate(v, axis=0) for k, v in save_dict.items()}
        return save_dict, fine_metrics
        
            
python/difffacto/datasets/shapenet55.py
diffFacto/diffFacto
import os
import torch
import torch.nn.functional as F
from difffacto.utils.registry import DATASETS
from difffacto.utils.misc import fps
from torch.utils.data import Dataset
import numpy as np
from tqdm.auto import tqdm
from .dataset_utils import pc_norm, shapenet_part_normal_cat_to_id, build_loss_funcs, DataLoaderWrapper, DataLoaderWrapperOne
from .evaluation_utils import EMD_CD, compute_all_metrics
import json
import random


@DATASETS.register_module()
def ShapeNetPart(
    batch_size, 
    root, 
    npoints, 
    num_workers=0, 
    scale_mode=None,
    eval_mode='ae',
    distributed=False, 
    shuffle=True, 
    drop_last=True,
    ):
    return ShapeNetParts(root, npoints, scale_mode=scale_mode, eval_mode=eval_mode)





@DATASETS.register_module()
class ShapeNetParts(Dataset):

    def __init__(
        self,
        root = '../datasets/shapenetcore_partanno_segmentation_benchmark_v0_normal',
        npoints=2500, 
        scale_mode=None,
        eval_mode='ae'
    ):
        self.scale_mode=scale_mode 
        self.npoints = npoints # Sampling points
        self.root = root # File root path
        self.eval_mode=eval_mode

        self.data = torch.load(self.root).numpy()

    def __getitem__(self, index):
        # resample
        point_set = self.data[index] # Sample by index
        choice = np.random.choice(point_set.shape[0], self.npoints, replace=True)
        point_set = point_set[choice] 
        point_set, shift, scale = pc_norm(point_set, self.scale_mode, to_torch=True)
        
        
        return {
            'input':point_set, 
            'ref':point_set, 
            'shift':shift,
            'scale':scale, 
            'id':index,
            } # pointset is point cloud data, cls has 16 categories, and seg is a small category corresponding to different points in the data

    def __len__(self):
        return self.data.shape[0]

    

    def evaluate(self, results, save_num_batch, device):
        save_dict = dict()
        preds = []
        refs = []
        for idx, pred_dict in enumerate(tqdm(results)):
            shift = pred_dict['shift']
            scale = pred_dict['scale']
            del pred_dict['shift']
            del pred_dict['scale']
            pred = pred_dict['pred']
            ref = pred_dict['input_ref']
            if pred.shape[1] > 2048:
                pred = fps(pred.to('cuda').contiguous(), 2048).cpu()
            if ref.shape[1] > 2048:
                ref = fps(ref.to('cuda').contiguous(), 2048).cpu()
            if self.eval_mode == 'ae':
                pred = pred * scale + shift
                ref = ref * scale + shift
            else:
                pred_max = pred.max(1)[0].reshape(-1, 1, 3) # (1, 3)
                pred_min = pred.min(1)[0].reshape(-1, 1, 3) # (1, 3)
                pred_shift = ((pred_min + pred_max) / 2).reshape(-1, 1, 3)
                pred_scale = (pred_max - pred_min).max(-1)[0].reshape(-1, 1, 1) / 2
                pred = (pred - pred_shift) / pred_scale
                ref_max = ref.max(1)[0].reshape(-1, 1, 3) # (1, 3)
                ref_min = ref.min(1)[0].reshape(-1, 1, 3) # (1, 3)
                ref_shift = ((ref_min + ref_max) / 2).reshape(-1, 1, 3)
                ref_scale = (ref_max - ref_min).max(-1)[0].reshape(-1, 1, 1) / 2
                ref = (ref - ref_shift) / ref_scale
                if self.using_whole_chair_only:
                    present = pred_dict['present'][:, :3].sum(1) == 3 
                    pred = pred[present]
                    ref = ref[present]
            preds.append(pred)
            refs.append(ref)
            pred_dict = {k: v * scale + shift if v.shape[-1] == 3 else v for k, v in pred_dict.items()}    
                     
            if idx < save_num_batch:
                for k, v in pred_dict.items():
                    if k not in save_dict.keys():
                        save_dict[k] = []
                    save_dict[k].append(v.detach().cpu().numpy())
        preds = torch.cat(preds, dim=0).cuda()
        refs = torch.cat(refs, dim=0).cuda()
        if self.eval_mode == 'ae':
            metrics = EMD_CD(preds, refs, 32)
        elif self.eval_mode == 'gen':
            # metrics = dict(l=torch.zeros(1))
            print(preds.shape, refs.shape)
            metrics = compute_all_metrics(preds, refs, 32)
        ssave_dict = dict()
        for k, v in save_dict.items():
            ssave_dict[k] = np.concatenate(v, axis=0)
        return ssave_dict, metrics

python/difffacto/datasets/shapenet_parts.py
diffFacto/diffFacto
import os
import torch
import torch.nn.functional as F
from difffacto.utils.registry import DATASETS
from difffacto.utils.misc import fps
from torch.utils.data import Dataset
import numpy as np
from tqdm.auto import tqdm
from .dataset_utils import pc_norm, shapenet_part_normal_cat_to_id, build_loss_funcs, DataLoaderWrapper, DataLoaderWrapperOne
from .evaluation_utils import EMD_CD, compute_all_metrics, compute_bbox_metric, compute_part_metric, compute_snapping_metric
import json
import random
from pointnet2_ops import pointnet2_utils


@DATASETS.register_module()
def ShapeNetSegPart(
    batch_size, 
    root, 
    npoints, 
    split, 
    num_workers=0, 
    eval_mode='ae',
    scale_mode='shape_unit', 
    part_scale_mode=None,
    distributed=False, 
    shuffle=True, 
    drop_last=True,
    class_choice='Chair',
    save_only=False,
    augment=False,
    augment_shift=False,
    augment_scale=False,
    using_whole_chair_only=False,
    clip=True,
    dropout_part=0.0):
    assert not distributed
    dataset = _ShapeNetSegParts(root, npoints, split, scale_mode=scale_mode, class_choice=class_choice, save_only=save_only, eval_mode=eval_mode, augment=augment, augment_shift=augment_shift, augment_scale=augment_scale, dropout_part=dropout_part, using_whole_chair_only=using_whole_chair_only, part_scale_mode=part_scale_mode, clip=clip)

    return DataLoaderWrapperOne(dataset, batch_size=batch_size, num_workers=num_workers, drop_last=drop_last, shuffle=shuffle), None






@DATASETS.register_module()
def ShapeNetSeg(
    batch_size, 
    root, 
    npoints, 
    split, 
    num_workers=0, 
    scale_mode='shape_unit', 
    eval_mode='ae',
    distributed=False, 
    drop_last=True,
    shuffle=True, 
    class_choice='Chair',
    save_only=False,
    augment=False, 
    augment_shift=False,
    augment_scale=False):
    assert not distributed
    dataset = _ShapeNetSeg(root, npoints, split, scale_mode=scale_mode, class_choice=class_choice, save_only=save_only, eval_mode=eval_mode, augment=augment, augment_shift=augment_shift, augment_scale=augment_scale)
    
    return DataLoaderWrapperOne(dataset, batch_size=batch_size, num_workers=num_workers, drop_last=drop_last, shuffle=shuffle), None




class _ShapeNetSeg(Dataset):

    def __init__(
        self,
        root = '../datasets/shapenetcore_partanno_segmentation_benchmark_v0_normal',
        npoints=2500, 
        split='train', 
        class_choice='Chair', 
        normal_channel=False, 
        scale_mode='shape_unit',
        eval_mode='ae',
        save_only=False,
        augment=False,
        augment_shift=False,
        augment_scale=False,
        dropout_part=0.0,
        using_whole_chair_only=False,
        part_scale_mode=None,
        clip=True,
    ):
        self.eval_mode=eval_mode
        self.part_scale_mode=part_scale_mode if part_scale_mode is not None else scale_mode
        self.augment=augment
        self.clip=clip
        if augment:
            augment_shift = augment_scale = True
        self.augment_shift = augment_shift
        self.augment_scale = augment_scale
        self.using_whole_chair_only=using_whole_chair_only
        assert not (self.augment and split == 'test')
        self.npoints = npoints # Sampling points
        self.root = root # File root path
        self.catfile = os.path.join(self.root, 'synsetoffset2category.txt') # Path corresponding to category and folder name
        self.cat = {}
        self.normal_channel = normal_channel # Whether to use rgb information
        self.scale_mode = scale_mode

        self.cat = shapenet_part_normal_cat_to_id #{'Airplane': '02691156', 'Bag': '02773838', 'Cap': '02954340', 'Car': '02958343', 'Chair': '03001627', 'Earphone': '03261776', 'Guitar': '03467517', 'Knife': '03624134', 'Lamp': '03636649', 'Laptop': '03642806', 'Motorbike': '03790512', 'Mug': '03797390', 'Pistol': '03948459', 'Rocket': '04099429', 'Skateboard': '04225987', 'Table': '04379243'}
        self.classes_original = dict(zip(self.cat, range(len(self.cat)))) #{'Airplane': 0, 'Bag': 1, 'Cap': 2, 'Car': 3, 'Chair': 4, 'Earphone': 5, 'Guitar': 6, 'Knife': 7, 'Lamp': 8, 'Laptop': 9, 'Motorbike': 10, 'Mug': 11, 'Pistol': 12, 'Rocket': 13, 'Skateboard': 14, 'Table': 15}

        self.cat = {class_choice:shapenet_part_normal_cat_to_id[class_choice]} if class_choice is not None else shapenet_part_normal_cat_to_id
        self.class_choice = class_choice
        self.save_only=save_only
        print(self.cat)
        self.noises = dict()


        self.meta = {} # Read the classified folder jason file and put their names in the list
        with open(os.path.join(self.root, 'train_test_split', 'shuffled_train_file_list.json'), 'r') as f:
            train_ids = set([str(d.split('/')[2]) for d in json.load(f)]) # '928c86eabc0be624c2bf2dc31ba1713' this is the first value
        with open(os.path.join(self.root, 'train_test_split', 'shuffled_val_file_list.json'), 'r') as f:
            val_ids = set([str(d.split('/')[2]) for d in json.load(f)])
        with open(os.path.join(self.root, 'train_test_split', 'shuffled_test_file_list.json'), 'r') as f:
            test_ids = set([str(d.split('/')[2]) for d in json.load(f)])
        for item in self.cat:
            self.meta[item] = []
            dir_point = os.path.join(self.root, self.cat[item]) # # Get the path corresponding to a folder, such as the first folder 02691156
            fns = sorted(os.listdir(dir_point))  # Get each txt file under the folder according to the path and put it in the list
            
            if split == 'trainval':
                fns = [fn for fn in fns if ((fn[0:-4] in train_ids) or (fn[0:-4] in val_ids))]
            elif split == 'train':
                fns = [fn for fn in fns if fn[0:-4] in train_ids] # Judge whether the txt file in the folder is in the training TXT. If so, the txt file obtained in fns is the file to be trained in all txt files in this category, and put it into fns
            elif split == 'val':
                fns = [fn for fn in fns if fn[0:-4] in val_ids]
            elif split == 'test':
                fns = [fn for fn in fns if fn[0:-4] in test_ids]
            elif split == 'all':
                fns = [fn for fn in fns if ((fn[0:-4] in test_ids) or (fn[0:-4] in train_ids) or (fn[0:-4] in val_ids))]
            else:
                print('Unknown split: %s. Exiting..' % (split))
                exit(-1)
            # print(os.path.basename(fns))
            for fn in fns:
                "The first i Secondary cycle  fns What you get is the third i Training eligible in folders txt Folder name"
                token = (os.path.splitext(os.path.basename(fn))[0])
                self.meta[item].append(os.path.join(dir_point, token + '.txt'))  # Generate a dictionary and combine the category name and the training path as the training data in a large class
                #After the above code is executed, you can put all the data to be trained or verified into a dictionary. The key of the dictionary is the category to which the data belongs, such as aircraft. Value is the full path of its corresponding data
                #{aircraft: [path 1, path 2.....]}

            
        # exit()
        # print(fns[0][0:-4])
        #####################################################################################################################################################
        self.datapath = []
        for item in self.cat: # self.cat is the dictionary corresponding to the category name and folder
            for fn in self.meta[item]:
                self.datapath.append((item, fn)) # Generate tuples of labels and point cloud paths, and convert the dictionary in self.met into a tuple
        print(len(self.datapath))
        self.classes = {}
        for i in self.cat.keys():
            self.classes[i] = self.classes_original[i]
        ## self.classes corresponds the name and index of the category, such as aircraft < --- > 0
        # Mapping from category ('Chair') to a list of int [10,11,12,13] as segmentation labels
        """
        shapenet There are 16 categories, and then each category has some components, such as aircraft 'Airplane': [0, 1, 2, 3] Among them, the four sub categories labeled 0 1 2 3 belong to the category of aircraft
        self.seg_classes Is to correspond the major class to the minor class
        """
        self.seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
                            'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46],
                            'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27],
                            'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40],
                            'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}

        # for cat in sorted(self.seg_classes.keys()):
        #     print(cat, self.seg_classes[cat])

        self.cache = {}  # from index to (point_set, cls, seg) tuple
        self.cache_size = 20000
        self.num_class = len(self.seg_classes[class_choice])
        self.dropout_part=dropout_part
        # self.compute_part_distribution()

    def store_noise(self, all_noises, all_ids):
        self.noises = dict()
        for noises, ids in zip(all_noises, all_ids):
            assert noises.shape[0] == len(ids)
            for i, id in enumerate(ids):
                self.noises[id.item()] = noises[i].detach().cpu().numpy()
    def compute_part_distribution(self):
        distribution = dict()
        for fn in self.datapath:
            data = np.loadtxt(fn[1]).astype(np.float32) # size 20488,7 read in this txt file, a total of 20488 points, each point xyz rgb + sub category label
            present = [0]*self.num_class
            
            seg = data[:, -1].astype(np.long) - self.seg_classes[self.class_choice][0] # Get the label of the sub category
            for i in range(self.num_class):
                idx = seg == i 
                if idx.any():
                    present[i] = 1
            if ''.join(map(str, present)) not in distribution.keys():
                distribution[''.join(map(str, present))] = 0
            distribution[''.join(map(str, present))] += 1
        print(distribution)
        print({k: v / len(self.datapath) for k, v in distribution.items()})
    def __getitem__(self, index):
        if index in self.cache: # The initial slef.cache is an empty dictionary, which is used to store the retrieved data and place it according to (point_set, cls, seg) while avoiding repeated sampling
            point_set, cls, seg,  = self.cache[index]
        else:
            fn = self.datapath[index] # The path to get the training data according to the index self.datepath is a tuple (class name, path)
            cat = self.datapath[index][0] # Get the class name
            cls = self.classes[cat] # Convert class name to index
            cls = np.array([cls]).astype(np.int32)
            data = np.loadtxt(fn[1]).astype(np.float32) # size 20488,7 read in this txt file, a total of 20488 points, each point xyz rgb + sub category label
            token = (os.path.splitext(os.path.basename(fn[1]))[0])
            
            if not self.normal_channel:  # Determine whether to use rgb information
                point_set = data[:, 0:3]
            else:
                point_set = data[:, 0:6]
            seg = data[:, -1].astype(np.long) - self.seg_classes[self.class_choice][0] # Get the label of the sub category


            if len(self.cache) < self.cache_size:
                self.cache[index] = (point_set, cls, seg, token)
        
        noise = self.noises.get(index, np.zeros(1)) #for cimle
        
        choice = np.random.choice(len(seg), self.npoints, replace=True) # Random sampling of data in a category returns the index, allowing repeated sampling
        # resample
        point_set = point_set[choice, :] # Sample by index
        seg = seg[choice]
        seg_onehot = F.one_hot(torch.from_numpy(seg), num_classes=self.num_class)
        # print(seg.shape, seg_onehot.shape)
        present = torch.zeros(self.num_class)
        part_means = np.zeros((self.num_class, 3), dtype=np.float32)
        part_scale = np.ones((self.num_class, 3), dtype=np.float32)

        point_set, shift, scale = pc_norm(point_set, self.scale_mode, to_torch=True)
        for i in range(self.num_class):
            idx = seg == i 
            if idx.any():
                present[i] = 1
                part_point_set = point_set[idx]
                part_point_set_mean = part_point_set.mean(0)
                part_point_set_std = part_point_set.std(0)
                part_means[i] = part_point_set_mean
                part_scale[i] = part_point_set_std
                
        if self.augment_shift or self.augment_scale:
            if self.augment_scale:
                rand_scale = torch.rand(1,3) / 2 + 0.7
            else:
                rand_scale = torch.ones(1,3) 
            if self.augment_shift:
                rand_shift = torch.rand(1, 3) - 0.5
            else:
                rand_shift = torch.zeros(1, 3)
            # ((o - s  - ss * scale ) / ssc * scale
            point_set = (point_set - rand_shift) / rand_scale
            shift = shift + scale * rand_shift
            scale = rand_scale * scale
        


        part_scale = part_scale.transpose(1,0)
        part_means = part_means.transpose(1,0)
        return {
            'present': present,
            'part_scale': part_scale, 
            'part_shift': part_means, 
            'input':point_set, 
            'attn_map': seg_onehot, 
            'ref_attn_map': seg_onehot, 
            'ref':point_set, 
            'class':cls, 
            'seg_mask':seg, 
            'ref_seg_mask':seg, 
            'shift':shift,
            'scale':scale, 
            'id':index,
            'token':token,
            'noise':torch.from_numpy(noise),
            } # pointset is point cloud data, cls has 16 categories, and seg is a small category corresponding to different points in the data

    def __len__(self):
        return len(self.datapath)

    
    # def evaluate(self, results, save_num_batch, device):
    #     save_dict = dict()
    #     for idx, pred_dict in enumerate(tqdm(results)):
    #         idxx = pred_dict['id']
    #         mean, logvar = pred_dict['mean'], pred_dict['logvar']
    #         for id in idxx:
    #             save_dict[id.item] = (mean, logvar)
    #     print(len(save_dict.keys()))
    #     return save_dict, dict(l=0)
    def evaluate(self, results, save_num_batch, device):
        save_dict = dict()
        preds = []
        preds_mask = []
        refs_mask = []
        refs = []
        for idx, pred_dict in enumerate(tqdm(results)):
            if self.save_only:
                # pred_dict = {k: v * scale + shift if v.shape[-1] == 3 else v for k, v in pred_dict.items()}    
                if idx < save_num_batch:
                    for k, v in pred_dict.items():
                        if k not in save_dict.keys():
                            save_dict[k] = []
                        if isinstance(v, torch.Tensor):
                            save_dict[k].append(v.detach().cpu().numpy())
                        else:
                            save_dict[k].append(v)
                continue
            shift = pred_dict['shift']
            scale = pred_dict['scale']
            del pred_dict['shift']
            del pred_dict['scale']
            
            pred = pred_dict['pred']
            pred_mask = pred_dict['pred_seg_mask']
            ref = pred_dict['input_ref']
            ref_mask = pred_dict['ref_seg_mask']
            if pred.shape[1] > 2048:
                pred, idx = fps(pred.to('cuda').contiguous(), 2048, ret_id=True).cpu()
                pred_mask = pointnet2_utils.gather_operation(pred_mask.unsqueeze(1).contiguous(), idx).reshape(-1, 2048).cpu()
            if ref.shape[1] > 2048:
                ref, ref_idx = fps(ref.to('cuda').contiguous(), 2048).cpu()
                ref_mask = pointnet2_utils.gather_operation(ref_mask.unsqueeze(1).contiguous(), ref_idx).reshape(-1, 2048).cpu()
            if self.eval_mode == 'ae':
                pred = pred * scale + shift
                ref = ref * scale + shift
            else:
                pred_max = pred.max(1)[0].reshape(-1, 1, 3) # (1, 3)
                pred_min = pred.min(1)[0].reshape(-1, 1, 3) # (1, 3)
                pred_shift = ((pred_min + pred_max) / 2).reshape(-1, 1, 3)
                pred_scale = (pred_max - pred_min).max(-1)[0].reshape(-1, 1, 1) / 2
                pred = (pred - pred_shift) / pred_scale
                ref_max = ref.max(1)[0].reshape(-1, 1, 3) # (1, 3)
                ref_min = ref.min(1)[0].reshape(-1, 1, 3) # (1, 3)
                ref_shift = ((ref_min + ref_max) / 2).reshape(-1, 1, 3)
                ref_scale = (ref_max - ref_min).max(-1)[0].reshape(-1, 1, 1) / 2
                ref = (ref - ref_shift) / ref_scale
                if self.using_whole_chair_only:
                    present = pred_dict['present'][:, :3].sum(1) == 3 
                    pred = pred[present]
                    ref = ref[present]
            preds.append(pred)
            refs.append(ref)
            preds_mask.append(pred_mask)
            refs_mask.append(ref_mask)
            pred_dict = {k: v * scale + shift if v.shape[-1] == 3 else v for k, v in pred_dict.items()}    
            # pred_dict['pred'] = pred 
            # pred_dict['input_ref'] = ref  
                     
            if idx < save_num_batch:
                for k, v in pred_dict.items():
                    if k not in save_dict.keys():
                        save_dict[k] = []
                    save_dict[k].append(v.detach().cpu().numpy())
        if self.save_only: 
            ssave_dict = dict()
            for k, v in save_dict.items():
                ssave_dict[k] = np.concatenate(v, axis=0)
            return ssave_dict, dict(l=0)
        preds = torch.cat(preds, dim=0).cuda()
        refs = torch.cat(refs, dim=0).cuda()
        preds_mask = torch.cat(preds_mask, dim=0).cuda()
        refs_mask = torch.cat(refs_mask, dim=0).cuda()
        if self.eval_mode == 'ae':
            metrics = EMD_CD(preds, refs, 32)
        elif self.eval_mode == 'gen_part':
            # metrics = dict(l=torch.zeros(1))
            print(preds.shape, refs.shape)
            snapping_metric = compute_snapping_metric(preds, preds_mask, cls=self.class_choice)
            for k, v in snapping_metric.items():
                print(f'{k}: {v}')
            oracle_snapping_metric = compute_snapping_metric(refs, refs_mask, cls=self.class_choice)
            bbox_metric = compute_bbox_metric(preds, preds_mask, refs, refs_mask, 32, n_class=self.num_class, metric='chamfer')
            part_metric = compute_part_metric(preds, preds_mask, refs, refs_mask, 32, n_class=self.num_class)
            metrics = compute_all_metrics(preds, refs, 32)
            metrics.update(snapping_metric)
            metrics.update({f"oracle_{k}":v for k, v in oracle_snapping_metric.items()})
            metrics.update(part_metric)
            metrics.update(bbox_metric)
        elif self.eval_mode == 'gen':
            # metrics = dict(l=torch.zeros(1))
            print(preds.shape, refs.shape)
            metrics = compute_all_metrics(preds, refs, 32)
        ssave_dict = dict()
        for k, v in save_dict.items():
            ssave_dict[k] = np.concatenate(v, axis=0)
        return ssave_dict, metrics

class _ShapeNetSegParts(_ShapeNetSeg):

    def __init__(
        self,
        root = '../datasets/shapenetcore_partanno_segmentation_benchmark_v0_normal',
        npoints=2500, 
        split='train', 
        class_choice='Chair', 
        normal_channel=False, 
        scale_mode='shape_unit',
        save_only=False,
        eval_mode='ae',
        augment=False,
        augment_shift=False,
        augment_scale=False,
        dropout_part=0.0,
        using_whole_chair_only=False,
        part_scale_mode=None,
        clip=True,
    ):
        super().__init__(
            root = root,
            npoints=npoints, 
            split=split, 
            class_choice=class_choice, 
            normal_channel=normal_channel, 
            scale_mode=scale_mode,
            save_only=save_only,
            eval_mode=eval_mode,
            augment=augment,
            augment_shift=augment_shift,
            augment_scale=augment_scale,
            dropout_part=dropout_part,
            using_whole_chair_only=using_whole_chair_only,
            part_scale_mode=part_scale_mode,
            clip=clip
        )
        
    def __getitem__(self, index):
        if index in self.cache: # The initial slef.cache is an empty dictionary, which is used to store the retrieved data and place it according to (point_set, cls, seg) while avoiding repeated sampling
            point_set, cls, seg, token = self.cache[index]
        else:
            fn = self.datapath[index] # The path to get the training data according to the index self.datepath is a tuple (class name, path)
            cat = self.datapath[index][0] # Get the class name
            cls = self.classes[cat] # Convert class name to index
            noise = self.noises.get(index, np.zeros(1))#for cimle
            cls = np.array([cls]).astype(np.int32)
            data = np.loadtxt(fn[1]).astype(np.float32) # size 20488,7 read in this txt file, a total of 20488 points, each point xyz rgb + sub category label
            token = (os.path.splitext(os.path.basename(fn[1]))[0])
            
            if not self.normal_channel:  # Determine whether to use rgb information
                point_set = data[:, 0:3]
            else:
                point_set = data[:, 0:6]
            seg = data[:, -1].astype(np.long) - self.seg_classes[self.class_choice][0] # Get the label of the sub category


            if len(self.cache) < self.cache_size:
                self.cache[index] = (point_set, cls, seg, token)
        
        # Random sampling of data in a category returns the index, allowing repeated sampling
        # resample
        noise = self.noises.get(index, np.zeros(1))#for cimle
        
        shifts = np.zeros((self.num_class, 3))
        scales = np.ones((self.num_class, 3))
        present = torch.zeros(self.num_class)
        choice = np.random.choice(point_set.shape[0], self.npoints, replace=True)
        ori_out = point_set[choice] 
        
        seg = seg[choice]
        
        ori_out, shift, scale = pc_norm(ori_out, self.scale_mode, to_torch=False) 
        out = np.zeros_like(ori_out)
        for i in range(self.num_class):
            idx = seg == i
            if idx.sum() >= 10: # if less than 10 points we ignore
                part_point_set = ori_out[idx]
                part_point_set_mean = part_point_set.mean(0)
                part_point_set_std = part_point_set.std(0)
                if np.any(part_point_set_std == 0.0):
                    present[i] = 0
                    part_point_set_std = np.ones_like(part_point_set_std)
                else:
                    present[i] = 1
                shifted_chosen, pshift, pscale = pc_norm(part_point_set, self.part_scale_mode, to_torch=False, clip=self.clip) 
                shifts[i] = pshift[0]
                scales[i] = pscale[0]
                out[idx] = shifted_chosen
            elif np.any(idx):
                # we change its label to their nearest neighbors
                part_point = ori_out[idx] # n, 3
                rest_point = ori_out[~idx]
                rest_seg = seg[~idx]
                dist = (part_point[:, None] - rest_point[None]) ** 2
                seg_idx = dist.sum(-1).argmin(1)
                seg[idx] = np.take_along_axis(rest_seg, seg_idx, axis=0)
                
        seg_onehot = F.one_hot(torch.from_numpy(seg), num_classes=self.num_class)
        part_idx = np.random.rand(self.num_class) < self.dropout_part
        drop_outed_present = present.clone()
        drop_outed_present[part_idx] = 0
        ori_out, shift, scale = map(lambda t: torch.from_numpy(t).float(), [ori_out, shift, scale])
        out, shifts, scales = map(lambda t: torch.from_numpy(t).float(), [out, shifts, scales])
        if self.augment_shift or self.augment_scale:
            if self.augment_scale:
                rand_scale = torch.rand(1,3) / 2 + 0.7
            else:
                rand_scale = torch.ones(1,3) 
            if self.augment_shift:
                rand_shift = torch.rand(1, 3) - 0.5
            else:
                rand_shift = torch.zeros(1, 3)
            # (o + s) * scale
            ori_out = (ori_out + rand_shift) * rand_scale
            shift = shift  + scale * rand_shift
            scale = rand_scale * scale

        # print(seg.shape, seg_onehot.shape)
        assert (shifts.shape == shifts.shape == (self.num_class, 3))
        assert ori_out.shape == (self.npoints, 3) == out.shape
        assert seg_onehot.shape == (self.npoints, self.num_class)
        
        




        return {
            'present': present,
            'dp_present': drop_outed_present,
            'part_scale':scales.transpose(0,1), 
            'part_shift':shifts.transpose(0,1), 
            'input':out, 
            'ref_attn_map': seg_onehot, 
            'attn_map': seg_onehot, 
            'ref':ori_out, 
            'class':cls, 
            'ref_seg_mask':seg, 
            'token':token,
            'seg_mask':seg, 
            'shift':shift, 
            'scale':scale, 
            'id':index,
            'noise':torch.from_numpy(noise),
            } # pointset is point cloud data, cls has 16 categories, and seg is a small category corresponding to different points in the data

python/difffacto/datasets/shapenet_seg.py
diffFacto/diffFacto
import os
import os.path as osp
import torch
import torch.nn.functional as F 
from difffacto.utils.registry import DATASETS
from torch.utils.data import Dataset, DataLoader
import numpy as np
from tqdm.auto import tqdm
from .dataset_utils import pc_norm, shapenet_part_normal_cat_to_id, build_loss_funcs, DataLoaderWrapperOne, worker_init_fn, augment
from .evaluation_utils import EMD_CD, compute_all_metrics
import json
import pickle
import random

import os


@DATASETS.register_module()
def ShapeNetSegSuperSegment(batch_size,data_root, part='pn_aware', split='train', n_class=4, scale_mode='shape_unit', num_workers=0, distributed=False, augment=False, contrastive_learning=False, vertical_only=False, eval_mode='ae', shift_only=False, augment_attn=False, normalize_attn=False, augment_prob=0.5, global_shift_prob=0.0):
    assert not distributed
    dataset = _ShapeNetSegSuperSegment(data_root, split, part, scale_mode, n_class, augment=augment,contrastive_learning=contrastive_learning, vertical_only=vertical_only, eval_mode=eval_mode, shift_only=shift_only, augment_attn=augment_attn, normalize_attn=normalize_attn, augment_prob=augment_prob, global_shift_prob=global_shift_prob)
    loader = DataLoaderWrapperOne(dataset, batch_size=batch_size, num_workers=num_workers, drop_last=True)
    return loader, None


class _ShapeNetSegSuperSegment(Dataset):
    def __init__(self, data_root, split, part='pn_aware',scale_mode='shape_unit', n_class=4, augment=False, contrastive_learning=False, vertical_only=False, eval_mode='ae', shift_only=False, augment_attn=False, normalize_attn=False, augment_prob=0.5, global_shift_prob=0.0):
        super().__init__()
        assert eval_mode in ['ae', 'gen']
        self.eval_mode = eval_mode
        self.segs_data = pickle.load(open(os.path.join(data_root, f"shapenet_pointcloud_{part}.pkl"), "rb"))
        self.attn_map = pickle.load(open(os.path.join(data_root, f"shapenet_label_{part}.pkl"), "rb"))
        self.scale_mode = scale_mode
        self.shift_only=shift_only
        self.augment_attn=augment_attn
        self.augment_prob=augment_prob
        self.normalize_attn=normalize_attn
        self.global_shift_prob=global_shift_prob
        self.n_class=n_class
        self.augment=augment
        self.contrastive_learning=contrastive_learning
        self.vertical_only=vertical_only
        self.get_stats()
        if split == 'test':
            assert not (augment or contrastive_learning)
        num_train = int(len(self.segs_data) * 0.8)
        if split == 'train':
            self.segs_data = self.segs_data[:num_train]
            self.attn_map = self.attn_map[:num_train]
        if split == 'test':
            self.segs_data = self.segs_data[num_train:]
            self.attn_map = self.attn_map[num_train:]
        print(self.segs_data.shape, self.attn_map.shape)
        
        
        
    def get_stats(self):
        self.stats = dict() 
        self.stats['mean'] = self.segs_data.reshape(-1, 3).mean(0)
        self.stats['std'] = self.segs_data.reshape(-1).std()
        part_means = [[] for _ in range(self.n_class)]
        # print(len(part_means))
        for i in range(self.segs_data.shape[0]):
            geos = self.segs_data[i]
            attn_map = self.attn_map[i]
            seg_mask = attn_map.argmax(1)
            # print(seg_mask.shape)
            # print(i)
            for j in range(self.n_class):
                part = geos[seg_mask == j]
                
                if part.shape[0] > 0:
                    # print(part.mean(0))
                    part_means[j].append(part.mean(0))
                    # print(part.mean(0))
                    # print(sum(part_means[j]))
        self.stats["part_means"] = np.stack([np.stack(part_means[i], axis=0).mean(0) for i in range(self.n_class)], axis=0)
        print("#####PART MEANS ######")
        print(self.stats["part_means"])
        print("#####PART MEANS ######")
            

    def __getitem__(self, idx):
        out_dict=dict()
        geos = self.segs_data[idx]
        geos, shift, scale = pc_norm(geos, self.scale_mode, to_torch=True, stats=self.stats)
        attn_map = torch.from_numpy(self.attn_map[idx])
        if random.random() < self.global_shift_prob:
            rand_shift = torch.rand(1, 3) - 0.5
            if self.vertical_only:
                rand_shift[:, [0,2]] = 0.
            geos = geos + rand_shift
            shift -= rand_shift / scale
        if self.normalize_attn:
            attn_map = F.softmax(attn_map, 1)
        seg_mask = attn_map.max(1)[1]
        if self.augment_attn:
            attn_map = attn_map + torch.randn_like(attn_map) * 0.2 - 0.1
            attn_map= attn_map.clamp(0., 1.)
        
        if self.augment:
            input, part_scale, part_shifts = augment(geos, seg_mask, self.n_class, vertical_only=self.vertical_only, shift_only=self.shift_only)
            out_dict['input'] = input
            out_dict['attn_map'] = attn_map
            out_dict['ref'] = geos
            out_dict['seg_mask'] = seg_mask
            out_dict['shift'] = shift 
            out_dict['scale'] = scale
            out_dict['part_scale'] = part_scale
            out_dict['part_shift'] = part_shifts
            

        elif self.contrastive_learning:
            if random.random() < self.augment_prob:
                pos_input = augment(geos, seg_mask, self.n_class, vertical_only=self.vertical_only)
            neg_id = random.randrange(0, len(self.segs_data))
            if neg_id == idx: 
                neg_id = (idx + 1) % len(self.segs_data)
            neg_input, neg_shift, neg_scale = pc_norm(self.segs_data[neg_id], self.scale_mode, to_torch=True, stats=self.stats)
            neg_attn_map = torch.from_numpy(self.attn_map[neg_id])

            combined_attn_map = torch.cat([attn_map.unsqueeze(0).expand(2, -1, -1), neg_attn_map.unsqueeze(0)], dim=0)
            combined_seg_mask = combined_attn_map.max(-1)[1]
            out_dict['shift'] = torch.stack([shift, shift, neg_shift], dim=0)
            out_dict['scale'] = torch.stack([scale, scale, neg_scale], dim=0)
            out_dict['attn_map'] = combined_attn_map
            out_dict['input'] = torch.stack([geos, pos_input, neg_input], dim=0)
            out_dict['ref'] = torch.stack([geos, geos, neg_input], dim=0)
            out_dict['seg_mask'] = combined_seg_mask
            return out_dict
        else:
            out_dict['input'] = geos
            out_dict['ref'] = geos
            out_dict['attn_map'] = attn_map
            out_dict['seg_mask'] = seg_mask
            out_dict['shift'] = shift 
            out_dict['scale'] = scale
            out_dict['part_scale'] = torch.ones([self.n_class, 3])
            out_dict['part_shift'] = torch.zeros([self.n_class, 3])
        
        out_dict['global_anchor_mean'] = (torch.from_numpy(self.stats["part_means"]) - shift) / scale

        return out_dict
    def __len__(self):
        return len(self.segs_data)

    def evaluate(self, results, save_num_batch, device):
        save_dict = dict()
        preds = []
        refs = []
        for idx, pred_dict in enumerate(tqdm(results)):
            shift = pred_dict['shift']
            scale = pred_dict['scale']
            del pred_dict['shift']
            del pred_dict['scale']
            if self.eval_mode == 'ae':
                pred = pred_dict['pred'] * scale + shift 
                ref = pred_dict['input_ref'] * scale + shift 
            preds.append(pred)
            refs.append(ref)
            
            pred_dict = {k: v * scale + shift if v.shape[-1] == 3 else v for k, v in pred_dict.items()}               
            if idx < save_num_batch:
                for k, v in pred_dict.items():
                    if k not in save_dict.keys():
                        save_dict[k] = []
                    save_dict[k].append(v.detach().cpu().numpy())
        preds = torch.cat(preds, dim=0).cuda()
        refs = torch.cat(refs, dim=0).cuda()
        print(preds.shape, refs.shape)
        if self.eval_mode == 'ae':
            metrics = EMD_CD(preds, refs, 32)
        elif self.eval_mode == 'gen':
            metrics = compute_all_metrics(preds, refs, 32)
        save_dict = {k: np.concatenate(v, axis=0) for k, v in save_dict.items()}
        return save_dict, metrics

python/difffacto/datasets/shapenet_sup_seg.py
diffFacto/diffFacto
import os
import os.path as osp
import torch
from difffacto.utils.registry import DATASETS
from torch.utils.data import Dataset, DataLoader
import numpy as np
from tqdm.auto import tqdm
from .dataset_utils import pc_norm, shapenet_part_normal_cat_to_id, build_loss_funcs, DataLoaderWrapperOne, worker_init_fn, augment
import json
import pickle
import random
import os


@DATASETS.register_module()
def ShapeNetSegSuperSegmentParglot(batch_size,data_root, part='pn_aware', split='train', n_class=4, scale_mode='shape_unit', num_workers=0, distributed=False, augment=False, vertical_only=False):
    assert not distributed
    dataset = _ShapeNetSegSuperSegmentPartglot(data_root, split, part, scale_mode, n_class, augment=augment, vertical_only=vertical_only)
    loader = DataLoaderWrapperOne(dataset, batch_size=batch_size, num_workers=num_workers, drop_last=True)
    return loader, None


class _ShapeNetSegSuperSegmentPartglot(Dataset):
    def __init__(self, data_root, split, part='pn_aware',scale_mode='shape_unit', n_class=4, augment=False, vertical_only=False):
        super().__init__()
        data = pickle.load(open(os.path.join(data_root, f"partglot_shapenet_seg_out_{part}_raw.pkl"), "rb"))
        self.segs_data = data["geos"].astype(np.float32)
        self.geo_mask = data["geos_masks"].astype(np.float32)
        self.attn_map = data['attn_maps'].astype(np.float32)
        self.scale_mode = scale_mode
        self.n_class=n_class
        self.augment=augment
        self.vertical_only=vertical_only
        if split == 'test':
            assert not augment
        num_train = int(len(self.segs_data) * 0.8)
        if split == 'train':
            self.segs_data = self.segs_data[:num_train]
            self.attn_map = self.attn_map[:num_train]
        if split == 'test':
            self.segs_data = self.segs_data[num_train:]
            self.attn_map = self.attn_map[num_train:]
        print(self.segs_data.shape, self.attn_map.shape)
        self.get_stats()
        
        
    def get_stats(self):
        self.stats = dict() 
        self.stats['mean'] = self.segs_data.reshape(-1, 3).mean(0)
        self.stats['std'] = self.segs_data.reshape(-1).std()

    def __getitem__(self, idx):
        out_dict=dict()
        geos = self.segs_data[idx]
        npoint = geos.shape[1]
        geo_mask = self.geo_mask[idx]
        part_geos, shift, scale = pc_norm(geos[geo_mask == 1].reshape(-1, 3), self.scale_mode, to_torch=True, stats=self.stats)
        geos[geo_mask == 1] = part_geos.reshape(-1, npoint, 3)
        attn_map = torch.from_numpy(self.attn_map[idx])
        seg_mask = attn_map.max(0)[1]
        
        if self.augment:
            # not working
            input = augment(geos, seg_mask, self.n_class, vertical_only=self.vertical_only)
            out_dict['input'] = input
            out_dict['attn_map'] = attn_map
            out_dict['ref'] = geos
            out_dict['seg_mask'] = seg_mask
            out_dict['shift'] = shift 
            out_dict['scale'] = scale
        #else:
        out_dict['input'] = geos
        out_dict['ref'] = geos
        out_dict['geo_mask'] = geo_mask
        out_dict['attn_map'] = attn_map
        out_dict['seg_mask'] = seg_mask
        out_dict['shift'] = shift 
        out_dict['scale'] = scale
        


        return out_dict
    def __len__(self):
        return len(self.segs_data)

    def evaluate(self, results, save_num_batch, device):
        save_dict = dict()
        fine_metrics = {'metric':0.}
        for idx, pred_dict in enumerate(tqdm(results)):
            shift = pred_dict['shift'][11].to(device)
            scale = pred_dict['scale'][11].to(device)
            del pred_dict['shift']
            del pred_dict['scale']
            for k, v in pred_dict.items():
                if v.shape[-1] == 3:
                    v = v.to(device) * scale.unsqueeze(0) + shift.unsqueeze(0)
            if idx < save_num_batch:
                for k, v in pred_dict.items():
                    if k not in save_dict.keys():
                        save_dict[k] = []
                    save_dict[k].append(v.detach().cpu().numpy())
            
        save_dict = {k: np.concatenate(v, axis=0) for k, v in save_dict.items()}
        return save_dict, fine_metrics

python/difffacto/datasets/shapenet_sup_seg_partglot.py
diffFacto/diffFacto
from .emd.emd_module import EMD
from .chamfer_dist import ChamferDistanceL1, ChamferDistanceL2, ChamferDistanceL2_split
from .common import L2Loss, L1Loss, SpectralContrastiveLoss, SmoothCrossEntropy, CrossEntropy
python/difffacto/metrics/__init__.py
diffFacto/diffFacto
# -*- coding: utf-8 -*-
# @Author: Thibault GROUEIX
# @Date:   2019-08-07 20:54:24
# @Last Modified by:   Haozhe Xie
# @Last Modified time: 2019-12-18 15:06:25
# @Email:  cshzxie@gmail.com

import torch

import chamfer
from difffacto.utils.registry import METRICS


class ChamferFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, xyz1, xyz2):
        dist1, dist2, idx1, idx2 = chamfer.forward(xyz1, xyz2)
        ctx.save_for_backward(xyz1, xyz2, idx1, idx2)

        return dist1, dist2

    @staticmethod
    def backward(ctx, grad_dist1, grad_dist2):
        xyz1, xyz2, idx1, idx2 = ctx.saved_tensors
        grad_xyz1, grad_xyz2 = chamfer.backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2)
        return grad_xyz1, grad_xyz2

@METRICS.register_module()
class ChamferDistanceL2(torch.nn.Module):
    f''' Chamder Distance L2
    '''
    def __init__(self, ignore_zeros=False, reduce=True):
        super().__init__()
        self.ignore_zeros = ignore_zeros
        self.reduce=reduce

    def forward(self, xyz1, xyz2):
        batch_size = xyz1.size(0)
        if batch_size == 1 and self.ignore_zeros:
            non_zeros1 = torch.sum(xyz1, dim=2).ne(0)
            non_zeros2 = torch.sum(xyz2, dim=2).ne(0)
            xyz1 = xyz1[non_zeros1].unsqueeze(dim=0)
            xyz2 = xyz2[non_zeros2].unsqueeze(dim=0)

        dist1, dist2 = ChamferFunction.apply(xyz1, xyz2)
        print(dist1, dist2)
        if self.reduce:
            dist1 = torch.mean(dist1)
            dist2 = torch.mean(dist2)
        return dist1 + dist2
    
@METRICS.register_module()
class ChamferDistanceL2_split(torch.nn.Module):
    f''' Chamder Distance L2
    '''
    def __init__(self, ignore_zeros=False, reduce=True):
        super().__init__()
        self.ignore_zeros = ignore_zeros
        self.reduce=reduce

    def forward(self, xyz1, xyz2):
        batch_size = xyz1.size(0)
        if batch_size == 1 and self.ignore_zeros:
            non_zeros1 = torch.sum(xyz1, dim=2).ne(0)
            non_zeros2 = torch.sum(xyz2, dim=2).ne(0)
            xyz1 = xyz1[non_zeros1].unsqueeze(dim=0)
            xyz2 = xyz2[non_zeros2].unsqueeze(dim=0)

        dist1, dist2 = ChamferFunction.apply(xyz1, xyz2)
        if self.reduce:
            dist1 = torch.mean(dist1)
            dist2 = torch.mean(dist2)
        return dist1, dist2
    
@METRICS.register_module()
class ChamferDistanceL1(torch.nn.Module):
    f''' Chamder Distance L1
    '''
    def __init__(self, ignore_zeros=False, reduce=True):
        super().__init__()
        self.ignore_zeros = ignore_zeros
        self.reduce=reduce

    def forward(self, xyz1, xyz2):
        batch_size = xyz1.size(0)
        if batch_size == 1 and self.ignore_zeros:
            non_zeros1 = torch.sum(xyz1, dim=2).ne(0)
            non_zeros2 = torch.sum(xyz2, dim=2).ne(0)
            xyz1 = xyz1[non_zeros1].unsqueeze(dim=0)
            xyz2 = xyz2[non_zeros2].unsqueeze(dim=0)

        dist1, dist2 = ChamferFunction.apply(xyz1, xyz2)
        # import pdb
        # pdb.set_trace()
        dist1 = torch.sqrt(dist1)
        dist2 = torch.sqrt(dist2)
        if self.reduce:
            dist1 = torch.mean(dist1)
            dist2 = torch.mean(dist2)
        return (dist1 + dist2)/2


python/difffacto/metrics/chamfer_dist/__init__.py
diffFacto/diffFacto
# -*- coding: utf-8 -*-
# @Author: Haozhe Xie
# @Date:   2019-12-10 10:38:01
# @Last Modified by:   Haozhe Xie
# @Last Modified time: 2019-12-26 14:21:36
# @Email:  cshzxie@gmail.com
#
# Note:
# - Replace float -> double, kFloat -> kDouble in chamfer.cu

import os
import sys
import torch
import unittest


from torch.autograd import gradcheck

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir)))
from difffacto.metrics.chamfer_dist import ChamferFunction


class ChamferDistanceTestCase(unittest.TestCase):
    def test_chamfer_dist(self):
        x = torch.rand(4, 64, 3)
        y = torch.rand(4, 128, 3)
        x.requires_grad = True
        y.requires_grad = True
        print(gradcheck(ChamferFunction.apply, [x.cuda(), y.cuda()]))



if __name__ == '__main__':
    # unittest.main()
    import pdb
    x = torch.rand(32,128,3)
    y = torch.rand(32,128,3)
    test = ChamferDistanceTestCase()
    test.test_chamfer_dist()
    #pdb.set_trace()

python/difffacto/metrics/chamfer_dist/test.py
diffFacto/diffFacto
import torch
import torch.nn.functional as F
from difffacto.utils.registry import METRICS
import math 
import random

def triplet_loss(anchor, pos, neg, thresh=0.1):
    pos_dist = (anchor - pos) ** 2
    neg_dist = (anchor - neg) ** 2
    l = pos_dist - neg_dist + thresh
    l = l.mean(1)
    return torch.clamp(l, min=0), pos_dist.mean(), neg_dist.mean()


@METRICS.register_module()
def dis_loss(d_real, d_fake, loss_type="wgan", weight=1., **kwargs):
    if loss_type.lower() == "wgan":
        loss_fake = d_fake.mean()
        loss_real = d_real.mean()
        wg_loss_orig = loss_fake - loss_real
        wg_loss = wg_loss_orig * weight
        return wg_loss, {
            "wgan_dis_l": wg_loss.clone().detach(),
            "wgan_dis_l_orig": wg_loss_orig.clone().detach(),
            "wgan_dis_l_real": loss_real.clone().detach(),
            "wgan_dis_l_fake": loss_fake.clone().detach()
        }
    elif loss_type.lower() == "hinge":
        d_loss_real = -torch.min(d_real - 1, d_real * 0).mean()
        d_loss_fake = -torch.min(-d_fake - 1, d_fake * 0).mean()
        d_correct = (d_real >= 0.).float().sum() + (d_fake < 0.).float().sum()
        d_acc = d_correct / float(d_real.size(0) + d_fake.size(0))

        d_loss = d_loss_real + d_loss_fake
        loss = d_loss * weight
        return loss, {
            "dis_acc": d_acc.clone().detach(),
            "dis_correct": d_correct.clone().detach(),
            "l_real": d_loss_real.clone().detach(),
            "l_fake": d_loss_fake.clone().detach(),
        }
    else:
        raise NotImplementedError("Not implement: %s" % loss_type)


@METRICS.register_module()
def gradient_penalty(x_real, x_fake, d_real, d_fake,
                     weight=1., gp_type='zero_center', interpolated=None, d_interpolated=None, seps=1e-8):
    if gp_type == "zero_center":
        bs = d_real.size(0)
        grad = torch.autograd.grad(
            outputs=d_real, inputs=x_real,
            grad_outputs=torch.ones_like(d_real).to(d_real),
            create_graph=True, retain_graph=True)[0]
        # [grad] should be either (B, D) or (B, #points, D)
        grad = grad.reshape(bs, -1)
        grad_norm = gp_orig = torch.sqrt(torch.sum(grad ** 2, dim=1)).mean()
        gp = gp_orig ** 2. * weight
        return gp, {
            'gp': gp.clone().detach().cpu(),
            'gp_orig': gp_orig.clone().detach().cpu(),
            'grad_norm': grad_norm.clone().detach().cpu()
        }
    elif gp_type == "interpolated":
        bs = d_real.size(0)
        gradients = torch.autograd.grad(outputs=d_interpolated, inputs=interpolated,
                               grad_outputs=torch.ones(d_interpolated.size()).cuda(),
                               create_graph=True, retain_graph=True)[0]

        # Gradients have shape (batch_size, num_channels, img_width, img_height),
        # so flatten to easily take norm per example in batch
        gradients = gradients.view(bs, -1)

        gradients_norm = torch.sqrt(torch.sum(gradients ** 2, dim=1) + seps)

        # Return gradient penalty
        gp =  weight * ((gradients_norm - 1) ** 2).mean()
        return gp, {
            'gp': gp.clone().detach().cpu(),
            'grad_norm': gradients_norm.mean().clone().detach().cpu()
        }
    else:
        raise NotImplemented("Invalid gp type:%s" % gp_type)

@METRICS.register_module()
def gen_loss(d_real, d_fake, loss_type="wgan", weight=1., **kwargs):
    if loss_type.lower() == "wgan":
        wg_loss_orig = - d_fake.mean()
        wg_loss = wg_loss_orig * weight
        return wg_loss, {
            "wgan_gen_l_orig": wg_loss_orig.clone().detach(),
        }
    elif loss_type.lower() == "hinge":
        g_loss = -d_fake.mean()
        d_correct = (d_real >= 0.).float().sum() + (d_fake < 0.).float().sum()
        d_acc = d_correct / float(d_real.size(0) + d_fake.size(0))

        loss = weight * g_loss
        return loss, {
            "dis_acc": d_acc.clone().detach(),
            "dis_correct": d_correct.clone().detach(),
        }

    else:
        raise NotImplementedError("Not implement: %s" % loss_type)
    
@METRICS.register_module()
class SmoothCrossEntropy:
    # https://arxiv.org/pdf/1512.00567.pdf
    def __init__(self, alpha=0.1):
        self.alpha=alpha
    def __call__(self, pred, target, ):
        n_class = pred.size(1)
        one_hot = torch.zeros_like(pred).scatter_(1, target.view(-1, 1), 1)
        one_hot = (
            one_hot * ((1.0 - self.alpha) + self.alpha / n_class) + (1.0 - one_hot) * self.alpha / n_class
        )  # smoothed

        log_prb = F.log_softmax(pred, dim=1)
        loss = -(one_hot * log_prb).sum(dim=1)
        return torch.mean(loss)
    
@METRICS.register_module()
class NegativeGaussianLogLikelihood:
    def __init__(self, dim=3, reduction='mean'):
        self.dim=dim 
        self.log_z = 0.5 * dim * math.log(2 * math.pi)
        self.reduction=reduction
    def __call__(self, z, mu, var=None):
        assert z.shape[-1] == mu.shape[-1] == self.dim
        if isinstance(var, float):
            var_term = 0
            var = 1.
        else:
            var_term = 0.5 * torch.sum(torch.log(var), dim=-1)
        
        if self.reduction == 'mean':
            return self.log_z + var_term + ((z - mu) ** 2 / var).mean() / 2
        elif self.reduction is None:
            return self.log_z + var_term + ((z - mu) ** 2 / var) / 2


@METRICS.register_module()
class CrossEntropy:
    def __init__(self, n_class=2, reduce=True):
        self.reduce=reduce
        self.n_class=n_class
    def __call__(self, pred, target):
        pred_prob = torch.sigmoid(pred)
        target_prob = torch.sigmoid(target)
        pred_dist = torch.stack([pred_prob, 1 - pred_prob], dim=1)
        target_dist = torch.stack([target_prob, 1 - target_prob], dim=1)

        log_prb = torch.log(pred_dist)
        loss = -(target_dist * log_prb).sum(dim=1)
        return torch.mean(loss) if self.reduce else loss

@METRICS.register_module()
class L2Loss:
    def __init__(self, reduction='mean'):
        self.reduction=reduction
    def __call__(self, source, target, var=None):
        if self.reduction == 'mean':
            diff = torch.pow(source - target, 2).mean()
        elif self.reduction is None:
            diff = torch.pow(source - target, 2).mean(-1)
        return diff

@METRICS.register_module()
class L1Loss:
    def __init__(self, reduction='mean'):
        self.reduction=reduction
    def __call__(self, source, target):
        return torch.abs(source - target).mean()


@METRICS.register_module()
class SpectralContrastiveLoss:
    def __init__(self):
        pass
    def __call__(self, input):
        anchor, pos, neg = input[:, 0], input[:, 1], input[:, 2] # [B, dim]
        anchor, pos, neg = map(lambda t: F.normalize(t, 2, 1), [anchor, pos, neg])
        B, dim=anchor.shape
        anchor_pos_dot = torch.bmm(anchor.unsqueeze(1), pos.unsqueeze(2)).reshape(B)
        anchor_neg_dot = torch.bmm(anchor.unsqueeze(1), neg.unsqueeze(2)).reshape(B)
        l = -2. * anchor_pos_dot + anchor_neg_dot ** 2
        return l.mean()
python/difffacto/metrics/common.py
diffFacto/diffFacto
# EMD approximation module (based on auction algorithm)
# memory complexity: O(n)
# time complexity: O(n^2 * iter) 
# author: Minghua Liu

# Input:
# xyz1, xyz2: [#batch, #points, 3]
# where xyz1 is the predicted point cloud and xyz2 is the ground truth point cloud 
# two point clouds should have same size and be normalized to [0, 1]
# #points should be a multiple of 1024
# #batch should be no greater than 512
# eps is a parameter which balances the error rate and the speed of convergence
# iters is the number of iteration
# we only calculate gradient for xyz1

# Output:
# dist: [#batch, #points],  sqrt(dist) -> L2 distance 
# assignment: [#batch, #points], index of the matched point in the ground truth point cloud
# the result is an approximation and the assignment is not guranteed to be a bijection

import time
import numpy as np
import torch
from torch import nn
from torch.autograd import Function
import emd
from difffacto.utils.registry import METRICS




class emdFunction(Function):
    @staticmethod
    def forward(ctx, xyz1, xyz2, eps, iters):

        batchsize, n, _ = xyz1.size()
        _, m, _ = xyz2.size()

        assert(n == m)
        assert(xyz1.size()[0] == xyz2.size()[0])
        assert(n % 1024 == 0)
        assert(batchsize <= 512)

        xyz1 = xyz1.contiguous().float().cuda()
        xyz2 = xyz2.contiguous().float().cuda()
        dist = torch.zeros(batchsize, n, device='cuda').contiguous()
        assignment = torch.zeros(batchsize, n, device='cuda', dtype=torch.int32).contiguous() - 1
        assignment_inv = torch.zeros(batchsize, m, device='cuda', dtype=torch.int32).contiguous() - 1
        price = torch.zeros(batchsize, m, device='cuda').contiguous()
        bid = torch.zeros(batchsize, n, device='cuda', dtype=torch.int32).contiguous()
        bid_increments = torch.zeros(batchsize, n, device='cuda').contiguous()
        max_increments = torch.zeros(batchsize, m, device='cuda').contiguous()
        unass_idx = torch.zeros(batchsize * n, device='cuda', dtype=torch.int32).contiguous()
        max_idx = torch.zeros(batchsize * m, device='cuda', dtype=torch.int32).contiguous()
        unass_cnt = torch.zeros(512, dtype=torch.int32, device='cuda').contiguous()
        unass_cnt_sum = torch.zeros(512, dtype=torch.int32, device='cuda').contiguous()
        cnt_tmp = torch.zeros(512, dtype=torch.int32, device='cuda').contiguous()

        emd.forward(xyz1, xyz2, dist, assignment, price, assignment_inv, bid, bid_increments, max_increments, unass_idx, unass_cnt, unass_cnt_sum, cnt_tmp, max_idx, eps, iters)

        ctx.save_for_backward(xyz1, xyz2, assignment)
        return dist, assignment

    @staticmethod
    def backward(ctx, graddist, gradidx):
        xyz1, xyz2, assignment = ctx.saved_tensors
        graddist = graddist.contiguous()

        gradxyz1 = torch.zeros(xyz1.size(), device='cuda').contiguous()
        gradxyz2 = torch.zeros(xyz2.size(), device='cuda').contiguous()

        emd.backward(xyz1, xyz2, gradxyz1, graddist, assignment)
        return gradxyz1, gradxyz2, None, None

@METRICS.register_module()
class EMD(nn.Module):
    def __init__(self, eps, iters, dist_only=False):
        super(EMD, self).__init__()
        self.eps = eps 
        self.iters = iters
        self.dist_only = dist_only

    def forward(self, input1, input2):
        if self.dist_only:
            return torch.sqrt(emdFunction.apply(input1, input2, self.eps, self.iters)[0]).mean(1)
        else:
            return emdFunction.apply(input1, input2, self.eps, self.iters)

def test_emd():
    x1 = torch.rand(20, 8192, 3).cuda() # please normalize your point cloud to [0, 1]
    x2 = torch.rand(20, 8192, 3).cuda()
    emd = EMD()
    start_time = time.perf_counter()
    dis, assigment = emd(x1, x2, 0.002, 10000) # 0.005, 50 for training 
    print("Input_size: ", x1.shape)
    print("Runtime: %lfs" % (time.perf_counter() - start_time))
    print("EMD: %lf" % np.sqrt(dis.cpu()).mean())
    print("|set(assignment)|: %d" % assigment.unique().numel())
    assigment = assigment.cpu().numpy()
    assigment = np.expand_dims(assigment, -1)
    x2 = np.take_along_axis(x2, assigment, axis = 1)
    d = (x1 - x2) * (x1 - x2)
    print("Verified EMD: %lf" % np.sqrt(d.cpu().sum(-1)).mean())        

python/difffacto/metrics/emd/emd_module.py
diffFacto/diffFacto
from .diffusions import *
from .encoders import *
from .losses import *
from .networks import *
from .decomposers import *
from .segmentors import *
from .modules import *
from .discriminators import *
from .generators import *
python/difffacto/models/__init__.py
diffFacto/diffFacto
from .mixer import ComponentMixer, ComponentMixerV2
from .transformer import PartCodeTransformer, PartCodeTransformerV2
python/difffacto/models/decomposers/__init__.py
diffFacto/diffFacto
import torch 
from torch import nn 
import torch.nn.functional as F
import math

class SinusoidalEmbedding3D(nn.Module):
    def __init__(self, out_dim):
        super().__init__()
        self.out_dim=out_dim 
        self.intermediate_dim = (out_dim // 6) * 6

        self.scaling_factor = torch.ones(self.intermediate_dim // 6, dtype=torch.float32) * (10000 ** (1 / self.intermediate_dim))
        self.scaling_factor = torch.pow(self.scaling_factor, torch.arange(self.intermediate_dim // 6) * 6)

    def forward(self, x):
        # x[B, N, 3]
        B, N, _ = x.shape
        out = torch.zeros([B, N, self.out_dim]).cuda()
        cos_x, sin_x = torch.cos(x.unsqueeze(3) / self.scaling_factor.reshape(1, 1, 1, -1).to(x.device)), torch.sin(x.unsqueeze(3) / self.scaling_factor.reshape(1, 1, 1, -1).to(x.device)) #[B, N, 3, D // 6]
        sinusoidal_pe = torch.stack([sin_x, cos_x], dim=-1).reshape(B, N, self.intermediate_dim) #[B, N, 3, D // 6, 2] -> [B, N, D]
        out[:, :, :self.intermediate_dim] += sinusoidal_pe
        return out
        




class MultiHeadSelfAttention(nn.Module):
    def __init__(self, in_channel, out_channel, qk_dim=512, nheads=8, qkv_bias=True, scale=True):
        super(MultiHeadSelfAttention, self).__init__()
        assert out_channel % nheads == 0
        self.ndim = qk_dim // nheads
        self.nheads=nheads
        self.out_channel=out_channel
        self.Wq = nn.Linear(in_channel, qk_dim, bias=qkv_bias)
        self.Wk = nn.Linear(in_channel, qk_dim, bias=qkv_bias)
        self.Wv = nn.Linear(in_channel, out_channel, bias=qkv_bias)
        self.scale= 1 / math.sqrt(self.ndim) if scale else 1. 

    def forward(self, x):
        ori_x = x
        B, N, C = x.shape
        Q, K, V = self.Wq(x), self.Wk(x), self.Wv(x)
        V = V.reshape(B, N, self.nheads, self.out_channel // self.nheads)
        Q, K = map(lambda w: w.reshape(B, N, self.nheads, self.ndim), [Q, K])
        QK = torch.matmul(Q.permute([0, 2, 1, 3]), K.permute([0, 2, 3, 1])) * self.scale # [B, head, N, N]
        QK = F.softmax(QK, dim=3)
        message = torch.matmul(QK, V.permute([0, 2, 1, 3])).transpose(1,2).reshape(B, N, self.out_channel)
        return message


class GAT(nn.Module):
    def __init__(self, in_channel, out_channel, nheads, dropout=0., alpha=0.2):
        """Dense version of GAT."""
        super(GAT, self).__init__()
        self.dropout = dropout
        assert out_channel % nheads == 0
        nhid = out_channel // nheads

        self.attentions = [GraphAttentionLayer(in_channel, nhid, dropout=dropout, alpha=alpha) for _ in range(nheads)]
        for i, attention in enumerate(self.attentions):
            self.add_module('attention_{}'.format(i), attention)

    def forward(self, x):
        ori_x = x
        x = F.dropout(x, self.dropout, training=self.training)
        x = torch.cat([att(x) for att in self.attentions], dim=2)
        x = F.dropout(x, self.dropout, training=self.training)
        return x


class GraphAttentionLayer(nn.Module):
    """
    Simple GAT layer, similar to https://arxiv.org/abs/1710.10903
    """
    def __init__(self, in_features, out_features, dropout, alpha):
        super(GraphAttentionLayer, self).__init__()
        self.dropout = dropout
        self.in_features = in_features
        self.out_features = out_features
        self.alpha = alpha

        self.W = nn.Parameter(torch.empty(size=(in_features, out_features)))
        nn.init.xavier_uniform_(self.W.data, gain=1.414)
        self.a = nn.Parameter(torch.empty(size=(2*out_features, 1)))
        nn.init.xavier_uniform_(self.a.data, gain=1.414)

        self.leakyrelu = nn.LeakyReLU(self.alpha)

    def forward(self, h):
        Wh = torch.einsum("bni, ij->bnj", h, self.W) # h.shape: (B, N, in_features), Wh.shape: (B, N, out_features)
        attention = self._prepare_attentional_mechanism_input(Wh) # [B, N, N]

        attention = F.softmax(attention, dim=2)
        attention = F.dropout(attention, self.dropout, training=self.training)
        h_prime = torch.einsum("bij, bjn->bin", attention, Wh)

        return F.elu(h_prime)

    def _prepare_attentional_mechanism_input(self, Wh):
        # Wh.shape (B, N, out_feature)
        # self.a.shape (2 * out_feature, 1)
        # Wh1&2.shape (B, N, 1)
        # e.shape (B, N, N)
        Wh1 = torch.einsum("bni, ij->bnj", Wh, self.a[:self.out_features, :])
        Wh2 = torch.einsum("bni, ij->bnj", Wh, self.a[self.out_features:, :])
        # broadcast add
        e = Wh1 + Wh2.transpose(1,2)
        return self.leakyrelu(e)


if __name__ == '__main__':
    pe_embedding = SinusoidalEmbedding3D(6)
    x = torch.tensor([[1,1,1], [0,1,0], [0,0,1]],dtype=torch.float32).reshape(1, 3, 3)
    embed = pe_embedding(x)
    print(embed)
    print(embed.shape)

python/difffacto/models/decomposers/common.py
diffFacto/diffFacto
import torch 
from torch import nn 
import numpy as np
from difffacto.utils.registry import DECOMPOSERS 
import torch.nn.functional as F
from .common import GAT, MultiHeadSelfAttention, SinusoidalEmbedding3D
from einops import rearrange, einsum, repeat

def zero_module(module):
    """
    Zero out the parameters of a module and return it.
    """
    for p in module.parameters():
        p.detach().zero_()
    return module


def Normalize(in_channels):
    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)

class CrossAttention(nn.Module):
    def __init__(self, query_dim, context_dim, heads=8, dim_head=64, dropout=0.):
        super().__init__()
        inner_dim = dim_head * heads

        self.scale = dim_head ** -0.5
        self.heads = heads

        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, query_dim),
            nn.Dropout(dropout)
        )

    def forward(self, x, context):
        h = self.heads

        q = self.to_q(x)
        k = self.to_k(context)
        v = self.to_v(context)

        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
        sim = einsum(q, k, 'b i d, b j d -> b i j') * self.scale

        # attention, what we cannot get enough of
        attn = sim.softmax(dim=-1)

        out = einsum(attn, v, 'b i j, b j d -> b i d')
        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
        return self.to_out(out)

# feedforward
class GEGLU(nn.Module):
    def __init__(self, dim_in, dim_out):
        super().__init__()
        self.proj = nn.Linear(dim_in, dim_out * 2)

    def forward(self, x):
        x, gate = self.proj(x).chunk(2, dim=-1)
        return x * F.gelu(gate)


class FeedForward(nn.Module):
    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
        super().__init__()
        inner_dim = int(dim * mult)
        dim_out = dim
        project_in = nn.Sequential(
            nn.Linear(dim, inner_dim),
            nn.GELU()
        ) if not glu else GEGLU(dim, inner_dim)

        self.net = nn.Sequential(
            project_in,
            nn.Dropout(dropout),
            nn.Linear(inner_dim, dim_out)
        )

    def forward(self, x):
        return self.net(x)

def identity(*args, **kwargs):
    return nn.Identity()
@DECOMPOSERS.register_module()
class ComponentMixer(nn.Module):
    def __init__(
        self, 
        num_anchors, 
        part_latent_dim, 
        point_dim=3, 
        include_attention=False, 
        nheads=8, 
        use_graph_attention=True, 
        use_abs_pe=False,
        include_global_feature=False,
        global_mlp_type=0,
        normalize_latent=False,
        deprecation=False,
        mlp_type=0,
        norm=None,
        mlp_norm=-1, 
        regressor_norm=-1,
        embed_channel=128,
        use_scale_shift_norm=False,
        pe_dp = 0.2,
        pe_norm=None,
        res=True,
        attn_ln=True,
    ):
        super().__init__()
        self.num_anchors=num_anchors 
        self.global_mlp_type=global_mlp_type
        self.point_dim = point_dim
        self.include_attention=include_attention
        self.use_abs_pe=use_abs_pe
        self.include_global_feature=include_global_feature
        self.normalize_latent=normalize_latent
        self.mlp_type=mlp_type
        if mlp_norm == -1 and regressor_norm == -1:
            if norm == 'bn':
                mlp_norm = nn.BatchNorm1d
                regressor_norm = nn.BatchNorm1d 
            elif norm == 'gn':
                mlp_norm = Normalize
                regressor_norm = Normalize 
            elif norm is None:
                mlp_norm = identity
                regressor_norm = identity
        else:
            if mlp_norm == 'bn':
                mlp_norm = nn.BatchNorm1d
            elif mlp_norm == 'gn':
                mlp_norm = Normalize
            elif mlp_norm is None:
                mlp_norm = identity
            if regressor_norm == 'bn':
                regressor_norm = nn.BatchNorm1d
            elif regressor_norm == 'gn':
                regressor_norm = Normalize
            elif regressor_norm is None:
                regressor_norm = identity
                
        if mlp_type == 1:
            self.mlp = nn.Sequential(
                mlp_norm(part_latent_dim),
                nn.ReLU(inplace=True),
                nn.Conv1d(part_latent_dim, 256, 1),
                mlp_norm(256),
                nn.ReLU(inplace=True),
                nn.Conv1d(256, part_latent_dim, 1),
            )
            self.anchor_regressor = nn.Sequential(
                regressor_norm(part_latent_dim),
                nn.ReLU(inplace=True),
                nn.Linear(part_latent_dim, 128 ),
                regressor_norm(128),
                nn.ReLU(inplace=True),
                nn.Linear(128 , point_dim * num_anchors)
            )
        elif mlp_type == 0:
            self.mlp = nn.Sequential(
                nn.Linear(part_latent_dim * num_anchors, 1024),
                mlp_norm(1024),
                nn.ReLU(inplace=True),
                nn.Linear(1024, 1024),
                mlp_norm(1024),
                nn.ReLU(inplace=True),
                nn.Linear(1024, point_dim * num_anchors)
            )
        self.deprecation=deprecation
        if deprecation:
            self.mlp_type=-1
            print("THIS VERSION IS DEPRECATED USE MLP TYPE INSTEAD")
            self.mlp = nn.Sequential(
                nn.Linear(part_latent_dim * num_anchors, 1024),
                nn.ReLU(inplace=True),
                nn.Linear(1024, 1024),
                nn.ReLU(inplace=True),
                nn.Linear(1024, point_dim * num_anchors)
            )
            if self.use_abs_pe:
                self.pe_encoder = SinusoidalEmbedding3D(part_latent_dim)
        
        if self.include_attention:
            self.res=res
            self.use_scale_shift_norm=use_scale_shift_norm
            if pe_norm == 'bn':
                pe_norm = nn.BatchNorm1d
            elif pe_norm == 'gn':
                pe_norm = Normalize
            elif pe_norm is None:
                pe_norm = identity
            if self.use_abs_pe:
                self.pe_encoder = SinusoidalEmbedding3D(embed_channel)
                self.emb_layers = nn.Sequential(
                    nn.SiLU(),
                    nn.Linear(
                        embed_channel,
                        2 * part_latent_dim if use_scale_shift_norm else part_latent_dim,
                    )
                )
                self.out_layers = nn.Sequential(
                    pe_norm(part_latent_dim),
                    nn.SiLU(),
                    nn.Dropout(p=pe_dp),
                    zero_module(
                        nn.Conv1d(part_latent_dim, part_latent_dim, 1)
                    ) if self.res else nn.Conv1d(part_latent_dim, part_latent_dim, 1),
                )

            self.attention = GAT(part_latent_dim, part_latent_dim, nheads) if use_graph_attention else MultiHeadSelfAttention(part_latent_dim, part_latent_dim, nheads=nheads)
            self.ff = FeedForward(part_latent_dim, part_latent_dim, glu=True)
            self.norm1 = nn.LayerNorm(part_latent_dim) if attn_ln else nn.Identity()
            self.norm2 = nn.LayerNorm(part_latent_dim) if attn_ln else nn.Identity()

        if self.include_global_feature:
            if self.global_mlp_type == 0:
                self.aggregator = nn.Sequential(
                    nn.Linear(part_latent_dim * num_anchors, 1024),
                    nn.ReLU(inplace=True),
                    nn.Linear(1024, 1024),
                    nn.ReLU(inplace=True),
                    nn.Linear(1024, part_latent_dim)
                )
            elif self.global_mlp_type == 1:
                self.aggregator = nn.Sequential(
                    nn.BatchNorm1d(part_latent_dim),
                    nn.ReLU(inplace=True),
                    nn.Conv1d(part_latent_dim, 256, 1),
                    nn.BatchNorm1d(256),
                    nn.ReLU(inplace=True),
                    nn.Conv1d(256, part_latent_dim, 1),
                )
                
                self.aggregator_mlp = nn.Sequential(
                    nn.Linear(part_latent_dim, 256),
                    nn.BatchNorm1d(part_latent_dim),
                    nn.ReLU(inplace=True),
                    nn.Linear(256, 256),
                    nn.BatchNorm1d(256),
                    nn.ReLU(inplace=True),
                    nn.Linear(256, part_latent_dim, 1),
                )
                

    def forward(self, part_latent):
        B = part_latent.shape[0]
        if self.mlp_type == 0:
            coarse = self.mlp(part_latent.reshape(B, -1)).reshape(B, self.num_anchors, self.point_dim)
        elif self.mlp_type == 1:
            coarse_feat = self.mlp(part_latent.reshape(B, self.num_anchors, -1).transpose(1,2)).max(-1)[0]
            coarse = self.anchor_regressor(coarse_feat).reshape(B, self.num_anchors, self.point_dim)

        if self.deprecation:
            coarse = self.mlp(part_latent.reshape(B, -1)).reshape(B, self.num_anchors, self.point_dim)
            if self.use_abs_pe:
                pe_embedding = self.pe_encoder(coarse)
                part_latent = part_latent + pe_embedding.detach()

        if self.include_attention:
            if self.use_abs_pe:
                pe_embedding = self.emb_layers(self.pe_encoder(coarse)).transpose(1,2)
                _part_latent = part_latent.transpose(1,2)
                if self.use_scale_shift_norm:
                    emb_norm, emb_rest = self.out_layers[0], self.out_layers[1:]
                    scale, shift = torch.chunk(pe_embedding, 2, dim=1)
                    part_latent_out = emb_norm(_part_latent) * (1 + scale) + shift
                    part_latent_out = emb_rest(part_latent_out).transpose(1,2)
                else:
                    part_latent_out = _part_latent + pe_embedding
                    part_latent_out = self.out_layers(part_latent_out).transpose(1,2)
                
                part_latent = part_latent + part_latent_out if self.res else part_latent_out
            part_latent = self.attention(self.norm1(part_latent)) + part_latent if self.res else self.attention(self.norm1(part_latent))
            part_latent = self.ff(self.norm2(part_latent)) + part_latent if self.res else self.ff(self.norm2(part_latent))
            
        if self.include_global_feature:
            if self.global_mlp_type == 0:
                global_feature = self.aggregator(part_latent.reshape(B, -1))
            elif self.global_mlp_type == 1:
                global_feature = self.aggregator(part_latent.transpose(1,2))
                global_feature = torch.max(global_feature, dim=-1)[0]
                global_feature = self.aggregator_mlp(global_feature)
        else:
            global_feature = None
            
        if self.normalize_latent:
            global_feature = F.normalize(global_feature)
            part_latent = F.normalize(part_latent.reshape(B* self.num_anchors, -1)).reshape(B, self.num_anchors, -1)
        
        
        return global_feature, coarse, part_latent




@DECOMPOSERS.register_module()
class ComponentMixerV2(nn.Module):
    def __init__(
        self, 
        num_anchors, 
        part_latent_dim,
        embed_channel, 
        use_scale_shift_norm=True,
        point_dim=3, 
        include_attention=False, 
        nheads=8, 
        use_graph_attention=True, 
        include_global_feature=False,
        dropout=0.2,
        normalize_latent=False,
        correct_anchors=True,
        mult=1,
        part_dp_prob=0.,
        global_dp_prob=0.
    ):
        super().__init__()
        self.num_anchors=num_anchors 
        self.correct_anchors=correct_anchors
        self.point_dim = point_dim
        self.include_attention=include_attention
        self.include_global_feature=include_global_feature
        self.normalize_latent=normalize_latent
        self.use_scale_shift_norm=use_scale_shift_norm
        self.part_latent_dropout = nn.Dropout(part_dp_prob)
        self.global_latent_dropout = nn.Dropout(global_dp_prob)

        
        self.mlp = nn.Sequential(
            Normalize(part_latent_dim),
            nn.ReLU(inplace=True),
            nn.Conv1d(part_latent_dim, 128 * mult, 1),
            Normalize(128 * mult),
            nn.ReLU(inplace=True),
            nn.Conv1d(128 * mult, part_latent_dim, 1),
        )
        self.anchor_regressor = nn.Sequential(
            Normalize(part_latent_dim* mult),
            nn.ReLU(inplace=True),
            nn.Linear(part_latent_dim, 128 * mult),
            Normalize(part_latent_dim* mult),
            nn.ReLU(inplace=True),
            nn.Linear(128 * mult, point_dim * num_anchors)
        )

        self.pe_encoder = SinusoidalEmbedding3D(embed_channel)
        self.emb_layers = nn.Sequential(
            nn.SiLU(),
            nn.Linear(
                embed_channel,
                2 * part_latent_dim if use_scale_shift_norm else part_latent_dim,
            )
        )
        self.out_layers = nn.Sequential(
            Normalize(part_latent_dim),
            nn.SiLU(),
            nn.Dropout(p=dropout),
            zero_module(
                nn.Conv1d(part_latent_dim, part_latent_dim, 1)
            ),
        )

        if self.include_attention:
            self.attention=GAT(part_latent_dim, part_latent_dim, nheads) if use_graph_attention else MultiHeadSelfAttention(part_latent_dim, part_latent_dim, nheads=nheads)
            self.ff = FeedForward(part_latent_dim, part_latent_dim, glu=True)
            self.norm1 = nn.LayerNorm(part_latent_dim)
            self.norm2 = nn.LayerNorm(part_latent_dim)

        if self.include_global_feature:
            self.aggregator = nn.Sequential(
                nn.Conv1d(part_latent_dim * 2, 256 * mult, 1),
                Normalize(256 * mult),
                nn.ReLU(inplace=True),
                nn.Conv1d(256 * mult, 256 * mult, 1),
                Normalize(256 * mult),
                nn.ReLU(inplace=True),
                nn.Conv1d(256 * mult, part_latent_dim, 1)
            )
        if self.correct_anchors:
            self.anchor_corrector = CrossAttention(self.point_dim, part_latent_dim, 8, 32)

    def forward(self, part_latent):
        B = part_latent.shape[0]
        part_latent = part_latent.reshape(B, self.num_anchors, -1).transpose(1,2)
        coarse_feat = self.mlp(part_latent).max(-1)[0]
        coarse = self.anchor_regressor(coarse_feat).reshape(B, self.num_anchors, self.point_dim)

        pe_embedding = self.emb_layers(self.pe_encoder(coarse)).transpose(1,2)
        if self.use_scale_shift_norm:
            emb_norm, emb_rest = self.out_layers[0], self.out_layers[1:]
            scale, shift = torch.chunk(pe_embedding, 2, dim=1)
            part_latent_out = emb_norm(part_latent) * (1 + scale) + shift
            part_latent_out = emb_rest(part_latent_out)
        else:
            part_latent_out = part_latent + pe_embedding
            part_latent_out = self.out_layers(part_latent_out)
        
        part_latent = part_latent + part_latent_out
        part_latent = part_latent.transpose(1,2)
        if self.include_attention:
            part_latent = self.attention(self.norm1(part_latent)) + part_latent
            part_latent = self.ff(self.norm2(part_latent)) + part_latent
        part_latent = part_latent.transpose(1,2)
        if self.include_global_feature:
            global_feature = self.aggregator(torch.cat([part_latent, coarse_feat.unsqueeze(2).expand(-1, -1, self.num_anchors)], dim=1)).max(2)[0]
            global_feature = self.global_latent_dropout(global_feature)
        else:
            global_feature = None
        part_latent = part_latent.transpose(1,2)
        part_latent = self.part_latent_dropout(part_latent)
        if self.correct_anchors:
            coarse = coarse + self.anchor_corrector(coarse, part_latent)
        if self.normalize_latent:
            global_feature = F.normalize(global_feature)
            part_latent = F.normalize(part_latent.reshape(B* self.num_anchors, -1)).reshape(B, self.num_anchors, -1)
        return global_feature, coarse, part_latent


python/difffacto/models/decomposers/mixer.py
diffFacto/diffFacto


"""
Adapted from: https://github.com/openai/openai/blob/55363aa496049423c37124b440e9e30366db3ed6/orc/orc/diffusion/vit.py
"""


import math
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

import torch
import torch.nn as nn
from difffacto.utils.registry import DECOMPOSERS
from difffacto.utils.misc import checkpoint
from einops import rearrange, repeat

def exists(val):
    return val is not None

def init_linear(l, stddev):
    nn.init.normal_(l.weight, std=stddev)
    if l.bias is not None:
        nn.init.constant_(l.bias, 0.0)


class MultiheadAttention(nn.Module):
    def __init__(
        self,
        *,
        device: torch.device,
        dtype: torch.dtype,
        n_ctx: int,
        width: int,
        heads: int,
        init_scale: float,
    ):
        super().__init__()
        self.n_ctx = n_ctx
        self.width = width
        self.heads = heads
        self.c_qkv = nn.Linear(width, width * 3, device=device, dtype=dtype)
        self.c_proj = nn.Linear(width, width, device=device, dtype=dtype)
        self.attention = QKVMultiheadAttention(device=device, dtype=dtype, heads=heads, n_ctx=n_ctx)
        init_linear(self.c_qkv, init_scale)
        init_linear(self.c_proj, init_scale)

    def forward(self, x, mask=None):
        x = self.c_qkv(x)
        x = checkpoint(self.attention, (x, mask), (), False)
        x = self.c_proj(x)
        return x


class MLP(nn.Module):
    def __init__(self, *, device: torch.device, dtype: torch.dtype, width: int, init_scale: float):
        super().__init__()
        self.width = width
        self.c_fc = nn.Linear(width, width * 4, device=device, dtype=dtype)
        self.c_proj = nn.Linear(width * 4, width, device=device, dtype=dtype)
        self.gelu = nn.GELU()
        init_linear(self.c_fc, init_scale)
        init_linear(self.c_proj, init_scale)

    def forward(self, x):
        return self.c_proj(self.gelu(self.c_fc(x)))


class QKVMultiheadAttention(nn.Module):
    def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int, n_ctx: int):
        super().__init__()
        self.device = device
        self.dtype = dtype
        self.heads = heads
        self.n_ctx = n_ctx

    def forward(self, qkv, mask=None):
        bs, n_ctx, width = qkv.shape
        attn_ch = width // self.heads // 3
        scale = 1 / math.sqrt(math.sqrt(attn_ch))
        qkv = qkv.view(bs, n_ctx, self.heads, -1)
        q, k, v = torch.split(qkv, attn_ch, dim=-1)
        weight = torch.einsum(
            "bthc,bshc->bhts", q * scale, k * scale
        )  # More stable with f16 than dividing afterwards
        if exists(mask):
            mask = repeat(mask, 'b j -> b h i j', i=n_ctx, h=self.heads)
            max_neg_value = -torch.finfo(weight.dtype).max
            weight.masked_fill_(~mask, max_neg_value)
        wdtype = weight.dtype
        weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
        return torch.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1)


class ResidualAttentionBlock(nn.Module):
    def __init__(
        self,
        *,
        device: torch.device,
        dtype: torch.dtype,
        n_ctx: int,
        width: int,
        heads: int,
        init_scale: float = 1.0,
    ):
        super().__init__()

        self.attn = MultiheadAttention(
            device=device,
            dtype=dtype,
            n_ctx=n_ctx,
            width=width,
            heads=heads,
            init_scale=init_scale,
        )
        self.ln_1 = nn.LayerNorm(width, device=device, dtype=dtype)
        self.mlp = MLP(device=device, dtype=dtype, width=width, init_scale=init_scale)
        self.ln_2 = nn.LayerNorm(width, device=device, dtype=dtype)

    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor]=None):
        x = x + self.attn(self.ln_1(x), mask=mask)
        x = x + self.mlp(self.ln_2(x))
        return x


class Transformer(nn.Module):
    def __init__(
        self,
        *,
        device: torch.device,
        dtype: torch.dtype,
        n_ctx: int,
        width: int,
        layers: int,
        heads: int,
        init_scale: float = 0.25,
    ):
        super().__init__()
        self.n_ctx = n_ctx
        self.width = width
        self.layers = layers
        init_scale = init_scale * math.sqrt(1.0 / width)
        self.resblocks = nn.ModuleList(
            [
                ResidualAttentionBlock(
                    device=device,
                    dtype=dtype,
                    n_ctx=n_ctx,
                    width=width,
                    heads=heads,
                    init_scale=init_scale,
                )
                for _ in range(layers)
            ]
        )

    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor]=None):
        for block in self.resblocks:
            x = block(x, mask=mask)
        return x

@DECOMPOSERS.register_module()
class PartCodeTransformer(nn.Module):
    def __init__(
        self,
        *,
        device: torch.device,
        input_channels: int = 3,
        output_channels: int = 3,
        n_ctx: int = 1024,
        width: int = 512,
        layers: int = 12,
        heads: int = 8,
        init_scale: float = 0.25,
        class_cond: bool = True,
        param_cond: bool = True,
        param_dim: int = 6,
        use_mask_in_transformer=False,
    ):
        super().__init__()
        dtype=torch.float32
        self.dtype=torch.float32
        input_channels += int(param_cond) * param_dim + int(class_cond) * n_ctx
        self.input_channels = input_channels
        self.output_channels = output_channels
        self.n_ctx = n_ctx
        self.use_mask_in_transformer=use_mask_in_transformer
        self.class_cond = class_cond
        self.param_cond = param_cond
        self.ln_pre = nn.LayerNorm(width, device=device, dtype=dtype)
        self.backbone = Transformer(
            device=device,
            dtype=dtype,
            n_ctx=n_ctx,
            width=width,
            layers=layers,
            heads=heads,
            init_scale=init_scale,
        )
        
        self.ln_post = nn.LayerNorm(width, device=device, dtype=dtype)
        self.input_proj = nn.Linear(input_channels, width, device=device, dtype=dtype)
        self.output_proj = nn.Linear(width, output_channels, device=device, dtype=dtype)
        self.aggregate = MLP(device=device, dtype=dtype, width=width, init_scale=init_scale)
        self.class_embedding = nn.Embedding(n_ctx, n_ctx, device=device, dtype=dtype)
        self.width=width
        with torch.no_grad():
            self.output_proj.weight.zero_()
            self.output_proj.bias.zero_()

    def forward(self, x: torch.Tensor, params: List[torch.Tensor], mask: torch.Tensor):
        """
        :param x: an [N x C x T] tensor.
        :param t: an [N] tensor.
        :return: an [N x C' x T] tensor.
        """
        assert x.shape[-1] == self.n_ctx
        if self.param_cond:
            x = torch.cat([x] + params, dim=1)
        class_embed = self.class_embedding(torch.arange(self.n_ctx).to(x.device)).unsqueeze(0).repeat_interleave(x.shape[0], dim=0)
        if self.class_cond:
            x = torch.cat([x, class_embed.transpose(1,2)], dim=1)
        assert self.input_channels == x.shape[1]
        h =  self._forward_with_cond(x, [], mask=mask.to(bool) if self.use_mask_in_transformer else None)
        h = h * mask.unsqueeze(1) # zero out the part code generated from nothing
        z = self.get_global_from_part(h, mask)
        return z, h

    def get_global_from_part(self, h: torch.Tensor, mask: torch.Tensor):
        h = h + (1 - mask.unsqueeze(1)) * -100000
        z = self.aggregate(h.max(2)[0])
        assert z.shape == (h.shape[0], self.output_channels)
        return z
    
    def _forward_with_cond(
        self, x: torch.Tensor, cond_as_token: List[Tuple[torch.Tensor, bool]], mask: Optional[torch.Tensor]=None
    ) -> List[torch.Tensor]:
        B = x.shape[0]
        h = self.input_proj(x.permute(0, 2, 1))  # NCL -> NLC
        for emb, as_token in cond_as_token:
            if not as_token:
                h = h + emb[:, None]
        extra_tokens = [
            (emb[:, None] if len(emb.shape) == 2 else emb)
            for emb, as_token in cond_as_token
            if as_token
        ]
        if len(extra_tokens):
            h = torch.cat(extra_tokens + [h], dim=1)

        h = self.ln_pre(h)
        h = self.backbone(h, mask=mask)
        h = self.ln_post(h)
        if len(extra_tokens):
            h = h[:, sum(h.shape[1] for h in extra_tokens) :]
        h = self.output_proj(h)
        
        
        assert h.shape == (B, self.n_ctx, self.output_channels)
        return h.transpose(1,2)


@DECOMPOSERS.register_module()
class PartCodeTransformerV2(PartCodeTransformer):
    def __init__(self,
        *,
        device: torch.device,
        input_channels: int = 3,
        output_channels: int = 3,
        n_ctx: int = 1024,
        width: int = 512,
        layers: int = 12,
        heads: int = 8,
        init_scale: float = 0.25,
        class_cond: bool = True,
        param_cond: bool = True,
        param_dim: int = 6,
        post_mlp: bool=False
        ):
        assert output_channels % n_ctx == 0
        out_ch_per_part = output_channels // n_ctx
        self.new_out_ch = output_channels
        super().__init__(
            device=device,
            input_channels=input_channels,
            output_channels=out_ch_per_part,
            n_ctx=n_ctx,
            width=width,
            layers=layers,
            heads=heads,
            init_scale=init_scale,
            class_cond=class_cond,
            param_cond=param_cond,
            param_dim=param_dim
        )
        if post_mlp:
            self.post_mlp = MLP(device=device, dtype=self.dtype, width=self.new_out_ch, init_scale=init_scale)
        else:
            self.post_mlp = None
        
    def get_global_from_part(self, h: torch.Tensor, mask: torch.Tensor):
        B, C, N = h.shape
        assert C*N == self.new_out_ch and N == self.n_ctx
        return h.reshape(B, self.new_out_ch) if self.post_mlp is None else self.post_mlp(h.reshape(B, self.new_out_ch))
    
python/difffacto/models/decomposers/transformer.py
diffFacto/diffFacto
from .point_diffusion import PointDiffusion
# from .latent_diffusion import LDM
from .anchored_diffusion import AnchoredDiffusion
from .nets import *
from . import diffusion_utils
python/difffacto/models/diffusions/__init__.py
diffFacto/diffFacto
import torch 
from torch.nn import Module
from difffacto.utils.registry import DIFFUSIONS, NETS, build_from_cfg
from difffacto.utils.constants import ModelMeanType, ModelVarType, LossType
from difffacto.utils.misc import *
from difffacto.utils.constants import *
from torch.distributions.multivariate_normal import MultivariateNormal
import math
import numpy as np
from .diffusion_utils import extract_into_tensor, betas_for_alpha_bar, from_numpy_to_device, get_cov_from_params, normal_kl

@DIFFUSIONS.register_module()
class AnchoredDiffusion(Module):

    def __init__(
        self, 
        net, 
        num_timesteps, 
        beta_1, 
        beta_T, 
        k=1.,
        res=True,
        mode='linear', 
        use_beta=True, 
        rescale_timesteps=False, 
        loss_type='mse', 
        model_mean_type='epsilon',
        model_var_type='fixed_small', 
        scale_loss=False,
        clip_xstart=False,
        include_anchors=True,
        include_cov=False,
        learn_anchor=True,
        learn_variance=False,
        classifier_weight=1.,
        guidance=False,

        ddim_sampling=False,
        ddim_nsteps=10,
        ddim_discretize='uniform',
        ddim_eta=1.
    ):
        super().__init__()
        assert mode in ('linear', 'cosine')
        self.clip_xstart=clip_xstart
        self.res=res
        self.model = build_from_cfg(net, NETS)
        self.model_mean_type = model_mean_dict[model_mean_type]
        self.model_var_type = model_var_dict[model_var_type]
        self.loss_type = model_loss_dict[loss_type]
        self.scale_loss = scale_loss
        self.learn_anchor=learn_anchor
        self.learn_variance=learn_variance
        self.include_cov=include_cov 
        self.ddim_sampling=ddim_sampling
        self.beta_1 = beta_1
        self.beta_T = beta_T
        self.use_beta = use_beta
            
        self.rescale_timesteps = rescale_timesteps
        self.num_timesteps = int(num_timesteps)
        if mode == 'linear':
            betas = np.linspace(self.beta_1, self.beta_T, num=self.num_timesteps, dtype=np.float64)
        elif mode == 'cosine':
            betas = betas_for_alpha_bar(self.num_timesteps, lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,)
            
        self.betas = betas
        assert len(betas.shape) == 1, "betas must be 1-D"
        assert (betas > 0).all() and (betas <= 1).all()
        self.classifier_weight=classifier_weight
        self.guidance=guidance
        

        alphas = 1.0 - betas
        self.alphas_cumprod = np.cumprod(alphas, axis=0)
        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)

        # calculations for diffusion q(x_t | x_{t-1}) and others
        if isinstance(k, list):
            self.k = np.array(k)
        else:
            self.k = np.array([k] * 3).astype(np.float32)
        self.k_squared = self.k
        self.log_k = np.log(self.k)
        self.include_anchors=include_anchors
        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)

        # calculations for posterior q(x_{t-1} | x_t, x_0)
        self.posterior_variance = (
            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )
        # log calculation clipped because the posterior variance is 0 at the
        # beginning of the diffusion chain.
        self.posterior_log_variance_clipped = np.log(
            np.append(self.posterior_variance[1], self.posterior_variance[1:])
        )
        self.posterior_mean_coef1 = (
            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )
        self.posterior_mean_coef2 = (
            (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod)
        )
        self.posterior_mean_coef3 = (
            1.0 + ((np.sqrt(self.alphas_cumprod) - 1.) * ( np.sqrt(self.alphas_cumprod_prev) + np.sqrt(alphas)))
            / (1.0 - self.alphas_cumprod)
        )
        
        if self.ddim_sampling:
            self.ddim_eta = ddim_eta
            self.xt_dir_coeff = np.sqrt(1. - self.alphas_cumprod - ddim_eta * ddim_eta * self.posterior_variance)
            if ddim_discretize == 'uniform':
                skip = self.num_timesteps // ddim_nsteps
                self.steps = list(range(0, self.num_timesteps, skip))
            elif ddim_discretize == 'quad':
                self.steps = (np.linspace(0., math.sqrt(self.num_timesteps * 0.8), ddim_nsteps) ** 2).astype(np.int32)
                self.steps = self.steps.tolist()
            else:
                NotImplementedError(ddim_discretize)
        else:
            self.steps = list(range(self.num_timesteps))

        
        
    # def q_mean_variance(self, x_start, t, anchors):
    #     """
    #     ## NOT USED WRONG
    #     Get the distribution q(x_t | x_0).
    #     :param x_0: the [N x 3 x n] tensor of noiseless inputs.
    #     :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
    #     :param x_0: the [N x 3 x n] tensor of anchors the tensor of shape as x_0 to which x_0's diffuses
    #     :return: A tuple (mean, variance, log_variance), all of x_start's shape.
    #     """
    #     mean = (
    #         extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * (x_start - anchors) + anchors
    #     )
    #     variance = extract_into_tensor((1.0 - self.alphas_cumprod), t, x_start.shape) * from_numpy_to_device(self.k_squared, x_start.device).reshape(1, 3, 1)
    #     log_variance = extract_into_tensor(
    #         self.log_one_minus_alphas_cumprod, t, x_start.shape
    #     ) + 2 * from_numpy_to_device(self.log_k, x_start.device).reshape(1, 3, 1)
    #     return mean, variance, log_variance
    
    def q_sample(self, x_start, t, anchors, noise=None, variance=None):
        """
        Diffuse the data for a given number of diffusion steps.
        In other words, sample from q(x_t | x_0).
        :param x_0: the initial data batch.
        :param anchors: anchors for each point in x_0
        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
        :param noise: if specified, the split-out normal noise.
        :return: A noisy version of x_0.
        """
        if self.learn_variance:
            assert variance is not None 
            assert variance.shape == anchors.shape
        if not self.learn_anchor:
            anchors = anchors * 0.0
        if noise is None:
            noise = torch.randn_like(x_start)
        assert noise.shape == x_start.shape
        B, C, N = x_start.shape
        if not self.learn_variance:
            variance = from_numpy_to_device(self.k_squared, x_start.device).reshape(1, 3, 1).expand(B, -1, N)
        L = torch.sqrt(variance)
        return (
            extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * (x_start - anchors) + anchors
            + extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
            * L * noise)
        
    def q_posterior_mean(self, x_start, x_t, t, anchors):
        """
        Compute the mean and variance of the diffusion posterior:
            q(x_{t-1} | x_t, x_0)
        """
        assert x_start.shape == x_t.shape == anchors.shape
        if not self.learn_anchor:
            anchors = anchors * 0.0
        
        posterior_mean = (
            extract_into_tensor(self.posterior_mean_coef1, t, x_start.shape) * x_start
            + extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t 
            + extract_into_tensor(self.posterior_mean_coef3, t, anchors.shape) * anchors
        )
        assert (
            posterior_mean.shape[0]
            == x_t.shape[0]
        )
        return posterior_mean
    
    def q_posterior_variance(self, x_start, x_t, t, anchors, variance):
        """
        Compute the mean and variance of the diffusion posterior:
            q(x_{t-1} | x_t, x_0)
        """
        assert x_start.shape == x_t.shape == anchors.shape == variance.shape

        B, C, N = x_start.shape

        posterior_variance = extract_into_tensor(self.posterior_variance, t, variance.shape) * variance
        posterior_log_variance_clipped = extract_into_tensor(
            self.posterior_log_variance_clipped, t, variance.shape
        ) + torch.log(variance)
        assert (
            posterior_variance.shape[0]
            == posterior_log_variance_clipped.shape[0]
            == x_t.shape[0]
        )
        return posterior_variance, posterior_log_variance_clipped
    
    def q_posterior_mean_variance(self, x_start, x_t, t, anchors, variance):
        """
        Compute the mean and variance of the diffusion posterior:
            q(x_{t-1} | x_t, x_0)
        """
        assert x_start.shape == x_t.shape == anchors.shape == variance.shape
        posterior_mean = self.q_posterior_mean(x_start, x_t, t, anchors)
        posterior_variance, posterior_log_variance_clipped = self.q_posterior_variance(x_start, x_t, t, anchors, variance=variance)
        # print("variance",  posterior_variance.permute([0, 3, 1,2]).reshape(-1, 3, 3))
        return posterior_mean, posterior_variance, posterior_log_variance_clipped
    
    
    def p_mean_variance(
        self, x, t, anchors, ctx=None, variance=None ,frozen_out=None, anchor_assignment=None, valid_id=None
    ):
        """
        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
        the initial x, x_0.
        :param model: the model, which takes a signal and a batch of timesteps
                      as input.
        :param x: the [N x C x ...] tensor at time t.
        :param t: a 1-D Tensor of timesteps.
        :param clip_denoised: if True, clip the denoised signal into [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample. Applies before
            clip_denoised.
        :return: a dict with the following keys:
                 - 'mean': the model mean output.
                 - 'variance': the model variance output.
                 - 'log_variance': the log of 'variance'.
                 - 'pred_xstart': the prediction for x_0.
        """
        B, C, N = x.shape
        assert t.shape == (B,)
        if self.res:
            input = x - anchors 
        else:
            input = x
        if self.include_anchors:
            input = torch.cat([input, anchors], dim=1)
        if self.include_cov:
            input = torch.cat([input, variance], dim=1)
        
        if frozen_out is not None:
            model_output_cond = frozen_out
        else:
            model_output_cond = self.model(input, self._scale_timesteps(t), ctx, anchors=anchors.transpose(1,2), anchor_assignment=anchor_assignment, variances=variance.transpose(1,2), valid_id=valid_id)

        if self.guidance:
            uncond_ctx = [torch.zeros_like(r) for r in ctx]
            model_output_uncond = self.model(input, self._scale_timesteps(t), uncond_ctx, anchors=anchors.transpose(1,2), anchor_assignment=anchor_assignment, variances=variance.transpose(1,2), valid_id=valid_id)
            model_output = (1. - self.classifier_weight) * model_output_uncond + self.classifier_weight * model_output_cond
        else:
            model_output = model_output_cond

        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
            assert model_output.shape == (B, C*2, N)
            model_output, model_var_values = torch.split(model_output, C, dim=1)
            if self.model_var_type == ModelVarType.LEARNED:
                model_log_variance = model_var_values
                model_variance = torch.exp(model_log_variance)
            else:
                min_log = extract_into_tensor(
                    self.posterior_log_variance_clipped, t, x.shape
                )
                max_log = extract_into_tensor(np.log(self.betas), t, x.shape)
                # The model_var_values is [-1, 1] for [min_var, max_var].
                frac = (model_var_values + 1) / 2
                model_log_variance = frac * max_log + (1 - frac) * min_log
                model_variance = torch.exp(model_log_variance)
        else:
            model_variance, model_log_variance = {
                # for fixedlarge, we set the initial (log-)variance like so
                # to get a better decoder log likelihood.
                ModelVarType.FIXED_LARGE: (
                    np.append(self.posterior_variance[1], self.betas[1:]),
                    np.log(np.append(self.posterior_variance[1], self.betas[1:])),
                ),
                ModelVarType.FIXED_SMALL: (
                    self.posterior_variance,
                    self.posterior_log_variance_clipped,
                ),
            }[self.model_var_type]

        if self.learn_variance:
            assert variance is not None 
            assert variance.shape == anchors.shape
        
        if not self.learn_variance:
            variance = from_numpy_to_device(self.k_squared, x.device).reshape(1, 3, 1).expand(B, -1, N)
        
        L = torch.sqrt(variance)
        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
            model_sqrt_log_variance = 0.5 * model_log_variance.reshape(B, C, N) + torch.log(L)
            model_variance = model_variance.reshape(B, C, N) * variance
            model_log_variance = model_log_variance.reshape(B, C, N) + torch.log(variance)
        else:
            model_sqrt_log_variance = 0.5 * extract_into_tensor(model_log_variance, t, [B, 3, N]) + torch.log(L)
            model_variance = extract_into_tensor(model_variance, t, [B, 3, N]) * variance
            model_log_variance = extract_into_tensor(model_log_variance, t, [B, 3, N]) + torch.log(variance)
        

        def process_xstart(x):
            if self.clip_xstart:
                return x.clip(-10, 10)
            return x

        if self.model_mean_type == ModelMeanType.PREVIOUS_X:
            pred_xstart = process_xstart(
                self._predict_xstart_from_xprev(x_t=x, t=t, anchors=anchors, xprev=model_output)
            )
            model_mean = model_output
        elif self.model_mean_type in [
            ModelMeanType.DRIFTED_EPSILON5, 
            ModelMeanType.SCALED_EPSILON, 
            ModelMeanType.START_X, 
            ModelMeanType.EPSILON, 
            ModelMeanType.EPSILON_AND_ANCHOR, 
            ModelMeanType.DRIFTED_EPSILON1, 
            ModelMeanType.DRIFTED_EPSILON2, 
            ModelMeanType.DRIFTED_EPSILON3, 
            ModelMeanType.DRIFTED_EPSILON4
            ]:
            if self.model_mean_type == ModelMeanType.START_X:
                pred_xstart = process_xstart(model_output)
            elif self.model_mean_type in [ModelMeanType.DRIFTED_EPSILON1, ModelMeanType.DRIFTED_EPSILON2, ModelMeanType.DRIFTED_EPSILON3, ModelMeanType.DRIFTED_EPSILON4]:
                if self.model_mean_type == ModelMeanType.DRIFTED_EPSILON1:
                    model_output = extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, model_output.shape) * model_output
                elif self.model_mean_type == ModelMeanType.DRIFTED_EPSILON3:
                    model_output = extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, model_output.shape) * model_output
                elif self.model_mean_type == ModelMeanType.DRIFTED_EPSILON4:
                    model_output = model_output - anchors
                elif self.model_mean_type == ModelMeanType.DRIFTED_EPSILON5:
                    model_output = L * extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, model_output.shape) * model_output
                
                pred_xstart = process_xstart(
                    self._predict_xstart_from_drifted_eps(x_t=x, t=t, anchors=anchors, eps=model_output, sqrt_variance=L)
                )
            elif self.model_mean_type == ModelMeanType.EPSILON_AND_ANCHOR:
                # print(model_output.shape)
                pred_output, pred_anchor = torch.split(model_output, 3, dim=1)
                # print(pred_output.shape, pred_anchor.shape)
                pred_xstart = process_xstart(
                    self._predict_xstart_from_eps(x_t=x, t=t, anchors=pred_anchor, eps=pred_output, sqrt_variance=L)
                )
            elif self.model_mean_type == ModelMeanType.SCALED_EPSILON:
                pred_xstart = process_xstart(
                    self._predict_xstart_from_eps(x_t=x, t=t, anchors=anchors, eps=model_output, sqrt_variance=None)
                )
            else:
                pred_xstart = process_xstart(
                    self._predict_xstart_from_eps(x_t=x, t=t, anchors=anchors, eps=model_output, sqrt_variance=L)
                )
            if self.ddim_sampling:
                # print(xt_dir_coeff.mean())
                # print(model_variance.mean())
                # print(extract_into_tensor(self.alphas_cumprod_prev, timesteps=t, broadcast_shape=model_output.shape).mean())
                # print()
                xt_dir = L * extract_into_tensor(self.xt_dir_coeff, t, model_output.shape) * model_output
                ddim_variance = model_sqrt_log_variance 
            else:
                xt_dir=None
                ddim_variance=None
            model_mean, _, _ = self.q_posterior_mean_variance(
                x_start=pred_xstart, x_t=x, t=t, anchors=anchors, variance=variance
            )
        else:
            raise NotImplementedError(self.model_mean_type)

        assert (
            model_mean.shape == pred_xstart.shape == x.shape
        )
        return {
            "mean": model_mean,
            "variance": model_variance,
            "log_variance": model_log_variance,
            "pred_xstart": pred_xstart,
            "sqrt_log_variance":model_sqrt_log_variance,
            "xt_dir": xt_dir,
            "ddim_variance":ddim_variance
        }
    
    def _predict_xstart_from_drifted_eps(self, x_t, t, anchors, eps, sqrt_variance):
        assert x_t.shape == eps.shape == anchors.shape
        
        return extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * (x_t - eps)
    def _predict_xstart_from_eps(self, x_t, t, anchors, eps, sqrt_variance):
        assert x_t.shape == eps.shape == anchors.shape
        if sqrt_variance is None:
            sqrt_variance = torch.ones_like(x_t)
        if not self.learn_anchor:
            anchors = anchors * 0.0
        return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * (x_t - anchors) + anchors
                -  extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, sqrt_variance.shape) * sqrt_variance * eps
                )

    def _predict_xstart_from_xprev(self, x_t, t, anchors, xprev):
        assert x_t.shape == xprev.shape == anchors.shape
        if not self.learn_anchor:
            anchors = anchors * 0.0
        return (  # (xprev - coef2*x_t - coef3 * anchors) / coef1
            extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev
            - extract_into_tensor(
                self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape
            ) * x_t
            - extract_into_tensor(
                self.posterior_mean_coef3 / self.posterior_mean_coef1, t, x_t.shape
            ) * anchors
        )
    
    def _predict_eps_from_xstart(self, x_t, t, anchors, pred_xstart, variance=None):
    # Not WOrking !!!!! Don't Use!!!!
        B, C, N = x_t.shape

        if self.learn_variance:
            assert variance is not None 
            assert variance.shape == anchors.shape
        if not self.learn_anchor:
            anchors = anchors * 0.0
            
        if not self.learn_variance:
            variance = from_numpy_to_device(self.k_squared, x_t.device).reshape(1, 3, 1).expand(B, -1, N)
        L = torch.sqrt(variance)
        return (
            extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * (x_t - anchors)
            - pred_xstart
        ) / extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, L.shape)
        
    def _scale_timesteps(self, t):
        if self.rescale_timesteps:
            return t.float() * (1000.0 / self.num_timesteps)
        if self.use_beta:
            return extract_into_tensor(self.betas, t, t.shape)
        return t

    def p_sample(
        self, x, t, anchors, ctx=None, variance=None, anchor_assignment=None, valid_id=None
    ):
        """
        Sample x_{t-1} from the model at the given timestep.
        :param model: the model to sample from.
        :param x: the current tensor at x_{t-1}.
        :param t: the value of t, starting at 0 for the first diffusion step.
        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample.
        :return: a dict containing the following keys:
                 - 'sample': a random sample from the model.
                 - 'pred_xstart': a prediction of x_0.
        """
        out = self.p_mean_variance(
            x,
            t,
            anchors,
            ctx=ctx,
            variance=variance,
            anchor_assignment=anchor_assignment,
            valid_id=valid_id
        )
        if not self.learn_anchor:
            anchors = anchors * 0.0
        noise = torch.randn_like(x)
        nonzero_mask = (
            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
        )  # no noise when t == 0
        if self.ddim_sampling:
            sample = (out["pred_xstart"] - anchors) * torch.sqrt(extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)) + anchors + out['xt_dir'] + self.ddim_eta * nonzero_mask * torch.sqrt(out["variance"]) * noise
        else:
            sample = out["mean"] + nonzero_mask * torch.sqrt(out["variance"]) * noise
        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
    
    def p_sample_loop(
        self,
        shape,
        anchors,
        ctx=None,
        noise=None,
        variance=None,
        anchor_assignment=None,
        valid_id=None,
        device=None,
        progress=False
    ):
        """
        Generate samples from the model.
        :param model: the model module.
        :param shape: the shape of the samples, (N, C, H, W).
        :param noise: if specified, the noise from the encoder to sample.
                      Should be of the same shape as `shape`.
        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample.
        :param device: if specified, the device to create the samples on.
                       If not specified, use a model parameter's device.
        :param progress: if True, show a tqdm progress bar.
        :return: a non-differentiable batch of samples.
        """
        final = None
        for t, sample in self.p_sample_loop_progressive(
            shape,
            noise=noise,
            anchors=anchors,
            ctx=ctx,
            variance=variance,
            anchor_assignment=anchor_assignment,
            valid_id=valid_id,
            device=device,
            progress=progress,
        ):
            final = sample

        return final["sample"]

    def p_sample_loop_progressive(
        self,
        shape,
        anchors,
        ctx=None,
        variance=None,
        anchor_assignment=None,
        valid_id=None,
        noise=None,
        device=None,
        progress=False,
    ):
        """
        Generate samples from the model and yield intermediate samples from
        each timestep of diffusion.
        Arguments are the same as p_sample_loop().
        Returns a generator over dicts, where each dict is the return value of
        p_sample().
        """
        if device is None:
            device = next(self.model.parameters()).device
        assert isinstance(shape, (tuple, list))
        if noise is not None:
            pcd = noise
        else:
            B, _, N = shape

            if self.learn_variance:
                assert variance is not None 
                assert variance.shape == anchors.shape

            if not self.learn_variance:
                cov = from_numpy_to_device(self.k_squared, device).reshape(1, 3, 1).expand(B, -1, N)
            else:
                cov = variance
            L = torch.sqrt(cov)
            pcd = L  * torch.randn(*shape, device=device) + anchors * int(self.learn_anchor)

        indices = self.steps[::-1]
        print("sampling trajectory: %s" % [s for s in indices])

        if progress:
            # Lazy import so that we don't depend on tqdm.
            from tqdm.auto import tqdm

            indices = tqdm(indices)
        yield self.num_timesteps, dict(sample=pcd)
        for i in indices:
            t = torch.tensor([i] * shape[0], device=device)
            with torch.no_grad():
                out = self.p_sample(
                    pcd,
                    t,
                    anchors,
                    ctx=ctx,
                    variance=variance,
                    anchor_assignment=anchor_assignment,
                    valid_id=valid_id
                )
                yield i, out
                pcd = out["sample"]
                
    def q_sample_loop(
        self,
        gt_pcd,
        anchors,
        variance=None,
        noise=None,
        device=None,
        progress=False,
    ):
        """
        forward sample from the model
        :param shape: the shape of the samples, (B, C, N).
        :param noise: if specified, the noise from the encoder to sample.
                      Should be of the same shape as `shape`.
        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample.
        :param device: if specified, the device to create the samples on.
                       If not specified, use a model parameter's device.
        :param progress: if True, show a tqdm progress bar.
        :return: a non-differentiable batch of samples.
        """
        final = None
        for t, sample in self.q_sample_loop_progressive(
            gt_pcd,
            noise=noise,
            anchors=anchors,
            variance=variance,
            device=device,
            progress=progress,
        ):
            final = sample

        return final

    def q_sample_loop_progressive(self, gt_pcd, anchors, noise=None, variance=None, device=None, progress=False):
        if device is None:
            device = next(self.model.parameters()).device
        indices = list(range(self.num_timesteps))[1:]
        if noise is None:
            noise = torch.randn(*gt_pcd.shape).to(device)
        if progress:
            # Lazy import so that we don't depend on tqdm.
            from tqdm.auto import tqdm

            indices = tqdm(indices)
        for i in indices:
            t = torch.tensor([i] * gt_pcd.shape[0], device=device)
            with torch.no_grad():
                yield i, self.q_sample(gt_pcd, t, anchors, noise=noise, variance=variance)

    def _vb_terms_bpd(
        self, x_start, x_t, t, anchors, ctx=None, variance=None, frozen_out=None, anchor_assignment=None, flags=None
    ):
        """
        Get a term for the variational lower-bound.
        This allows for comparison to other papers.
        :return: a dict with the following keys:
                 - 'output': a shape [N] tensor of NLLs or KLs.
                 - 'pred_xstart': the x_0 predictions.
        """
        B, C, N = x_start.shape
        
        true_mean, true_posterior_variance, _ = self.q_posterior_mean_variance(
            x_start=x_start, x_t=x_t, t=t, anchors=anchors, variance=variance
        )
        true_posterior_diag_var = torch.diagonal(true_posterior_variance.permute([0, 3, 1, 2]).reshape(B*N, C, C), offset=0, dim1=-2, dim2=-1)
        log_true_posterior_diag_var = torch.log(true_posterior_diag_var)
        out = self.p_mean_variance(
            x_t, t, anchors, ctx=ctx, variance=variance, frozen_out=frozen_out, anchor_assignment=anchor_assignment
        )
        # print(true_mean.shape, out["mean"].shape, true_log_variance_clipped.shape, out["log_variance"].shape)
        # print(true_posterior_variance)
        kl = normal_kl(
            true_mean.transpose(1,2).reshape(B*N, C), log_true_posterior_diag_var, out["mean"].transpose(1,2).reshape(B*N, C), out["log_variance"].transpose(1,2).reshape(B*N, C), dim=C
        ).sum(-1)
        # print(kl.reshape(B, N).mean(-1))
        # print()
        if flags is not None:
            kl = (kl.reshape(B, 1, N) * flags).sum(dim=(1,2)) / (flags.sum(dim=(1,2)) * np.log(2.0))
        else:
            kl = kl.reshape(B, N).mean(-1) / np.log(2.0)
        # print(torch.all(torch.det(out["variance"].permute([0, 3, 1, 2]).reshape(B*N, C, C)) > 0))
        decoder_nll = -gaussian_log_likelihood(z=x_start.transpose(1,2).reshape(B*N, C), mean=out["mean"].transpose(1,2).reshape(B*N, C), logvar=out["log_variance"].transpose(1,2).reshape(B*N, C), dim=C).sum(-1)
        # print(decoder_nll.reshape(B, N).mean(-1))
        # print(t)
        if flags is not None:
            decoder_nll = (decoder_nll.reshape(B, 1, N) * flags).sum(dim=(1,2)) / (flags.sum(dim=(1,2)) * np.log(2.0))
        else:
            decoder_nll = decoder_nll.reshape(B, N).mean(-1) / np.log(2.0)
        
            
        # if kl > 1000:
            # print(decoder_nll, kl, out["mean"], out["variance"].permute([0, 3, 1, 2]).reshape(B*N, C, C), t)
            # exit()
        # At the first timestep return the decoder NLL,
        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
        output = torch.where((t == 0), decoder_nll, kl)
        return {"output": output, "pred_xstart": out["pred_xstart"]}
    """
    def training_losses(self, x_start, t, anchors, x_start_feats, code=None, partial_pc=None, noise=None):
        Compute training losses for a single timestep.
        :param model: the model to evaluate loss on.
        :param x_start: the [B x 3 x N] tensor of inputs.
        :param t: a batch of timestep indices.
        :param noise: if specified, the specific Gaussian noise to try to remove.
        :return: a dict with the key "loss" containing a tensor of shape [N].
                 Some mean or variance settings may also have other keys.
        if noise is None:
            noise = torch.randn_like(x_start)
        x_t = self.q_sample(x_start, t, anchors, noise=noise)

        terms = {}

        if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:
            terms["loss"] = self._vb_terms_bpd(
                x_start=x_start,
                x_t=x_t,
                t=t,
                anchors=anchors,
                x_start_feats=x_start_feats,
                code=code,
                clip_denoised=False,
            )["output"]
            if self.loss_type == LossType.RESCALED_KL:
                terms["loss"] *= self.num_timesteps
        elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:
            model_output = self.model(x_t, self._scale_timesteps(t), x_start_feats, anchors, code, partial_pc)

            if self.model_var_type in [
                ModelVarType.LEARNED,
                ModelVarType.LEARNED_RANGE,
            ]:
                B, C = x_t.shape[:2]
                assert model_output.shape == (B, C * 2, *x_t.shape[2:])
                model_output, model_var_values = torch.split(model_output, C, dim=1)
                # Learn the variance using the variational bound, but don't let
                # it affect our mean prediction.
                frozen_out = torch.cat([model_output.detach(), model_var_values], dim=1)
                terms["vb"] = self._vb_terms_bpd(
                    x_start=x_start,
                    x_t=x_t,
                    t=t,
                    anchors=anchors, 
                    x_start_feats=x_start_feats,
                    code=code,
                    model_output=frozen_out
                )["output"]
                if self.loss_type == LossType.RESCALED_MSE:
                    # Divide by 1000 for equivalence with initial implementation.
                    # Without a factor of 1/1000, the VB term hurts the MSE term.
                    terms["vb"] *= self.num_timesteps / 1000.0

            target = {
                ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(
                    x_start=x_start, x_t=x_t, t=t, anchors=anchors
                )[0],
                ModelMeanType.START_X: x_start,
                ModelMeanType.EPSILON: noise,
            }[self.model_mean_type]
            assert model_output.shape == target.shape == x_start.shape
            terms["mse"] = mean_flat((target - model_output) ** 2)
            if "vb" in terms:
                terms["recons_loss"] = terms["mse"] + terms["vb"]
            else:
                terms["recons_loss"] = terms["mse"]
        else:
            raise NotImplementedError(self.loss_type)
        return terms["recons_loss"].mean()
    """
    def training_losses(self, x_start, t, anchors=None, variance=None, ctx=None, reduce=True, anchor_assignment=None, valid_id=None, flags=None, noise=None):
        """
        Compute training losses for a single timestep.
        :param model: the model to evaluate loss on.
        :param x_start: the [B x 3 x N] tensor of inputs.
        :param t: a batch of timestep indices.
        :param noise: if specified, the specific Gaussian noise to try to remove.
        :return: a dict with the key "loss" containing a tensor of shape [N].
                 Some mean or variance settings may also have other keys.
        """
        loss_dict = dict()
        if anchors is None:
            anchors = torch.zeros_like(x_start)
        if self.learn_variance:
            assert variance is not None 
            assert variance.shape == anchors.shape
            # print(variance.mean())

        if noise is None:
            noise = torch.randn_like(x_start)
        B, C, N = x_start.shape
        ori_anchors = anchors
        if not self.learn_anchor:
            anchors = anchors * 0.0
        x_t = self.q_sample(x_start, t, anchors, noise=noise, variance=variance)
        model_in=x_t
        if self.res:
            model_in = model_in - ori_anchors
        if self.include_anchors:
            model_in = torch.cat([model_in, ori_anchors], dim=1)
        if self.include_cov:
            model_in = torch.cat([model_in, variance], dim=1)

        
        model_output = self.model(model_in, self._scale_timesteps(t), ctx, anchor_assignment=anchor_assignment, anchors=ori_anchors.transpose(1,2), variances=variance.transpose(1,2) if variance is not None else None, valid_id=valid_id)
        if self.model_var_type in [
            ModelVarType.LEARNED,
            ModelVarType.LEARNED_RANGE,
        ]:
            # print(model_output.shape)
            assert model_output.shape == (B, C*2, N)
            model_output, model_var_values = torch.split(model_output, C, dim=1)
            loss_dict['model_var_value'] = model_var_values.mean()
            # Learn the variance using the variational bound, but don't let
            # it affect our mean prediction.
            frozen_out = torch.cat([model_output.detach(), model_var_values], dim=1)
            vb_loss= self._vb_terms_bpd(
                x_start=x_start,
                x_t=x_t,
                t=t,
                anchors=ori_anchors,
                ctx=ctx,
                variance=variance,
                frozen_out=frozen_out,
                anchor_assignment=anchor_assignment,
                flags=flags
            )["output"].mean()
            # print(vb_loss)
            if vb_loss > 1000:
                print("vb_loss is", vb_loss.mean())
                print("model var is ", model_var_values.mean())
                exit()
            vb_loss *= self.num_timesteps / 1000.0
            loss_dict['vb_loss'] = vb_loss
            
        if not self.learn_variance:
            variance = from_numpy_to_device(self.k_squared, x_start.device).reshape(1, 3, 1).expand(B, -1, N)
        L = torch.sqrt(variance)
        target = {
                ModelMeanType.START_X: x_start,
                ModelMeanType.DRIFTED_EPSILON1: extract_into_tensor(1. - self.sqrt_alphas_cumprod, t, noise.shape) * anchors + extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, noise.shape) * noise,
                ModelMeanType.DRIFTED_EPSILON2: extract_into_tensor(self.sqrt_recip_alphas_cumprod - 1, t, noise.shape) * anchors + extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, noise.shape) * noise,
                ModelMeanType.DRIFTED_EPSILON3: extract_into_tensor((1. - self.sqrt_alphas_cumprod) / self.sqrt_one_minus_alphas_cumprod, t, noise.shape) * anchors + L * noise,
                ModelMeanType.DRIFTED_EPSILON4: extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, noise.shape) * anchors + extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, noise.shape) * noise,
                ModelMeanType.DRIFTED_EPSILON5: extract_into_tensor((1. - self.sqrt_alphas_cumprod) / self.sqrt_one_minus_alphas_cumprod, t, noise.shape) * anchors / L + noise,
                ModelMeanType.SCALED_EPSILON: L * noise,
                ModelMeanType.EPSILON: noise,
                ModelMeanType.EPSILON_AND_ANCHOR: torch.cat([noise, anchors], dim=1)
            }[self.model_mean_type]
        assert model_output.shape == target.shape
        diffusion_loss = (target - model_output) ** 2
        if flags is not None:
            diffusion_loss = diffusion_loss * flags
        if reduce:
            if flags is not None:
                diffusion_loss = diffusion_loss.mean(1).sum() / flags.sum()
            else:
                diffusion_loss = diffusion_loss.mean()
            
        if self.scale_loss:
            diffusion_loss = diffusion_loss * extract_into_tensor( self.betas * self.betas / (2 * (1. - self.betas) * (1. - self.alphas_cumprod)), t, diffusion_loss.shape)
        loss_dict['mse_loss'] = diffusion_loss
        
        return loss_dict
python/difffacto/models/diffusions/anchored_diffusion.py
diffFacto/diffFacto
import numpy as np
import torch


def normal_kl(mean1, logvar1, mean2, logvar2, dim=3):
    """
    Compute the KL divergence between two gaussians.
    Shapes are automatically broadcasted, so batches can be compared to
    scalars, among other use cases.
    """


    return 0.5 * (
        -dim
        + logvar2
        - logvar1
        + torch.exp(logvar1 - logvar2)
        + ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
    )

def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
    """
    credit: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
    
    Create a beta schedule that discretizes the given alpha_t_bar function,
    which defines the cumulative product of (1-beta) over time from t = [0,1].
    :param num_diffusion_timesteps: the number of betas to produce.
    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
                    produces the cumulative product of (1-beta) up to that
                    part of the diffusion process.
    :param max_beta: the maximum beta to use; use values lower than 1 to
                    prevent singularities.
    """
    betas = []
    for i in range(num_diffusion_timesteps):
        t1 = i / num_diffusion_timesteps
        t2 = (i + 1) / num_diffusion_timesteps
        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
    return np.array(betas, dtype=np.float64)


def extract_into_tensor(arr, timesteps, broadcast_shape):
    """
    Extract values from a 1-D numpy array for a batch of indices.
    :param arr: the 1-D numpy array.
    :param timesteps: a tensor of indices into the array to extract.
    :param broadcast_shape: a larger shape of K dimensions with the batch
                            dimension equal to the length of timesteps.
    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
    """
    res = from_numpy_to_device(arr, timesteps.device)[timesteps]
    while len(res.shape) < len(broadcast_shape):
        res = res[..., None]
    return res.expand(broadcast_shape)

def from_numpy_to_device(arr, device):
    """
    Extract values from a 1-D numpy array for a batch of indices.
    :param arr: the 1-D numpy array.
    :param timesteps: a tensor of indices into the array to extract.
    :param broadcast_shape: a larger shape of K dimensions with the batch
                            dimension equal to the length of timesteps.
    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
    """
    res = torch.from_numpy(arr).to(device=device).float()
    return res

def get_cov_from_params(eigens, thetas=None):
    # thetas: [B, 3], eigens: [B, 3]
    B = eigens.shape[0]
    if thetas is None:
        thetas = torch.zeros_like(eigens)
    assert B == thetas.shape[0]
    tensor_0 = torch.zeros(B).to(eigens.device)
    tensor_1 = torch.ones(B).to(eigens.device)
    cos_thetas = torch.cos(thetas)
    sin_thetas = torch.sin(thetas)

    RX = torch.stack([
                    torch.stack([tensor_1, tensor_0, tensor_0], dim=1),
                    torch.stack([tensor_0, cos_thetas[:, 0], -sin_thetas[:, 0]], dim=1),
                    torch.stack([tensor_0, sin_thetas[:, 0], cos_thetas[:, 0]], dim=1)], dim=1).reshape(B,3,3)

    RY = torch.stack([
                    torch.stack([cos_thetas[:, 1], tensor_0, sin_thetas[:, 1]], dim=1),
                    torch.stack([tensor_0, tensor_1, tensor_0], dim=1),
                    torch.stack([-sin_thetas[:, 1], tensor_0, cos_thetas[:, 1]], dim=1)], dim=1).reshape(B, 3,3)

    RZ = torch.stack([
                    torch.stack([cos_thetas[:,2], -sin_thetas[:,2], tensor_0], dim=1),
                    torch.stack([sin_thetas[:,2], cos_thetas[:,2], tensor_0], dim=1),
                    torch.stack([tensor_0, tensor_0, tensor_1], dim=1)], dim=1).reshape(B,3,3)

    R = torch.bmm(RZ, RY)
    R = torch.bmm(R, RX)
    cov = torch.bmm(R, eigens.reshape(B, 3,1) * R.transpose(1,2))
    sqrt_cov = R * torch.sqrt(eigens.reshape(B,1,3))
    return cov, sqrt_cov
python/difffacto/models/diffusions/diffusion_utils.py
diffFacto/diffFacto
import torch 
from torch.nn import Module
from anchor_diff.utils.registry import DIFFUSIONS, NETS, build_from_cfg
from anchor_diff.utils.constants import ModelMeanType, ModelVarType, LossType
from anchor_diff.utils.misc import *
from anchor_diff.utils.constants import *
import math
import numpy as np
from .diffusion_utils import extract_into_tensor, betas_for_alpha_bar, from_numpy_to_device, get_cov_from_params, normal_kl

@DIFFUSIONS.register_module()
class LDM(Module):

    def __init__(
        self, 
        net, 
        num_timesteps, 
        beta_1, 
        beta_T, 
        mode='linear', 
        use_beta=True, 
        rescale_timesteps=False, 
        loss_type='mse', 
        model_mean_type='epsilon',
        model_var_type='fixed_small', 
        clip_xstart=False,
        classifier_weight=1.,
        guidance=False,

        ddim_sampling=False,
        ddim_nsteps=10,
        ddim_discretize='uniform',
        ddim_eta=1.
    ):
        super().__init__()
        assert mode in ('linear', 'cosine')
        self.clip_xstart=clip_xstart
        self.model = build_from_cfg(net, NETS)
        self.model_mean_type = model_mean_dict[model_mean_type]
        self.model_var_type = model_var_dict[model_var_type]
        self.loss_type = model_loss_dict[loss_type]
        self.ddim_sampling=ddim_sampling
        self.beta_1 = beta_1
        self.beta_T = beta_T
        self.use_beta = use_beta
            
        self.rescale_timesteps = rescale_timesteps
        self.num_timesteps = int(num_timesteps)
        if mode == 'linear':
            betas = np.linspace(self.beta_1, self.beta_T, num=self.num_timesteps, dtype=np.float64)
        elif mode == 'cosine':
            betas = betas_for_alpha_bar(self.num_timesteps, lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,)
            
        self.betas = betas
        assert len(betas.shape) == 1, "betas must be 1-D"
        assert (betas > 0).all() and (betas <= 1).all()
        self.classifier_weight=classifier_weight
        self.guidance=guidance
        

        alphas = 1.0 - betas
        self.alphas_cumprod = np.cumprod(alphas, axis=0)
        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)

        # calculations for diffusion q(x_t | x_{t-1}) and others
        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)

        # calculations for posterior q(x_{t-1} | x_t, x_0)
        self.posterior_variance = (
            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )
        # log calculation clipped because the posterior variance is 0 at the
        # beginning of the diffusion chain.
        self.posterior_log_variance_clipped = np.log(
            np.append(self.posterior_variance[1], self.posterior_variance[1:])
        )
        self.posterior_mean_coef1 = (
            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )
        self.posterior_mean_coef2 = (
            (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod)
        )
        
        if self.ddim_sampling:
            self.ddim_eta = ddim_eta
            self.xt_dir_coeff = np.sqrt(1. - self.alphas_cumprod - ddim_eta * ddim_eta * self.posterior_variance)
            if ddim_discretize == 'uniform':
                skip = self.num_timesteps // ddim_nsteps
                self.steps = list(range(0, self.num_timesteps, skip))
            elif ddim_discretize == 'quad':
                self.steps = (np.linspace(0., math.sqrt(self.num_timesteps * 0.8), ddim_nsteps) ** 2).astype(np.int32)
                self.steps = self.steps.tolist()
            else:
                NotImplementedError(ddim_discretize)
        else:
            self.steps = list(range(self.num_timesteps))

        
        
    
    def q_sample(self, x_start, t, noise=None):
        """
        Diffuse the data for a given number of diffusion steps.
        In other words, sample from q(x_t | x_0).
        :param x_0: the initial data batch.
        :param anchors: anchors for each point in x_0
        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
        :param noise: if specified, the split-out normal noise.
        :return: A noisy version of x_0.
        """
        if noise is None:
            noise = torch.randn_like(x_start)
        assert noise.shape == x_start.shape
        return (
            extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
            + extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
            * noise)
        
    def q_posterior_mean(self, x_start, x_t, t):
        """
        Compute the mean and variance of the diffusion posterior:
            q(x_{t-1} | x_t, x_0)
        """
        assert x_start.shape == x_t.shape
        posterior_mean = (
            extract_into_tensor(self.posterior_mean_coef1, t, x_start.shape) * x_start
            + extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t 
        )
        assert (
            posterior_mean.shape[0]
            == x_t.shape[0]
        )
        return posterior_mean
    
    def q_posterior_variance(self, x_start, x_t, t):
        """
        Compute the mean and variance of the diffusion posterior:
            q(x_{t-1} | x_t, x_0)
        """
        assert x_start.shape == x_t.shape


        posterior_variance = extract_into_tensor(self.posterior_variance, t, x_start.shape) 
        posterior_log_variance_clipped = extract_into_tensor(
            self.posterior_log_variance_clipped, t, x_start.shape
        )
        assert (
            posterior_variance.shape[0]
            == posterior_log_variance_clipped.shape[0]
            == x_t.shape[0]
        )
        return posterior_variance, posterior_log_variance_clipped
    
    def q_posterior_mean_variance(self, x_start, x_t, t):
        """
        Compute the mean and variance of the diffusion posterior:
            q(x_{t-1} | x_t, x_0)
        """
        assert x_start.shape == x_t.shape
        posterior_mean = self.q_posterior_mean(x_start, x_t, t)
        posterior_variance, posterior_log_variance_clipped = self.q_posterior_variance(x_start, x_t, t)
        return posterior_mean, posterior_variance, posterior_log_variance_clipped
    
    
    def p_mean_variance(
        self, x, t, ctx=None, frozen_out=None, valid_id=None
    ):
        """
        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
        the initial x, x_0.
        :param model: the model, which takes a signal and a batch of timesteps
                      as input.
        :param x: the [N x C x ...] tensor at time t.
        :param t: a 1-D Tensor of timesteps.
        :param clip_denoised: if True, clip the denoised signal into [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample. Applies before
            clip_denoised.
        :return: a dict with the following keys:
                 - 'mean': the model mean output.
                 - 'variance': the model variance output.
                 - 'log_variance': the log of 'variance'.
                 - 'pred_xstart': the prediction for x_0.
        """
        B, C, N = x.shape
        assert t.shape == (B,)
        input = x
        if frozen_out is not None:
            model_output_cond = frozen_out
        else:
            model_output_cond = self.model(input, self._scale_timesteps(t), ctx, valid_id=valid_id)

        if self.guidance:
            uncond_ctx = [torch.zeros_like(r) for r in ctx]
            model_output_uncond = self.model(input, self._scale_timesteps(t), uncond_ctx, valid_id=valid_id)
            model_output = (1. - self.classifier_weight) * model_output_uncond + self.classifier_weight * model_output_cond
        else:
            model_output = model_output_cond

        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
            assert model_output.shape == (B, C*2, N)
            model_output, model_var_values = torch.split(model_output, C, dim=1)
            if self.model_var_type == ModelVarType.LEARNED:
                model_log_variance = model_var_values
                model_variance = torch.exp(model_log_variance)
            else:
                min_log = extract_into_tensor(
                    self.posterior_log_variance_clipped, t, x.shape
                )
                max_log = extract_into_tensor(np.log(self.betas), t, x.shape)
                # The model_var_values is [-1, 1] for [min_var, max_var].
                frac = (model_var_values + 1) / 2
                model_log_variance = frac * max_log + (1 - frac) * min_log
                model_variance = torch.exp(model_log_variance)
        else:
            model_variance, model_log_variance = {
                # for fixedlarge, we set the initial (log-)variance like so
                # to get a better decoder log likelihood.
                ModelVarType.FIXED_LARGE: (
                    np.append(self.posterior_variance[1], self.betas[1:]),
                    np.log(np.append(self.posterior_variance[1], self.betas[1:])),
                ),
                ModelVarType.FIXED_SMALL: (
                    self.posterior_variance,
                    self.posterior_log_variance_clipped,
                ),
            }[self.model_var_type]

        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
            model_sqrt_log_variance = 0.5 * model_log_variance.reshape(B, C, N)
            model_variance = model_variance.reshape(B, C, N)
            model_log_variance = model_log_variance.reshape(B, C, N)
        else:
            model_sqrt_log_variance = 0.5 * extract_into_tensor(model_log_variance, t, [B, C, N])
            model_variance = extract_into_tensor(model_variance, t, [B, C, N])
            model_log_variance = extract_into_tensor(model_log_variance, t, [B, C, N])
        

        def process_xstart(x):
            if self.clip_xstart:
                return x.clip(-10, 10)
            return x

        if self.model_mean_type == ModelMeanType.PREVIOUS_X:
            pred_xstart = process_xstart(
                self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output)
            )
            model_mean = model_output
        elif self.model_mean_type in [
            ModelMeanType.START_X, 
            ModelMeanType.EPSILON, 
            ]:
            if self.model_mean_type == ModelMeanType.START_X:
                pred_xstart = process_xstart(model_output)
            else:
                pred_xstart = process_xstart(
                    self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
                )
            if self.ddim_sampling:
                # print(xt_dir_coeff.mean())
                # print(model_variance.mean())
                # print(extract_into_tensor(self.alphas_cumprod_prev, timesteps=t, broadcast_shape=model_output.shape).mean())
                # print()
                xt_dir = extract_into_tensor(self.xt_dir_coeff, t, model_output.shape) * model_output
                ddim_variance = model_sqrt_log_variance 
            else:
                xt_dir=None
                ddim_variance=None
            model_mean, _, _ = self.q_posterior_mean_variance(
                x_start=pred_xstart, x_t=x, t=t
            )
        else:
            raise NotImplementedError(self.model_mean_type)

        assert (
            model_mean.shape == pred_xstart.shape == x.shape
        )
        return {
            "mean": model_mean,
            "variance": model_variance,
            "log_variance": model_log_variance,
            "pred_xstart": pred_xstart,
            "sqrt_log_variance":model_sqrt_log_variance,
            "xt_dir": xt_dir,
            "ddim_variance":ddim_variance
        }
    
    def _predict_xstart_from_drifted_eps(self, x_t, t, eps):
        assert x_t.shape == eps.shape
        
        return extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * (x_t - eps)
    def _predict_xstart_from_eps(self, x_t, t, eps):
        assert x_t.shape == eps.shape
        return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
                -  extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
                )

    def _predict_xstart_from_xprev(self, x_t, t, xprev):
        assert x_t.shape == xprev.shape
        return (  # (xprev - coef2*x_t - coef3 * anchors) / coef1
            extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev
            - extract_into_tensor(
                self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape
            ) * x_t
        )
    
    def _predict_eps_from_xstart(self, x_t, t, pred_xstart,):
    # Not WOrking !!!!! Don't Use!!!!
        B, C, N = x_t.shape

        return (
            extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
            - pred_xstart
        ) / extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
        
    def _scale_timesteps(self, t):
        if self.rescale_timesteps:
            return t.float() * (1000.0 / self.num_timesteps)
        if self.use_beta:
            return extract_into_tensor(self.betas, t, t.shape)
        return t

    def p_sample(
        self, x, t, ctx=None, valid_id=None
    ):
        """
        Sample x_{t-1} from the model at the given timestep.
        :param model: the model to sample from.
        :param x: the current tensor at x_{t-1}.
        :param t: the value of t, starting at 0 for the first diffusion step.
        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample.
        :return: a dict containing the following keys:
                 - 'sample': a random sample from the model.
                 - 'pred_xstart': a prediction of x_0.
        """
        out = self.p_mean_variance(
            x,
            t,
            ctx=ctx,
            valid_id=valid_id
        )
        noise = torch.randn_like(x)
        nonzero_mask = (
            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
        )  # no noise when t == 0
        if self.ddim_sampling:
            sample = out["pred_xstart"] * torch.sqrt(extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)) + out['xt_dir'] + self.ddim_eta * nonzero_mask * torch.sqrt(out["variance"]) * noise
        else:
            sample = out["mean"] + nonzero_mask * torch.sqrt(out["variance"]) * noise
        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
    
    def p_sample_loop(
        self,
        shape,
        ctx=None,
        noise=None,
        valid_id=None,
        device=None,
        progress=False
    ):
        """
        Generate samples from the model.
        :param model: the model module.
        :param shape: the shape of the samples, (N, C, H, W).
        :param noise: if specified, the noise from the encoder to sample.
                      Should be of the same shape as `shape`.
        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample.
        :param device: if specified, the device to create the samples on.
                       If not specified, use a model parameter's device.
        :param progress: if True, show a tqdm progress bar.
        :return: a non-differentiable batch of samples.
        """
        final = None
        for t, sample in self.p_sample_loop_progressive(
            shape,
            noise=noise,
            ctx=ctx,
            valid_id=valid_id,
            device=device,
            progress=progress,
        ):
            final = sample

        return final["sample"]

    def p_sample_loop_progressive(
        self,
        shape,
        ctx=None,
        valid_id=None,
        noise=None,
        device=None,
        progress=False,
    ):
        """
        Generate samples from the model and yield intermediate samples from
        each timestep of diffusion.
        Arguments are the same as p_sample_loop().
        Returns a generator over dicts, where each dict is the return value of
        p_sample().
        """
        if device is None:
            device = next(self.model.parameters()).device
        assert isinstance(shape, (tuple, list))
        if noise is not None:
            pcd = noise
        else:
            B, _, N = shape

            pcd = torch.randn(*shape, device=device)

        indices = self.steps[::-1]
        print("sampling trajectory: %s" % [s for s in indices])

        if progress:
            # Lazy import so that we don't depend on tqdm.
            from tqdm.auto import tqdm

            indices = tqdm(indices)
        yield self.num_timesteps, dict(sample=pcd)
        for i in indices:
            t = torch.tensor([i] * shape[0], device=device)
            with torch.no_grad():
                out = self.p_sample(
                    pcd,
                    t,
                    ctx=ctx,
                    valid_id=valid_id
                )
                yield i, out
                pcd = out["sample"]
                
    
    def _vb_terms_bpd(
        self, x_start, x_t, t,ctx=None, frozen_out=None, flags=None
    ):
        """
        Get a term for the variational lower-bound.
        This allows for comparison to other papers.
        :return: a dict with the following keys:
                 - 'output': a shape [N] tensor of NLLs or KLs.
                 - 'pred_xstart': the x_0 predictions.
        """
        B, C, N = x_start.shape
        
        true_mean, true_posterior_diag_var, _ = self.q_posterior_mean_variance(
            x_start=x_start, x_t=x_t, t=t, 
        )
        log_true_posterior_diag_var = torch.log(true_posterior_diag_var)
        out = self.p_mean_variance(
            x_t, t, ctx=ctx, frozen_out=frozen_out, 
        )
        # print(true_mean.shape, out["mean"].shape, true_log_variance_clipped.shape, out["log_variance"].shape)
        # print(true_posterior_variance)
        kl = normal_kl(
            true_mean.transpose(1,2).reshape(B*N, C), log_true_posterior_diag_var, out["mean"].transpose(1,2).reshape(B*N, C), out["log_variance"].transpose(1,2).reshape(B*N, C), dim=C
        ).sum(-1)
        # print(kl.reshape(B, N).mean(-1))
        # print()
        if flags is not None:
            kl = (kl.reshape(B, 1, N) * flags).sum(dim=(1,2)) / (flags.sum(dim=(1,2)) * np.log(2.0))
        else:
            kl = kl.reshape(B, N).mean(-1) / np.log(2.0)
        # print(torch.all(torch.det(out["variance"].permute([0, 3, 1, 2]).reshape(B*N, C, C)) > 0))
        decoder_nll = -gaussian_log_likelihood(z=x_start.transpose(1,2).reshape(B*N, C), mean=out["mean"].transpose(1,2).reshape(B*N, C), logvar=out["log_variance"].transpose(1,2).reshape(B*N, C), dim=C).sum(-1)
        # print(decoder_nll.reshape(B, N).mean(-1))
        # print(t)
        if flags is not None:
            decoder_nll = (decoder_nll.reshape(B, 1, N) * flags).sum(dim=(1,2)) / (flags.sum(dim=(1,2)) * np.log(2.0))
        else:
            decoder_nll = decoder_nll.reshape(B, N).mean(-1) / np.log(2.0)
        
            
        # if kl > 1000:
            # print(decoder_nll, kl, out["mean"], out["variance"].permute([0, 3, 1, 2]).reshape(B*N, C, C), t)
            # exit()
        # At the first timestep return the decoder NLL,
        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
        output = torch.where((t == 0), decoder_nll, kl)
        return {"output": output, "pred_xstart": out["pred_xstart"]}

    def training_losses(self, x_start, t, ctx=None, reduce=True, valid_id=None, noise=None):
        """
        Compute training losses for a single timestep.
        :param model: the model to evaluate loss on.
        :param x_start: the [B x 3 x N] tensor of inputs.
        :param t: a batch of timestep indices.
        :param noise: if specified, the specific Gaussian noise to try to remove.
        :return: a dict with the key "loss" containing a tensor of shape [N].
                 Some mean or variance settings may also have other keys.
        """
        loss_dict = dict()
        if noise is None:
            noise = torch.randn_like(x_start)
        B, C, N = x_start.shape
        x_t = self.q_sample(x_start, t, noise=noise)
        
        model_output = self.model(x_t, self._scale_timesteps(t), ctx, valid_id=valid_id)
        if self.model_var_type in [
            ModelVarType.LEARNED,
            ModelVarType.LEARNED_RANGE,
        ]:
            # print(model_output.shape)
            assert model_output.shape == (B, C*2, N)
            model_output, model_var_values = torch.split(model_output, C, dim=1)
            loss_dict['model_var_value'] = model_var_values.mean()
            # Learn the variance using the variational bound, but don't let
            # it affect our mean prediction.
            frozen_out = torch.cat([model_output.detach(), model_var_values], dim=1)
            vb_loss= self._vb_terms_bpd(
                x_start=x_start,
                x_t=x_t,
                t=t,
                ctx=ctx,
                frozen_out=frozen_out,
            )["output"].mean()
            # print(vb_loss)
            if vb_loss > 1000:
                print("vb_loss is", vb_loss.mean())
                print("model var is ", model_var_values.mean())
                exit()
            vb_loss *= self.num_timesteps / 1000.0
            loss_dict['vb_loss'] = vb_loss
            
        target = {
                ModelMeanType.START_X: x_start,
                ModelMeanType.EPSILON: noise,
            }[self.model_mean_type]
        assert model_output.shape == target.shape
        diffusion_loss = (target - model_output) ** 2
        if valid_id is not None:
            diffusion_loss = diffusion_loss * valid_id.unsqueeze(1)
        if reduce:
            if valid_id is not None:
                diffusion_loss = diffusion_loss.mean(1).sum() / valid_id.sum()
            else:
                diffusion_loss = diffusion_loss.mean()
            
        loss_dict['mse_loss'] = diffusion_loss
        
        return loss_dict
python/difffacto/models/diffusions/latent_diffusion.py
diffFacto/diffFacto
import torch 
from torch.nn import Module
from anchor_diff.utils.registry import DIFFUSIONS, NETS, build_from_cfg
from anchor_diff.utils.constants import ModelMeanType, ModelVarType, LossType
from anchor_diff.utils.misc import *
from anchor_diff.utils.constants import *
import math
import numpy as np
from .diffusion_utils import extract_into_tensor, betas_for_alpha_bar, normal_kl


@DIFFUSIONS.register_module()
class LatentDiffusionDrift(Module):
    #TODO
    def __init__(
        self, 
        net, 
        num_timesteps, 
        beta_1, 
        beta_T, 
        mode='linear', 
        use_beta=False, 
        rescale_timesteps=False, 
        loss_type='mse', 
        model_mean_type='epsilon', 
        model_var_type='fixed_small', 
        classifier_weight=1.,
    ):
        super().__init__()
        assert mode in ('linear', 'cosine')
        self.model = build_from_cfg(net, NETS)
        self.model_mean_type = model_mean_dict[model_mean_type]
        self.model_var_type = model_var_dict[model_var_type]
        self.loss_type = model_loss_dict[loss_type]
        self.beta_1 = beta_1
        self.beta_T = beta_T
        self.use_beta = use_beta
        self.rescale_timesteps = rescale_timesteps
        self.classifier_weight=classifier_weight
        self.mode = mode
        if mode == 'linear':
            betas = np.linspace(beta_1, beta_T, num=num_timesteps, dtype=np.float64)
        elif mode == 'cosine':
            betas = betas_for_alpha_bar(num_timesteps, lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,)
            
        self.betas = betas
        assert len(betas.shape) == 1, "betas must be 1-D"
        assert (betas > 0).all() and (betas <= 1).all()

        self.num_timesteps = int(num_timesteps)

        alphas = 1.0 - betas
        self.alphas_cumprod = np.cumprod(alphas, axis=0)
        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)

        # calculations for diffusion q(x_t | x_{t-1}) and others
        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)

        # calculations for posterior q(x_{t-1} | x_t, x_0)
        self.posterior_variance = (
            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )
        # log calculation clipped because the posterior variance is 0 at the
        # beginning of the diffusion chain.
        self.posterior_log_variance_clipped = np.log(
            np.append(self.posterior_variance[1], self.posterior_variance[1:])
        )
        self.posterior_mean_coef1 = (
            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )
        self.posterior_mean_coef2 = (
            (1.0 - self.alphas_cumprod_prev)
            * np.sqrt(alphas)
            / (1.0 - self.alphas_cumprod)
        )
        self.posterior_mean_coef3 = (
            1.0 + ((self.alphas_cumprod - 1.) * ( np.sqrt(self.alphas_cumprod_prev) + np.sqrt(alphas)))
            / (1.0 - self.alphas_cumprod)
        )
        
        
    def q_mean_variance(self, x_start, t, prior):
        """
        Get the distribution q(x_t | x_0).
        :param x_0: the [N x 3 x n] tensor of noiseless inputs.
        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
        :param x_0: the [N x 3 x n] tensor of anchors the tensor of shape as x_0 to which x_0's diffuses
        :return: A tuple (mean, variance, log_variance), all of x_start's shape.
        """
        mean = (
            extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * (x_start - prior) + prior
        )
        variance = extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
        log_variance = extract_into_tensor(
            self.log_one_minus_alphas_cumprod, x_start.shape
        )
        return mean, variance, log_variance
    
    def q_sample(self, x_start, t, prior, noise=None):
        """
        Diffuse the data for a given number of diffusion steps.
        In other words, sample from q(x_t | x_0).
        :param x_0: the initial data batch.
        :param anchors: anchors for each point in x_0
        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
        :param noise: if specified, the split-out normal noise.
        :return: A noisy version of x_0.
        """
        if noise is None:
            noise = torch.randn_like(x_start)
        assert noise.shape == x_start.shape
        return (
            extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * (x_start - prior) + prior
            + extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
            * noise
        )
        
    def q_posterior_mean(self, x_start, x_t, t, prior):
        """
        Compute the mean and variance of the diffusion posterior:
            q(x_{t-1} | x_t, x_0)
        """
        assert x_start.shape == x_t.shape
        posterior_mean = (
            extract_into_tensor(self.posterior_mean_coef1, t, x_start.shape) * x_start
            + extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t 
            + extract_into_tensor(self.posterior_mean_coef3, t, x_t.shape) * prior
        )
        assert (
            posterior_mean.shape[0]
            == x_t.shape[0]
        )
        return posterior_mean
    
    def q_posterior_variance(self, x_start, x_t, t):
        """
        Compute the mean and variance of the diffusion posterior:
            q(x_{t-1} | x_t, x_0)
        """
        assert x_start.shape == x_t.shape
        
        posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape)
        posterior_log_variance_clipped = extract_into_tensor(
            self.posterior_log_variance_clipped, t, x_t.shape
        )
        assert (
            posterior_variance.shape[0]
            == posterior_log_variance_clipped.shape[0]
            == x_t.shape[0]
        )
        return posterior_variance, posterior_log_variance_clipped
    
    def q_posterior_mean_variance(self, x_start, x_t, t, prior):
        """
        Compute the mean and variance of the diffusion posterior:
            q(x_{t-1} | x_t, x_0)
        """
        assert x_start.shape == x_t.shape == prior.shape
    
        
        posterior_mean = self.q_posterior_mean(x_start, x_t, t, prior)
        posterior_variance, posterior_log_variance_clipped = self.q_posterior_variance(x_start, x_t, t)
        return posterior_mean, posterior_variance, posterior_log_variance_clipped
    
    
    def p_mean_variance(
        self, x, t, prior, code, part_indicator=None
    ):
        """
        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
        the initial x, x_0.
        :param model: the model, which takes a signal and a batch of timesteps
                      as input.
        :param x: the [N x C x ...] tensor at time t.
        :param t: a 1-D Tensor of timesteps.
        :param clip_denoised: if True, clip the denoised signal into [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample. Applies before
            clip_denoised.
        :return: a dict with the following keys:
                 - 'mean': the model mean output.
                 - 'variance': the model variance output.
                 - 'log_variance': the log of 'variance'.
                 - 'pred_xstart': the prediction for x_0.
        """

        B, C = x.shape[:2]
        assert t.shape == (B,)
        model_output_cond = self.model(x,  self._scale_timesteps(t), prior, code, part_indicator=part_indicator)

        model_output_uncond = self.model(x, self._scale_timesteps(t), torch.zeros_like(prior).cuda(), torch.zeros_like(code).cuda(), part_indicator=part_indicator)
        model_output = (1. - self.classifier_weight) * model_output_uncond + self.classifier_weight * model_output_cond
        model_variance, model_log_variance = {
            # for fixedlarge, we set the initial (log-)variance like so
            # to get a better decoder log likelihood.
            ModelVarType.FIXED_LARGE: (
                np.append(self.posterior_variance[1], self.betas[1:]),
                np.log(np.append(self.posterior_variance[1], self.betas[1:])),
            ),
            ModelVarType.FIXED_SMALL: (
                self.posterior_variance,
                self.posterior_log_variance_clipped,
            ),
        }[self.model_var_type]
        model_variance = extract_into_tensor(model_variance, t, x.shape)
        model_log_variance = extract_into_tensor(model_log_variance, t, x.shape)

        def process_xstart(x):
            return x
        prior = torch.zeros_like(prior)
        if self.model_mean_type == ModelMeanType.PREVIOUS_X:
            pred_xstart = process_xstart(
                self._predict_xstart_from_xprev(x_t=x, t=t, prior=prior, xprev=model_output)
            )
            model_mean = model_output
        elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]:
            if self.model_mean_type == ModelMeanType.START_X:
                pred_xstart = process_xstart(model_output)
            else:
                pred_xstart = process_xstart(
                    self._predict_xstart_from_eps(x_t=x, t=t, prior=prior, eps=model_output)
                )
            model_mean, _, _ = self.q_posterior_mean_variance(
                x_start=pred_xstart, x_t=x, t=t, prior=prior
            )
        else:
            raise NotImplementedError(self.model_mean_type)

        assert (
            model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
        )
        return {
            "mean": model_mean,
            "variance": model_variance,
            "log_variance": model_log_variance,
            "pred_xstart": pred_xstart,
        }
        
    def _predict_xstart_from_eps(self, x_t, t, prior, eps):
        assert x_t.shape == eps.shape == prior.shape
        
        return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * (x_t - prior) + prior
                - extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, eps.shape) * eps
                )

    def _predict_xstart_from_xprev(self, x_t, t, prior, xprev):
        assert x_t.shape == xprev.shape == prior.shape
        return (  # (xprev - coef2*x_t - coef3 * anchors) / coef1
            extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev
            - extract_into_tensor(
                self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape
            ) * x_t
            - extract_into_tensor(
                self.posterior_mean_coef3 / self.posterior_mean_coef1, t, x_t.shape
            ) * prior
        )
    
    def _predict_eps_from_xstart(self, x_t, t, prior, pred_xstart):
        return (
            extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * (x_t - prior)
            - pred_xstart
        ) / extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
        
    def _scale_timesteps(self, t):
        if self.rescale_timesteps:
            return t.float() * (1000.0 / self.num_timesteps)
        if self.use_beta:
            return extract_into_tensor(self.betas, t, t.shape)
        return t

    def p_sample(
        self, x, t, prior, code, part_indicator=None
    ):
        """
        Sample x_{t-1} from the model at the given timestep.
        :param model: the model to sample from.
        :param x: the current tensor at x_{t-1}.
        :param t: the value of t, starting at 0 for the first diffusion step.
        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample.
        :return: a dict containing the following keys:
                 - 'sample': a random sample from the model.
                 - 'pred_xstart': a prediction of x_0.
        """
        out = self.p_mean_variance(
            x,
            t,
            prior,
            code,
            part_indicator=part_indicator
        )
        noise = torch.randn_like(x)
        nonzero_mask = (
            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
        )  # no noise when t == 0
        sample = out["mean"] + nonzero_mask * torch.exp(0.5 * out["log_variance"]) * noise
        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
    
    def p_sample_loop(
        self,
        shape,
        prior,
        code,
        part_indicator=None,
        noise=None,
        device=None,
        progress=False,
    ):
        """
        Generate samples from the model.
        :param model: the model module.
        :param shape: the shape of the samples, (N, C, H, W).
        :param noise: if specified, the noise from the encoder to sample.
                      Should be of the same shape as `shape`.
        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample.
        :param device: if specified, the device to create the samples on.
                       If not specified, use a model parameter's device.
        :param progress: if True, show a tqdm progress bar.
        :return: a non-differentiable batch of samples.
        """
        final = None
        for t, sample in self.p_sample_loop_progressive(
            shape,
            prior=prior,
            noise=noise,
            code=code,
            part_indicator=part_indicator,
            device=device,
            progress=progress,
        ):
            final = sample
        return final["sample"]

    def p_sample_loop_progressive(
        self,
        shape,
        prior,
        code,
        part_indicator=None,
        noise=None,
        device=None,
        progress=False,
    ):
        """
        Generate samples from the model and yield intermediate samples from
        each timestep of diffusion.
        Arguments are the same as p_sample_loop().
        Returns a generator over dicts, where each dict is the return value of
        p_sample().
        """
        if device is None:
            device = next(self.model.parameters()).device
        assert isinstance(shape, (tuple, list))
        if noise is not None:
            pcd = noise + prior
        else:
            pcd = torch.randn(*shape, device=device) + prior
        indices = list(range(self.num_timesteps))[::-1]

        if progress:
            # Lazy import so that we don't depend on tqdm.
            from tqdm.auto import tqdm

            indices = tqdm(indices)

        for i in indices:
            t = torch.tensor([i] * shape[0], device=device)
            with torch.no_grad():
                out = self.p_sample(
                    pcd,
                    t,
                    prior,
                    code,
                    part_indicator
                )
                yield i, out
                pcd = out["sample"]
                
    def q_sample_loop(
        self,
        gt_pcd,
        prior,
        noise=None,
        device=None,
        progress=False,
    ):
        """
        forward sample from the model
        :param shape: the shape of the samples, (B, C, N).
        :param noise: if specified, the noise from the encoder to sample.
                      Should be of the same shape as `shape`.
        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample.
        :param device: if specified, the device to create the samples on.
                       If not specified, use a model parameter's device.
        :param progress: if True, show a tqdm progress bar.
        :return: a non-differentiable batch of samples.
        """
        final = None
        for t, sample in self.q_sample_loop_progressive(
            gt_pcd,
            noise=noise,
            prior=prior,
            device=device,
            progress=progress,
        ):
            final = sample

        return final 
    def q_sample_loop_progressive(self, gt_pcd, prior, noise=None, device=None, progress=False):
        if device is None:
            device = next(self.model.parameters()).device
        indices = list(range(self.num_timesteps))[1:]
        if noise is None:
            noise = torch.randn(*gt_pcd.shape).to(device)
        if progress:
            # Lazy import so that we don't depend on tqdm.
            from tqdm.auto import tqdm

            indices = tqdm(indices)
        for i in indices:
            t = torch.tensor([i] * gt_pcd.shape[0], device=device)
            with torch.no_grad():
                yield i, self.q_sample(gt_pcd, t, prior, noise=noise)
        
    def _vb_terms_bpd(
        self, x_start, x_t, t, code, model_output=None
    ):
        """
        Get a term for the variational lower-bound.
        This allows for comparison to other papers.
        :return: a dict with the following keys:
                 - 'output': a shape [N] tensor of NLLs or KLs.
                 - 'pred_xstart': the x_0 predictions.
        """
        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
            x_start=x_start, x_t=x_t, t=t
        )
        out = self.p_mean_variance(
            x_t, t, code, model_output=model_output
        )
        kl = normal_kl(
            true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
        )
        kl = mean_flat(kl) / np.log(2.0)

        decoder_nll = -gaussian_log_likelihood(
            x_start, mean=out["mean"], logvar=out["log_variance"]
        )
        assert decoder_nll.shape == x_start.shape
        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)

        # At the first timestep return the decoder NLL,
        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
        output = torch.where((t == 0), decoder_nll, kl)
        return {"output": output, "pred_xstart": out["pred_xstart"]}

    def training_losses(self, x_start, t, prior, code, part_indicator=None, noise=None):
        """
        Compute training losses for a single timestep.
        :param model: the model to evaluate loss on.
        :param x_start: the [B x 3 x N] tensor of inputs.
        :param t: a batch of timestep indices.
        :param noise: if specified, the specific Gaussian noise to try to remove.
        :return: a dict with the key "loss" containing a tensor of shape [N].
                 Some mean or variance settings may also have other keys.
        """
        if noise is None:
            noise = torch.randn_like(x_start)
        x_t = self.q_sample(x_start, t, torch.zeros_like(prior), noise=noise)
        # x_t = self.q_sample(x_start, t, prior, noise=noise)

        terms = {}

        if self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:
            model_output = self.model(x_t, self._scale_timesteps(t), prior, code, part_indicator=part_indicator)
            target = {
                ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(
                    x_start=x_start, x_t=x_t, t=t
                )[0],
                ModelMeanType.START_X: x_start,
                ModelMeanType.EPSILON: noise,
            }[self.model_mean_type]
            assert model_output.shape == target.shape == x_start.shape
            terms["mse"] = mean_flat((target - model_output) ** 2)
            if "vb" in terms:
                terms["recons_loss"] = terms["mse"] + terms["vb"]
            else:
                terms["recons_loss"] = terms["mse"]
        else:
            raise NotImplementedError(self.loss_type)
        
        return terms["recons_loss"]
python/difffacto/models/diffusions/latent_diffusion_drift.py
diffFacto/diffFacto
from .pointwisenet import PointwiseNet, PointwiseNetLatent
from .unet import UNet
from .pointwisenet2 import PointwiseNet2
from .attention import TransformerNet
python/difffacto/models/diffusions/nets/__init__.py
diffFacto/diffFacto
from inspect import isfunction
import math
import torch
import torch.nn.functional as F
from torch import nn, einsum
from einops import rearrange, repeat
from typing import Optional, Any, Union, List, Tuple

from difffacto.utils.misc import checkpoint
from .unet import timestep_embedding
from .utils import SineLayer
from difffacto.utils.registry import NETS


try:
    import xformers
    import xformers.ops
    print(torch.cuda.get_device_name(0))
    XFORMERS_IS_AVAILBLE = (torch.cuda.get_device_name(0) not in  ['NVIDIA TITAN RTX', 'TITAN Xp', 'NVIDIA TITAN Xp', 'Tesla V100-DGXS-16GB', 'NVIDIA TITAN X (Pascal)'])
except:
    XFORMERS_IS_AVAILBLE = False    

print(XFORMERS_IS_AVAILBLE)
def exists(val):
    return val is not None


def uniq(arr):
    return{el: True for el in arr}.keys()


def default(val, d):
    if exists(val):
        return val
    return d() if isfunction(d) else d


def max_neg_value(t):
    return -torch.finfo(t.dtype).max


def init_(tensor):
    dim = tensor.shape[-1]
    std = 1 / math.sqrt(dim)
    tensor.uniform_(-std, std)
    return tensor


# feedforward
class GEGLU(nn.Module):
    def __init__(self, dim_in, dim_out):
        super().__init__()
        self.proj = nn.Linear(dim_in, dim_out * 2)

    def forward(self, x):
        x, gate = self.proj(x).chunk(2, dim=-1)
        return x * F.gelu(gate)


class MLP(nn.Module):

    def forward(self, x, *_):
        return self.net(x)

    def __init__(self, ch: Union[List[int], Tuple[int, ...]], use_linear: bool = False, act: nn.Module = nn.ReLU,
                 weight_norm=False):
        super(MLP, self).__init__()
        layers = []
        for i in range(len(ch) - 1):
            layers.append(nn.Linear(ch[i], ch[i + 1]) if use_linear else nn.Conv1d(ch[i], ch[i + 1], 1))
            if weight_norm:
                layers[-1] = nn.utils.weight_norm(layers[-1])
            if i < len(ch) - 2:
                layers.append(act(True))
        self.net = nn.Sequential(*layers)

class FeedForward(nn.Module):
    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
        super().__init__()
        inner_dim = int(dim * mult)
        dim_out = default(dim_out, dim)
        project_in = nn.Sequential(
            nn.Linear(dim, inner_dim),
            nn.GELU()
        ) if not glu else GEGLU(dim, inner_dim)

        self.net = nn.Sequential(
            project_in,
            nn.Dropout(dropout),
            nn.Linear(inner_dim, dim_out)
        )

    def forward(self, x):
        return self.net(x)


def zero_module(module):
    """
    Zero out the parameters of a module and return it.
    """
    for p in module.parameters():
        p.detach().zero_()
    return module


def Normalize(in_channels):
    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)


class SpatialSelfAttention(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.in_channels = in_channels

        self.norm = Normalize(in_channels)
        self.q = torch.nn.Conv1d(in_channels,
                                 in_channels,
                                 kernel_size=1,
                                 stride=1,
                                 padding=0)
        self.k = torch.nn.Conv1d(in_channels,
                                 in_channels,
                                 kernel_size=1,
                                 stride=1,
                                 padding=0)
        self.v = torch.nn.Conv1d(in_channels,
                                 in_channels,
                                 kernel_size=1,
                                 stride=1,
                                 padding=0)
        self.proj_out = torch.nn.Conv1d(in_channels,
                                        in_channels,
                                        kernel_size=1,
                                        stride=1,
                                        padding=0)

    def forward(self, x):
        h_ = x
        h_ = self.norm(h_)
        q = self.q(h_)
        k = self.k(h_)
        v = self.v(h_)

        # compute attention
        b,c,h,w = q.shape
        q = rearrange(q, 'b c n -> b n c')
        k = rearrange(k, 'b c n -> b c n')
        w_ = torch.einsum('bij,bjk->bik', q, k)

        w_ = w_ * (int(c)**(-0.5))
        w_ = torch.nn.functional.softmax(w_, dim=2)

        # attend to values
        w_ = rearrange(w_, 'b i j -> b j i')
        h_ = torch.einsum('bij,bjk->bik', v, w_)
        h_ = self.proj_out(h_)

        return x+h_


class CrossAttention(nn.Module):
    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
        super().__init__()
        inner_dim = dim_head * heads
        context_dim = default(context_dim, query_dim)

        self.scale = dim_head ** -0.5
        self.heads = heads

        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, query_dim),
            nn.Dropout(dropout)
        )

    def forward(self, x, context=None, mask=None):
        h = self.heads

        q = self.to_q(x)
        context = default(context, x)
        k = self.to_k(context)
        v = self.to_v(context)

        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))

        sim = einsum('b i d, b j d -> b i j', q, k)* self.scale
        del q, k

        if exists(mask):
            assert mask.shape == (x.shape[0], context.shape[1])
            mask = rearrange(mask, 'b ... -> b (...)')
            max_neg_value = -torch.finfo(sim.dtype).max
            mask = repeat(mask, 'b j -> (b h) () j', h=h).to(bool)
            sim.masked_fill_(~mask, max_neg_value)

        # attention, what we cannot get enough of
        sim = sim.softmax(dim=-1)

        out = einsum('b i j, b j d -> b i d', sim, v)
        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
        return self.to_out(out)


class MemoryEfficientCrossAttention(nn.Module):
    # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):
        super().__init__()
        print(f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using "
              f"{heads} heads.")
        inner_dim = dim_head * heads
        context_dim = default(context_dim, query_dim)

        self.heads = heads
        self.dim_head = dim_head

        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)

        self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
        self.attention_op: Optional[Any] = None

    def forward(self, x, context=None, mask=None):
        q = self.to_q(x)
        context = default(context, x)
        k = self.to_k(context)
        v = self.to_v(context)

        b, _, _ = q.shape
        q, k, v = map(
            lambda t: t.unsqueeze(3)
            .reshape(b, t.shape[1], self.heads, self.dim_head)
            .permute(0, 2, 1, 3)
            .reshape(b * self.heads, t.shape[1], self.dim_head)
            .contiguous(),
            (q, k, v),
        )

        # actually compute the attention, what we cannot get enough of
        if exists(mask):
            assert mask.shape == (b, context.shape[1])
            max_neg_value = -torch.finfo(x.dtype).max 
            mask = (1 - mask) * max_neg_value
            mask = mask.repeat_interleave(self.heads, dim=0).view(b*self.heads, 1, context.shape[1]).expand(-1, x.shape[1], -1)
        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask, op=self.attention_op)

        out = (
            out.unsqueeze(0)
            .reshape(b, self.heads, out.shape[1], self.dim_head)
            .permute(0, 2, 1, 3)
            .reshape(b, out.shape[1], self.heads * self.dim_head)
        )
        return self.to_out(out)


class BasicTransformerBlock(nn.Module):
    ATTENTION_MODES = {
        "softmax": CrossAttention,  # vanilla attention
        "softmax-xformers": MemoryEfficientCrossAttention
    }
    def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=False,
                 single_attn=False, adaln=False, y_dim=32):
        super().__init__()
        attn_mode = "softmax-xformers" if XFORMERS_IS_AVAILBLE else "softmax"
        assert attn_mode in self.ATTENTION_MODES
        attn_cls = self.ATTENTION_MODES[attn_mode]
        self.single_attn = single_attn
        self.adaln=adaln
        
        if not self.single_attn:
            self.attn1 = attn_cls(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout,
                                context_dim=context_dim)  # is a self-attention if not self.disable_self_attn
            self.norm1 = nn.LayerNorm(dim)
                
        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
        self.attn2 = attn_cls(query_dim=dim, context_dim=context_dim,
                              heads=n_heads, dim_head=d_head, dropout=dropout)  # is self-attn if context is none
        self.norm2 = nn.LayerNorm(dim)
        self.norm3 = nn.LayerNorm(dim)
        if adaln:
            self.emb_layers = nn.Sequential(
                nn.SiLU(),
                nn.Linear(
                    y_dim,
                    2 * dim,
                ),
            )
        self.checkpoint = checkpoint

    def forward(self, x, context=None, mask=None, y=None):
        return checkpoint(self._forward, (x, context, mask, y), self.parameters(), self.checkpoint)

    def _forward(self, x, context=None, mask=None, y=None):
        if not self.single_attn:
            x = self.attn1(self.norm1(x), context=context , mask=mask) + x
        x = self.attn2(self.norm2(x), context=context, mask=mask) + x
        if self.adaln:
            shift, scale = torch.chunk(self.emb_layers(y), 2, dim=-1)
            _x = self.norm3(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
        else:
            _x = self.norm3(x)
        x = self.ff(_x) + x
        return x

@NETS.register_module()
class TransformerNet(nn.Module):
    """
    Transformer block for image-like data.
    First, project the input (aka embedding)
    and reshape to b, t, d.
    Then apply standard transformer action.
    Finally, reshape to image
    NEW: use_linear for more efficiency instead of the 1x1 convs
    """
    def __init__(self, in_channels, n_heads, d_head, out_channels,
                 depth=1, dropout=0., context_dim=None, use_linear=False,
                 use_checkpoint=False, single_attn=False, class_cond=False, n_class=4,
                 cat_params_to_x=False, mask_out_unreferenced_code=True, cat_class_to_x=False, 
                 use_sine_proj_in=False, add_t_to_x=False, res=False, add_class_cond=False,
                 context_proj=False, include_std=False
                ):
        super().__init__()
        self.class_cond=class_cond
        self.add_class_cond=add_class_cond
        self.cat_class_to_x=cat_class_to_x
        in_channels = in_channels + int(cat_params_to_x) * 6 + int(cat_class_to_x) * n_class
        self.in_channels = in_channels
        self.mask_out_unreferenced_code=mask_out_unreferenced_code
        inner_dim = n_heads * d_head
        self.inner_dim = inner_dim
        self.context_dim =  context_dim + 256 + int(class_cond and not add_class_cond) * n_class if not add_t_to_x else context_dim + int(class_cond and not add_class_cond) * n_class
        if class_cond and add_class_cond:
            self.class_emb = nn.Embedding(n_class, inner_dim)
        self.n_class=n_class
        self.include_std=include_std
        self.cat_params_to_x=cat_params_to_x
        self.add_t_to_x=add_t_to_x
        self.res=res
        self.context_proj=context_proj
        if context_proj:
            self.ctx_proj_in = nn.Linear(self.context_dim,inner_dim)
            self.ctx_norm = nn.LayerNorm(inner_dim)
        self.pre_norm = nn.LayerNorm(inner_dim)
        self.post_norm = nn.LayerNorm(inner_dim)
        if not use_linear:
            self.proj_in = nn.Conv1d(in_channels,
                                     inner_dim,
                                     kernel_size=1,
                                     stride=1,
                                     padding=0) if not use_sine_proj_in else SineLayer(in_channels, inner_dim, bias=True, is_first=True, use_linear=True)
        else:
            self.proj_in = nn.Linear(in_channels, inner_dim) if not use_sine_proj_in else SineLayer(in_channels, inner_dim, bias=True, is_first=True, use_linear=False)

        self.time_embed = FeedForward(256, dropout=dropout, glu=True) if not add_t_to_x else FeedForward(inner_dim, dropout=dropout, glu=True)
        self.transformer_blocks = nn.ModuleList(
            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=self.context_dim if not self.context_proj else inner_dim,
                                   checkpoint=use_checkpoint, single_attn=single_attn, adaln=False)
                for d in range(depth)]
        )
        if not use_linear:
            if use_sine_proj_in:
                self.proj_out = MLP([inner_dim] + [inner_dim * 2] * 5 + [out_channels], use_linear=False)
            else:
                self.proj_out = zero_module(nn.Conv1d(inner_dim,
                                                    out_channels,
                                                    kernel_size=1,
                                                    stride=1,
                                                    padding=0)) \
                                            if in_channels == out_channels or res else \
                                            nn.Conv1d(inner_dim,
                                                    out_channels,
                                                    kernel_size=1,
                                                    stride=1,
                                                    padding=0)
        else:
            if use_sine_proj_in:
                self.proj_out = MLP([inner_dim] + [inner_dim * 2] * 5 + [out_channels], use_linear=True)
            else:
                self.proj_out = zero_module(nn.Linear(inner_dim, out_channels)) if in_channels == out_channels or res else nn.Linear(inner_dim, out_channels)
        self.use_linear = use_linear

    def forward(self, x, t, ctx, anchors=None, variances=None, valid_id=None, anchor_assignment=None, **kwargs):
        if isinstance(ctx, list):
            ctx = torch.cat(ctx, dim=1)
        ctx = rearrange(ctx, 'b c n -> b n c').contiguous()
        
        if self.class_cond and not self.add_class_cond:
            class_embed = torch.eye(self.n_class).to(x).unsqueeze(0).repeat_interleave(x.shape[0], dim=0) # B, n_ctx, n_ctx
            ctx = torch.cat([ctx, class_embed], dim=-1)
        t_embed = self.time_embed(timestep_embedding(t, 256 if not self.add_t_to_x else self.inner_dim))
        if not self.add_t_to_x:
            ctx = torch.cat([ctx, t_embed.unsqueeze(1).expand(-1, ctx.shape[1], -1)], dim=-1)
        
        assert ctx.shape[-1] == self.context_dim
        if self.cat_params_to_x:
            if not self.include_std:
                x = torch.cat([x, anchors.transpose(1,2), variances.transpose(1,2)], dim=1)
            else:
                x = torch.cat([x, anchors.transpose(1,2), torch.sqrt(variances).transpose(1,2)], dim=1)
                
            
        if self.cat_class_to_x:
            cls_onehot = F.one_hot(anchor_assignment.to(torch.long), num_classes=self.n_class).transpose(1,2)
            x = torch.cat([x, cls_onehot], dim=1)
        assert x.shape[1] == self.in_channels
        return self._forward_attn(x, ctx, mask=valid_id if self.mask_out_unreferenced_code else None, t_embed=t_embed)
    
    def _forward_attn(self, x, context=None, mask=None, t_embed=None):
        # note: if no context is given, cross-attention defaults to self-attention
        b, c, n = x.shape
        x_in = x
        if self.context_proj:
            context = self.ctx_proj_in(context)
            if self.add_class_cond:
                context = context + self.class_emb(torch.arange(self.n_class, device=x.device)).unsqueeze(0)
            context = self.ctx_norm(context)
        if not self.use_linear:
            x = self.proj_in(x)
        x = rearrange(x, 'b c n -> b n c').contiguous()
        if self.use_linear:
            x = self.proj_in(x)
        if self.add_t_to_x:
            x = x + t_embed.unsqueeze(1)
        x = self.pre_norm(x)
        for i, block in enumerate(self.transformer_blocks):
            x = block(x, context=context, mask=mask)
        x = self.post_norm(x)
        if self.use_linear:
            x = self.proj_out(x)
        x = rearrange(x, 'b n c -> b c n').contiguous()
        if not self.use_linear:
            x = self.proj_out(x)
        if x_in.shape[1] == x.shape[1]:
            return x + x_in
        if self.res:
            return x + x_in[:, :x.shape[1]]
        return x
    

@NETS.register_module()
class LDMNet(nn.Module):
    """
    Transformer block for image-like data.
    First, project the input (aka embedding)
    and reshape to b, t, d.
    Then apply standard transformer action.
    Finally, reshape to image
    NEW: use_linear for more efficiency instead of the 1x1 convs
    """
    def __init__(self, in_channels, n_heads, d_head, out_channels,
                 depth=1, dropout=0., context_dim=None, use_linear=False,
                 use_checkpoint=False, single_attn=False, class_cond=False, n_class=4,
                 mask_out_unreferenced_code=True,  cond_time_as_token=True
                ):
        super().__init__()
        self.class_cond=class_cond
        in_channels = in_channels
        self.in_channels = in_channels
        self.mask_out_unreferenced_code=mask_out_unreferenced_code
        inner_dim = n_heads * d_head
        self.inner_dim = inner_dim
        self.context_dim =  context_dim
        if class_cond:
            self.class_emb = nn.Embedding(n_class, inner_dim)
        self.n_class=n_class
        self.cond_time_as_token=cond_time_as_token
        self.pre_norm = nn.LayerNorm(inner_dim)
        self.post_norm = nn.LayerNorm(inner_dim)
        if not use_linear:
            self.proj_in = nn.Conv1d(in_channels,
                                     inner_dim,
                                     kernel_size=1,
                                     stride=1,
                                     padding=0) 
        else:
            self.proj_in = nn.Linear(in_channels, inner_dim) 

        self.time_embed = FeedForward(inner_dim, dropout=dropout, glu=True) 
        self.transformer_blocks = nn.ModuleList(
            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=self.context_dim,
                                   checkpoint=use_checkpoint, single_attn=single_attn)
                for d in range(depth)]
        )
        if not use_linear:
            self.proj_out = zero_module(nn.Conv1d(inner_dim,
                                                out_channels,
                                                kernel_size=1,
                                                stride=1,
                                                padding=0)) 
        else:
            self.proj_out = zero_module(nn.Linear(inner_dim, out_channels))
        self.use_linear = use_linear

    def forward(self, x, t, ctx, valid_id=None,  **kwargs):
        if isinstance(ctx, list):
            ctx = torch.cat(ctx, dim=1)
        if ctx is not None:
            ctx = rearrange(ctx, 'b c n -> b n c').contiguous()
        
        if self.class_cond:
            class_embed = self.class_emb(torch.arange(self.n_class).to(x.device)).unsqueeze(0).repeat_interleave(x.shape[0], dim=0) # B, n_ctx, n_ctx
        t_embed = self.time_embed(timestep_embedding(t, 256))
        
                
            
        assert x.shape[1] == self.in_channels
        return self._forward_with_cond(x, cond_as_token=[(t_embed, self.cond_time_as_token), (class_embed, False)] if self.class_cond else [(t_embed, self.cond_time_as_token)], ctx=ctx, mask=valid_id if self.mask_out_unreferenced_code else None)
    
    def _forward_with_cond(
        self, x: torch.Tensor, cond_as_token: List[Tuple[torch.Tensor, bool]], ctx: torch.Tensor=None, mask: torch.Tensor=None
    ) -> torch.Tensor:
        if not self.use_linear:
            x = self.proj_in(x)
        x = rearrange(x, 'b c n -> b n c').contiguous()
        if self.use_linear:
            x = self.proj_in(x)
        for emb, as_token in cond_as_token:
            if not as_token:
                x = x + emb[:, None] if len(emb.shape) == 2 else emb
        extra_tokens = [
            (emb[:, None] if len(emb.shape) == 2 else emb)
            for emb, as_token in cond_as_token
            if as_token
        ]
        if len(extra_tokens):
            x = torch.cat(extra_tokens + [x], dim=1)
            if mask is not None:
                new_mask = torch.ones(x.shape[0], len(extra_tokens) + self.n_class).to(x)
                new_mask[:, len(extra_tokens):] = mask
                mask = new_mask

        x = self.pre_norm(x)
        for i, block in enumerate(self.transformer_blocks):
            x = block(x, context=ctx, mask=mask)
        x = self.post_norm(x)
        if len(extra_tokens):
            x = x[:, sum(h.shape[1] for h in extra_tokens) :]
        if self.use_linear:
            x = self.proj_out(x)
        x = rearrange(x, 'b n c -> b c n').contiguous()
        if not self.use_linear:
            x = self.proj_out(x)
        return x
python/difffacto/models/diffusions/nets/attention.py
diffFacto/diffFacto
from torch.nn import Module, ModuleList, Linear
import torch
from difffacto.utils.registry import NETS
import torch.nn.functional as F
from difffacto.utils.constants import *
from .utils import ConcatSquashLinear

@NETS.register_module()
class PointwiseNetLatent(Module):

    def __init__(self, in_channels, out_channels, context_dim, res=True, use_part_ind=False, **kwargs):
        super().__init__()
        self.act = F.leaky_relu   
        self.res=res     

        self.layers = ModuleList([
            ConcatSquashLinear(in_channels, 128, context_dim+3),
            ConcatSquashLinear(128, 256, context_dim+3),
            ConcatSquashLinear(256, 512, context_dim+3),
            ConcatSquashLinear(512, 256, context_dim+3),
            ConcatSquashLinear(256, 128, context_dim+3),
            ConcatSquashLinear(128, out_channels, context_dim+3)
        ])
        self.use_part_ind=use_part_ind
    def forward(self, x, beta, prior, code, part_indicator, anchors=None):
        """
        Args:
            x:  Point clouds at some timestep t, (B, d, N).
            beta:     Time. (B, ).
            context:  Shape latents. (B, F) or (B, N, F).
        """
        batch_size, _, num_points = x.size()
        beta = beta.view(batch_size, 1, 1)          # (B, 1, 1)
        time_emb = torch.cat([beta, torch.sin(beta), torch.cos(beta)], dim=-1)  # (B, 1, 3)
        ori_x = x
        out = x.transpose(2, 1)
        prior = prior.transpose(1, 2)
        
        if len(code.shape) == 2:
            code = code.unsqueeze(1)
        elif len(code.shape) == 3:
            code = code.transpose(1,2)
        # print(code.shape, prior.shape)
        code = code.expand(-1, prior.shape[1], -1)
        if self.use_part_ind:
            code = code * part_indicator.unsqueeze(-1) # zero out the language latent for unrelavent parts
        ctx_emb = torch.cat([time_emb.expand(-1, prior.shape[1], -1), prior, code], dim=-1)
        assert len(ctx_emb.shape) == 3
        assert ctx_emb.shape[1] == num_points
        
        
        
        
        for i, layer in enumerate(self.layers):
            out = layer(ctx=ctx_emb, x=out)
            if i < len(self.layers) - 1:
                out = self.act(out)

        return out.transpose(2, 1) + ori_x if self.res else out.transpose(2, 1)
    
@NETS.register_module()
class PointwiseNet(Module):

    def __init__(self, in_channels, out_channels, context_dim, res=True):
        super().__init__()
        self.act = F.leaky_relu   
        self.res=res     
        point_dim = 3 


        self.layers = ModuleList([
            ConcatSquashLinear(in_channels, 128, context_dim+point_dim),
            ConcatSquashLinear(128, 256, context_dim+point_dim),
            ConcatSquashLinear(256, 512, context_dim+point_dim),
            ConcatSquashLinear(512, 256, context_dim+point_dim),
            ConcatSquashLinear(256, 128, context_dim+point_dim),
            ConcatSquashLinear(128, out_channels, context_dim+point_dim)
        ])
    def forward(self, x, beta, context, **kwargs):
        """
        Args:
            x:  Point clouds at some timestep t, (B, d, N).
            beta:     Time. (B, ).
            context:  Shape latents. (B, F) or (B, N, F).
        """
        batch_size, _, num_points = x.size()
        beta = beta.view(batch_size, 1, 1)          # (B, 1, 1)
        time_emb = torch.cat([beta, torch.sin(beta), torch.cos(beta)], dim=-1)  # (B, 1, 3)
        ori_x = x
        out = x.transpose(2, 1)
        if(context[0].shape[1] == 1):
            ctx_emb = torch.cat([time_emb] + context, dim=-1)    # (B, N, F+3)
        else:
            ctx_emb = torch.cat([time_emb.repeat_interleave(num_points, dim=1)] + context, dim=-1)    # (B, N, F+3)
        for i, layer in enumerate(self.layers):
            out = layer(ctx=ctx_emb, x=out)
            if i < len(self.layers) - 1:
                out = self.act(out)

        return out.transpose(2, 1) + ori_x if self.res else out.transpose(2, 1)
python/difffacto/models/diffusions/nets/pointwisenet.py
diffFacto/diffFacto
from torch.nn import Module, ModuleList, Linear
import torch
from difffacto.utils.registry import NETS
import torch.nn.functional as F
from difffacto.utils.constants import *
from pointnet2_ops.pointnet2_utils import gather_operation
from .utils import ConcatSquashLinear

@NETS.register_module()
class PointwiseNet2(Module):

    def __init__(self, 
                 in_channels, 
                 out_channels, 
                 context_dim, 
                 num_anchors, 
                 res=True, 
                 include_anchors=False, 
                 gather_layers=[1, 3], 
                 add_to_context=False,
                 K=1
                 ):
        super().__init__()
        self.act = F.leaky_relu   
        self.res=res     
        point_dim = 6 if include_anchors else 3
        self.num_anchors = num_anchors
        self.K=K
        self.gather_layers=gather_layers
        out_chs = [in_channels, 128, 256, 512, 256, 128, out_channels]
        self.layers = ModuleList()
        self.add_to_context=add_to_context
        self.context_dim = context_dim + point_dim
        if add_to_context:
            ch = 0
            for i in range(len(out_chs) - 1):
                if i - 1 in gather_layers:
                    ch = out_chs[i]
                self.layers.append(ConcatSquashLinear(out_chs[i], out_chs[i + 1], context_dim + ch + point_dim))
        else:
            for i in range(len(out_chs) - 1):
                self.layers.append(ConcatSquashLinear(out_chs[i] if i - 1 not in gather_layers else out_chs[i] * 2, out_chs[i + 1], context_dim+point_dim))
            
        self.include_anchors=include_anchors
    def gather_idx(self, x, anchor_assignment):
        # compute pointwise distance
        if x.shape[1] > 3:
            x = x[:, :3]
        dist = torch.pow(x.unsqueeze(-1) - x.unsqueeze(-2), exponent=2).sum(1) #[B, N, N]
        print(dist.shape)
        out_id = []
        for i in range(self.num_anchors):
            mask = (anchor_assignment != i).to(torch.float32) * 1e9 
            dist_ = dist + mask.unsqueeze(-1) 
            out_id.append(torch.topk(dist_, k=self.K, dim=1, largest=False)[1]) # [B, K, N]
        return torch.cat(out_id, dim=1) 
    def forward(self, x, beta, context, anchor_assignment, anchors=None, variances=None):
        """
        Args:
            x:  Point clouds at some timestep t, (B, d, N).
            beta:     Time. (B, ).
            context:  Shape latents. (B, F) or (B, N, F).
        """
        B, C, N = x.shape
        batch_size, _, num_points = x.size()
        beta = beta.view(batch_size, 1, 1)          # (B, 1, 1)
        time_emb = torch.cat([beta, torch.sin(beta), torch.cos(beta)], dim=-1)  # (B, 1, 3)
        ori_x = x
        out = x.transpose(2, 1)

        if len(context.shape) == 2:
            context = context.view(batch_size, 1, -1)   # (B, 1, F)
            ctx_emb = torch.cat([time_emb, context], dim=-1)    # (B, 1, F+3)
        else:
            assert len(context.shape) == 3
            assert context.shape[1] == num_points
            ctx_emb = torch.cat([time_emb.expand(-1, context.shape[1], -1), context], dim=-1)    # (B, N, F+3)
        
        if self.include_anchors:
            ctx_emb = torch.cat([ctx_emb, anchors], dim=-1)
        
        
        for i, layer in enumerate(self.layers):
            # print(ctx_emb.shape)
            out = layer(ctx=ctx_emb, x=out)
            if i < len(self.layers) - 1 and i not in self.gather_layers:
                out = self.act(out)
            elif i in self.gather_layers:
                feats = [] #[B, C, 4]
                for j in range(self.num_anchors):
                    ids = anchor_assignment != j
                    mask = ids.to(torch.float32) * -10000
                    feat = torch.max(out * mask.unsqueeze(-1), dim=1)[0]
                    feats.append(feat)
                feats = torch.stack(feats, dim=2).contiguous()
                neighboring_feat = gather_operation(feats, anchor_assignment).transpose(1,2)
                # print(neighboring_feat.shape)
                if self.add_to_context:
                    ctx_emb = ctx_emb[..., :self.context_dim]
                    ctx_emb = torch.cat([ctx_emb, neighboring_feat], dim=-1)
                else:
                    out = torch.cat([out, neighboring_feat], dim=-1)

        return out.transpose(2, 1) + ori_x if self.res else out.transpose(2, 1)
python/difffacto/models/diffusions/nets/pointwisenet2.py
diffFacto/diffFacto
from torch.nn import Module, ModuleList, Linear
from torch import nn, einsum
import torch
from difffacto.utils.registry import NETS
import torch.nn.functional as F
from difffacto.utils.constants import *
from .utils import ConcatSquashLinear, timestep_embedding
from inspect import isfunction 
import math
from einops import rearrange, repeat


class SinusoidalEmbedding3D(nn.Module):
    def __init__(self, out_dim):
        super().__init__()
        self.out_dim=out_dim 
        if out_dim % 6 :
            self.intermediate_dim = int(math.ceil(out_dim / 6) * 6)
        else:
            self.intermediate_dim = out_dim 
        self.linear = nn.Linear(self.intermediate_dim, out_dim, bias=False)
        self.scaling_factor = torch.ones(self.intermediate_dim // 6, dtype=torch.float32) * (10000 ** (1 / self.intermediate_dim))
        self.scaling_factor = torch.pow(self.scaling_factor, torch.arange(self.intermediate_dim // 6) * 6)

    def forward(self, x):
        # x[B, N, 3]
        B, N, _ = x.shape
        cos_x, sin_x = torch.cos(x.unsqueeze(3) / self.scaling_factor.reshape(1, 1, 1, -1).to(x.device)), torch.sin(x.unsqueeze(3) / self.scaling_factor.reshape(1, 1, 1, -1).to(x.device)) #[B, N, 3, D // 6]
        sinusoidal_pe = torch.stack([sin_x, cos_x], dim=-1).reshape(B, N, self.intermediate_dim) #[B, N, 3, D // 6, 2] -> [B, N, D]
        sinusoidal_pe = self.linear(sinusoidal_pe)
        return sinusoidal_pe

def exists(val):
    return val is not None


def uniq(arr):
    return{el: True for el in arr}.keys()


def default(val, d):
    if exists(val):
        return val
    return d() if isfunction(d) else d


def max_neg_value(t):
    return -torch.finfo(t.dtype).max


def init_(tensor):
    dim = tensor.shape[-1]
    std = 1 / math.sqrt(dim)
    tensor.uniform_(-std, std)
    return tensor


# feedforward
class GEGLU(nn.Module):
    def __init__(self, dim_in, dim_out):
        super().__init__()
        self.proj = nn.Linear(dim_in, dim_out * 2)

    def forward(self, x):
        x, gate = self.proj(x).chunk(2, dim=-1)
        return x * F.gelu(gate)


class FeedForward(nn.Module):
    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
        super().__init__()
        inner_dim = int(dim * mult)
        dim_out = default(dim_out, dim)
        project_in = nn.Sequential(
            nn.Linear(dim, inner_dim),
            nn.GELU()
        ) if not glu else GEGLU(dim, inner_dim)

        self.net = nn.Sequential(
            project_in,
            nn.Dropout(dropout),
            nn.Linear(inner_dim, dim_out)
        )

    def forward(self, x):
        return self.net(x)


def zero_module(module):
    """
    Zero out the parameters of a module and return it.
    """
    for p in module.parameters():
        p.detach().zero_()
    return module


def Normalize(in_channels):
    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)


class LinearAttention(nn.Module):
    def __init__(self, dim, heads=4, dim_head=32):
        super().__init__()
        self.heads = heads
        hidden_dim = dim_head * heads
        self.to_qkv = nn.Conv1d(dim, hidden_dim * 3, 1, bias = False)
        self.to_out = nn.Conv1d(hidden_dim, dim, 1)

    def forward(self, x):
        b, c, h, w = x.shape
        qkv = self.to_qkv(x)
        q, k, v = rearrange(qkv, 'b (qkv heads c) p -> qkv b heads c p', heads = self.heads, qkv=3)
        k = k.softmax(dim=-1)  
        context = torch.einsum('bhdn,bhen->bhde', k, v)
        out = torch.einsum('bhde,bhdn->bhen', context, q)
        out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
        return self.to_out(out)


class SpatialSelfAttention(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.in_channels = in_channels

        self.norm = Normalize(in_channels)
        self.q = torch.nn.Conv1d(in_channels,
                                 in_channels,
                                 kernel_size=1,
                                 stride=1,
                                 padding=0)
        self.k = torch.nn.Conv1d(in_channels,
                                 in_channels,
                                 kernel_size=1,
                                 stride=1,
                                 padding=0)
        self.v = torch.nn.Conv1d(in_channels,
                                 in_channels,
                                 kernel_size=1,
                                 stride=1,
                                 padding=0)
        self.proj_out = torch.nn.Conv1d(in_channels,
                                        in_channels,
                                        kernel_size=1,
                                        stride=1,
                                        padding=0)

    def forward(self, x):
        h_ = x
        h_ = self.norm(h_)
        q = self.q(h_)
        k = self.k(h_)
        v = self.v(h_)

        # compute attention
        b,c,p = q.shape
        q = rearrange(q, 'b c p -> b p c')
        w_ = torch.einsum('bij,bjk->bik', q, k)

        w_ = w_ * (int(c)**(-0.5))
        w_ = torch.nn.functional.softmax(w_, dim=2)

        # attend to values
        w_ = rearrange(w_, 'b i j -> b j i')
        h_ = torch.einsum('bij,bjk->bik', v, w_)
        h_ = self.proj_out(h_)

        return x+h_


class CrossAttention(nn.Module):
    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., rel_pe=False, num_part=4, use_anchor_as_pe=False, embed_channel=128, use_scale_shift_norm=True):
        super().__init__()
        inner_dim = dim_head * heads
        context_dim = default(context_dim, query_dim)

        self.scale = dim_head ** -0.5
        self.heads = heads

        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
        self.rel_pe = rel_pe
        if self.rel_pe and not use_anchor_as_pe:
            self.relative_position_bias_table = nn.Parameter(torch.zeros(heads, num_part, num_part))

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, query_dim),
            nn.Dropout(dropout)
        )

    def forward(self, x, context=None, mask=None):
        h = self.heads

        q = self.to_q(x)
        context = default(context, x)
        k = self.to_k(context)
        v = self.to_v(context)

        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), (q, k, v))

        sim = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
        
        if self.rel_pe:
            sim = sim + self.relative_position_bias_table.unsqueeze(0)

        if exists(mask):
            mask = rearrange(mask, 'b ... -> b (...)')
            max_neg_value = -torch.finfo(sim.dtype).max
            mask = repeat(mask, 'b j -> (b h) () j', h=h)
            sim.masked_fill_(~mask, max_neg_value)

        # attention, what we cannot get enough of
        attn = sim.softmax(dim=-1)

        out = einsum('b h i j, b h j d -> b h i d', attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)', h=h)
        return self.to_out(out)


class BasicTransformerBlock(nn.Module):
    def __init__(self, dim, n_heads, d_head, dropout=0., emb_channels=128, use_scale_shift_norm=True, context_dim=None, gated_ff=True, rel_pe=False, num_part=4, include_anchor_pe=False):
        super().__init__()
        self.use_scale_shift_norm=use_scale_shift_norm
        
        self.emb_layers = nn.Sequential(
            nn.SiLU(),
            nn.Linear(
                emb_channels,
                2 * dim if use_scale_shift_norm else dim,
            ),
        )
        self.out_layers = nn.Sequential(
            Normalize(dim),
            nn.SiLU(),
            nn.Dropout(p=dropout),
            zero_module(
                nn.Conv1d(dim, dim, 1)
            ),
        )
        self.include_anchor_pe=include_anchor_pe
        if self.include_anchor_pe:
            self.pe_encoder = SinusoidalEmbedding3D(emb_channels)
            self.pe_emb_layers = nn.Sequential(
                nn.SiLU(),
                nn.Linear(
                    emb_channels,
                    2 * dim if use_scale_shift_norm else dim,
                )
            )
            self.pe_out_layers = nn.Sequential(
                Normalize(dim),
                nn.SiLU(),
                nn.Dropout(p=dropout),
                zero_module(
                    nn.Conv1d(dim, dim, 1)
                ),
            )
        self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout, rel_pe=rel_pe, num_part=num_part)  # is a self-attention
        self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
        self.attn2 = CrossAttention(query_dim=dim, context_dim=context_dim,
                                    heads=n_heads, dim_head=d_head, dropout=dropout, rel_pe=rel_pe, num_part=num_part)  # is self-attn if context is none
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.norm3 = nn.LayerNorm(dim)

    def forward(self, x, emb, context=None, anchors=None):
        if self.include_anchor_pe:
            pe_embedding = self.emb_layers(self.pe_encoder(anchors.transpose(1,2))).transpose(1,2)
        else:
            pe_embedding = 0.
        emb_out = self.emb_layers(emb).type(x.dtype)
        while len(emb_out.shape) < len(x.shape):
            emb_out = emb_out[..., None]
        # print(pe_embedding.shape)
        emb_out = emb_out + pe_embedding
        if self.use_scale_shift_norm:
            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
            scale, shift = torch.chunk(emb_out, 2, dim=1)
            h = out_norm(x) * (1 + scale) + shift
            h = out_rest(h)
        else:
            h = x + emb_out
            h = self.out_layers(h)
        x = x + h
        x = x.transpose(1,2)
        if exists(context):
            context = context.transpose(1,2)
        x = self.attn1(self.norm1(x)) + x
        x = self.attn2(self.norm2(x), context=context) + x
        x = self.ff(self.norm3(x)) + x
        x = x.transpose(1,2)
        return x

@NETS.register_module()
class UNet(nn.Module):
    """
    Transformer block for image-like data.
    First, project the input (aka embedding)
    and reshape to b, t, d.
    Then apply standard transformer action.
    Finally, reshape to image
    """
    def __init__(self, in_channels, n_heads, d_head,
                 depth=1, embed_dim=128, use_scale_shift_norm=True, dropout=0., prior_dim=512, language_dim=64, gated_ff=True, rel_pe=False, num_part=4, include_anchor_pe=False):
        super().__init__()
        self.in_channels = in_channels
        self.embed_dim=embed_dim
        inner_dim = n_heads * d_head
        self.in_layers = nn.Sequential(
            Normalize(in_channels),
            nn.SiLU(),
            nn.Conv1d(in_channels, inner_dim, 1),
        )
        self.time_embed = nn.Sequential(
            nn.Linear(embed_dim, embed_dim),
            nn.SiLU(),
            nn.Linear(embed_dim, embed_dim),
        )
        self.transformer_blocks_prior = nn.ModuleList(
            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=prior_dim, use_scale_shift_norm=use_scale_shift_norm, gated_ff=gated_ff, emb_channels=embed_dim, rel_pe=rel_pe, num_part=num_part, include_anchor_pe=include_anchor_pe)
                for d in range(depth)]
        )
        self.transformer_blocks_language = nn.ModuleList(
            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=language_dim, use_scale_shift_norm=use_scale_shift_norm, gated_ff=gated_ff, emb_channels=embed_dim, rel_pe=rel_pe, num_part=num_part, include_anchor_pe=include_anchor_pe)
                for d in range(depth)]
        )

        self.proj_out = zero_module(nn.Conv1d(inner_dim,
                                              in_channels, 1))

    def forward(self, x, t, prior=None, language_code=None, part_indicator=None, anchors=None):
        # note: if no context is given, cross-attention defaults to self-attention
        x_in = x
        x = self.in_layers(x)
        t_emb = timestep_embedding(t, self.embed_dim, repeat_only=False)
        emb = self.time_embed(t_emb)
        for i in range(len(self.transformer_blocks_prior)):
            x = self.transformer_blocks_prior[i](x, emb, context=prior, anchors=anchors)
            x = self.transformer_blocks_language[i](x, emb, context=language_code, anchors=anchors)
        x = self.proj_out(x)
        return x + x_in 
python/difffacto/models/diffusions/nets/unet.py
diffFacto/diffFacto
import torch 
from torch.nn import Module, Linear, Conv1d
import numpy as np
import math
from einops import repeat

def timestep_embedding(timesteps, dim, max_period=10000):
    """
    Create sinusoidal timestep embeddings.
    :param timesteps: a 1-D Tensor of N indices, one per batch element.
                      These may be fractional.
    :param dim: the dimension of the output.
    :param max_period: controls the minimum frequency of the embeddings.
    :return: an [N x dim] Tensor of positional embeddings.
    """
    half = dim // 2
    freqs = torch.exp(
        -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
    ).to(device=timesteps.device)
    args = timesteps[:, None].to(timesteps.dtype) * freqs[None]
    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
    if dim % 2:
        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
    return embedding


class ConcatSquashLinear(Module):
    def __init__(self, dim_in, dim_out, dim_ctx):
        super(ConcatSquashLinear, self).__init__()
        self._layer = Linear(dim_in, dim_out)
        self._hyper_bias = Linear(dim_ctx, dim_out, bias=False)
        self._hyper_gate = Linear(dim_ctx, dim_out)

    def forward(self, ctx, x):
        gate = torch.sigmoid(self._hyper_gate(ctx))
        bias = self._hyper_bias(ctx)
        # if x.dim() == 3:
        #     gate = gate.unsqueeze(1)
        #     bias = bias.unsqueeze(1)
        ret = self._layer(x) * gate + bias
        return ret
    

class SineLayer(Module):
    """
    From the siren repository
    https://colab.research.google.com/github/vsitzmann/siren/blob/master/explore_siren.ipynb
    """
    def __init__(self, in_features, out_features, bias=True,
                 is_first=False, omega_0=30, use_linear=False):
        super().__init__()
        self.omega_0 = omega_0
        self.is_first = is_first

        self.in_features = in_features
        self.linear = Linear(in_features, out_features, bias=bias) if use_linear else Conv1d(in_features, out_features, 1, bias=bias)
        self.output_channels = out_features
        self.init_weights()

    def init_weights(self):
        with torch.no_grad():
            if self.is_first:
                self.linear.weight.uniform_(-1 / self.in_features,
                                            1 / self.in_features)
            else:
                self.linear.weight.uniform_(-np.sqrt(6 / self.in_features) / self.omega_0,
                                            np.sqrt(6 / self.in_features) / self.omega_0)

    def forward(self, input):
        return torch.sin(self.omega_0 * self.linear(input))
python/difffacto/models/diffusions/nets/utils.py
diffFacto/diffFacto
import torch 
from torch.nn import Module
from difffacto.utils.registry import DIFFUSIONS, NETS, build_from_cfg
from difffacto.utils.constants import ModelMeanType, ModelVarType, LossType, model_mean_dict, model_var_dict
from difffacto.utils.misc import *
from difffacto.utils.constants import *
import math
import numpy as np
from .diffusion_utils import extract_into_tensor, betas_for_alpha_bar, normal_kl






@DIFFUSIONS.register_module()
class PointDiffusion(Module):

    def __init__(
        self, 
        net, 
        num_timesteps, 
        beta_1, 
        beta_T, 
        mode='linear', 
        use_beta=True, 
        rescale_timesteps=False, 
        loss_type='mse', 
        model_mean_type='epsilon', 
        model_var_type='fixed_small', 
        include_global_latent=False, 
        include_anchor_latent=True,
        scaled_loss=False,
        include_anchors=False
    ):
        super().__init__()
        assert mode in ('linear', 'cosine')
        self.model = build_from_cfg(net, NETS)
        self.model_mean_type = model_mean_dict[model_mean_type]
        self.include_global_latent=include_global_latent
        self.include_anchor_latent=include_anchor_latent
        self.include_anchors=include_anchors
        self.model_var_type = model_var_dict[model_var_type]
        self.loss_type = model_loss_dict[loss_type]
        self.beta_1 = beta_1
        self.beta_T = beta_T
        self.scaled_loss=scaled_loss
        self.use_beta = use_beta
        self.rescale_timesteps = rescale_timesteps
        if mode == 'linear':
            betas = np.linspace(beta_1, beta_T, num=num_timesteps, dtype=np.float64)
        elif mode == 'cosine':
            betas = betas_for_alpha_bar(num_timesteps, lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,)
            
        self.betas = betas
        assert len(betas.shape) == 1, "betas must be 1-D"
        assert (betas > 0).all() and (betas <= 1).all()

        self.num_timesteps = int(num_timesteps)

        alphas = 1.0 - betas
        self.alphas_cumprod = np.cumprod(alphas, axis=0)
        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)

        # calculations for diffusion q(x_t | x_{t-1}) and others
        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)

        # calculations for posterior q(x_{t-1} | x_t, x_0)
        self.posterior_variance = (
            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )
        # log calculation clipped because the posterior variance is 0 at the
        # beginning of the diffusion chain.
        self.posterior_log_variance_clipped = np.log(
            np.append(self.posterior_variance[1], self.posterior_variance[1:])
        )
        self.posterior_mean_coef1 = (
            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )
        self.posterior_mean_coef2 = (
            (1.0 - self.alphas_cumprod_prev)
            * np.sqrt(alphas)
            / (1.0 - self.alphas_cumprod)
        )        
        
    def q_mean_variance(self, x_start, t):
        """
        Get the distribution q(x_t | x_0).
        :param x_0: the [N x 3 x n] tensor of noiseless inputs.
        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
        :param x_0: the [N x 3 x n] tensor of anchors the tensor of shape as x_0 to which x_0's diffuses
        :return: A tuple (mean, variance, log_variance), all of x_start's shape.
        """
        mean = (
            extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
        )
        variance = extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
        log_variance = extract_into_tensor(
            self.log_one_minus_alphas_cumprod, x_start.shape
        )
        return mean, variance, log_variance
    
    def q_sample(self, x_start, t, noise=None):
        """
        Diffuse the data for a given number of diffusion steps.
        In other words, sample from q(x_t | x_0).
        :param x_0: the initial data batch.
        :param anchors: anchors for each point in x_0
        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
        :param noise: if specified, the split-out normal noise.
        :return: A noisy version of x_0.
        """
        if noise is None:
            noise = torch.randn_like(x_start)
        assert noise.shape == x_start.shape
        return (
            extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
            + extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
            * noise
        )
        
    def q_posterior_mean(self, x_start, x_t, t):
        """
        Compute the mean and variance of the diffusion posterior:
            q(x_{t-1} | x_t, x_0)
        """
        assert x_start.shape == x_t.shape
        posterior_mean = (
            extract_into_tensor(self.posterior_mean_coef1, t, x_start.shape) * x_start
            + extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t 
        )
        assert (
            posterior_mean.shape[0]
            == x_t.shape[0]
        )
        return posterior_mean
    
    def q_posterior_variance(self, x_start, x_t, t):
        """
        Compute the mean and variance of the diffusion posterior:
            q(x_{t-1} | x_t, x_0)
        """
        assert x_start.shape == x_t.shape
        
        posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape)
        posterior_log_variance_clipped = extract_into_tensor(
            self.posterior_log_variance_clipped, t, x_t.shape
        )
        assert (
            posterior_variance.shape[0]
            == posterior_log_variance_clipped.shape[0]
            == x_t.shape[0]
        )
        return posterior_variance, posterior_log_variance_clipped
    
    def q_posterior_mean_variance(self, x_start, x_t, t):
        """
        Compute the mean and variance of the diffusion posterior:
            q(x_{t-1} | x_t, x_0)
        """
        assert x_start.shape == x_t.shape
    
        
        posterior_mean = self.q_posterior_mean(x_start, x_t, t)
        posterior_variance, posterior_log_variance_clipped = self.q_posterior_variance(x_start, x_t, t)
        return posterior_mean, posterior_variance, posterior_log_variance_clipped
    
    
    def p_mean_variance(
        self, x, t, anchors=None, code=None, pointwise_latent=None
    ):
        """
        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
        the initial x, x_0.
        :param model: the model, which takes a signal and a batch of timesteps
                      as input.
        :param x: the [N x C x ...] tensor at time t.
        :param t: a 1-D Tensor of timesteps.
        :param clip_denoised: if True, clip the denoised signal into [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample. Applies before
            clip_denoised.
        :return: a dict with the following keys:
                 - 'mean': the model mean output.
                 - 'variance': the model variance output.
                 - 'log_variance': the log of 'variance'.
                 - 'pred_xstart': the prediction for x_0.
        """

        B, C = x.shape[:2]
        assert t.shape == (B,)
        if self.include_anchors:
            input = torch.cat([x, anchors], dim=1)
        else:
            input = x
        if self.include_anchor_latent:
            assert pointwise_latent is not None 
            ctx = pointwise_latent
        elif self.include_global_latent:
            assert code is not None
            ctx = code 
        else:
            ctx = None
        model_output = self.model(input, self._scale_timesteps(t), ctx)

        model_variance, model_log_variance = {
            # for fixedlarge, we set the initial (log-)variance like so
            # to get a better decoder log likelihood.
            ModelVarType.FIXED_LARGE: (
                np.append(self.posterior_variance[1], self.betas[1:]),
                np.log(np.append(self.posterior_variance[1], self.betas[1:])),
            ),
            ModelVarType.FIXED_SMALL: (
                self.posterior_variance,
                self.posterior_log_variance_clipped,
            ),
        }[self.model_var_type]
        model_variance = extract_into_tensor(model_variance, t, x.shape)
        model_log_variance = extract_into_tensor(model_log_variance, t, x.shape)

        def process_xstart(x):
            return x

        if self.model_mean_type == ModelMeanType.PREVIOUS_X:
            pred_xstart = process_xstart(
                self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output)
            )
            model_mean = model_output
        elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]:
            if self.model_mean_type == ModelMeanType.START_X:
                pred_xstart = process_xstart(model_output)
            else:
                pred_xstart = process_xstart(
                    self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
                )
            model_mean, _, _ = self.q_posterior_mean_variance(
                x_start=pred_xstart, x_t=x, t=t
            )
        else:
            raise NotImplementedError(self.model_mean_type)

        assert (
            model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
        )
        return {
            "mean": model_mean,
            "variance": model_variance,
            "log_variance": model_log_variance,
            "pred_xstart": pred_xstart,
        }
        
    def _predict_xstart_from_eps(self, x_t, t, eps):
        assert x_t.shape == eps.shape
        
        return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
                - extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, eps.shape) * eps
                )

    def _predict_xstart_from_xprev(self, x_t, t, xprev):
        assert x_t.shape == xprev.shape
        return (  # (xprev - coef2*x_t - coef3 * anchors) / coef1
            extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev
            - extract_into_tensor(
                self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape
            ) * x_t
        )
    
    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
        return (
            extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
            - pred_xstart
        ) / extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
        
    def _scale_timesteps(self, t):
        if self.rescale_timesteps:
            return t.float() * (1000.0 / self.num_timesteps)
        if self.use_beta:
            return extract_into_tensor(self.betas, t, t.shape)
        return t

    def p_sample(
        self, x, t, anchors=None, code=None,pointwise_latent=None
    ):
        """
        Sample x_{t-1} from the model at the given timestep.
        :param model: the model to sample from.
        :param x: the current tensor at x_{t-1}.
        :param t: the value of t, starting at 0 for the first diffusion step.
        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample.
        :return: a dict containing the following keys:
                 - 'sample': a random sample from the model.
                 - 'pred_xstart': a prediction of x_0.
        """
        out = self.p_mean_variance(
            x,
            t,
            anchors=anchors,
            code=code,
            pointwise_latent=pointwise_latent
        )
        noise = torch.randn_like(x)
        nonzero_mask = (
            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
        )  # no noise when t == 0
        sample = out["mean"] + nonzero_mask * torch.exp(0.5 * out["log_variance"]) * noise
        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
    
    def p_sample_loop(
        self,
        shape,
        anchors=None,
        code=None,
        pointwise_latent=None,
        noise=None,
        device=None,
        progress=False,
    ):
        """
        Generate samples from the model.
        :param model: the model module.
        :param shape: the shape of the samples, (N, C, H, W).
        :param noise: if specified, the noise from the encoder to sample.
                      Should be of the same shape as `shape`.
        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
        :param denoised_fn: if not None, a function which applies to the
            x_start prediction before it is used to sample.
        :param device: if specified, the device to create the samples on.
                       If not specified, use a model parameter's device.
        :param progress: if True, show a tqdm progress bar.
        :return: a non-differentiable batch of samples.
        """
        final = None
        for t, sample in self.p_sample_loop_progressive(
            shape,
            noise=noise,
            anchors=anchors,
            code=code,
            pointwise_latent=pointwise_latent,
            device=device,
            progress=progress,
        ):
            final = sample
        return final["sample"]

    def p_sample_loop_progressive(
        self,
        shape,
        anchors=None,
        code=None,
        pointwise_latent=None,
        noise=None,
        device=None,
        progress=False,
    ):
        """
        Generate samples from the model and yield intermediate samples from
        each timestep of diffusion.
        Arguments are the same as p_sample_loop().
        Returns a generator over dicts, where each dict is the return value of
        p_sample().
        """
        if device is None:
            device = next(self.model.parameters()).device
        assert isinstance(shape, (tuple, list))
        if noise is not None:
            pcd = noise
        else:
            pcd = torch.randn(*shape, device=device)
        indices = list(range(self.num_timesteps))[::-1]

        if progress:
            # Lazy import so that we don't depend on tqdm.
            from tqdm.auto import tqdm

            indices = tqdm(indices)

        for i in indices:
            t = torch.tensor([i] * shape[0], device=device)
            with torch.no_grad():
                out = self.p_sample(
                    pcd,
                    t,
                    anchors=anchors,
                    code=code,
                    pointwise_latent=pointwise_latent
                )
                yield i, out
                pcd = out["sample"]
                
                
    def _vb_terms_bpd(
        self, x_start, x_t, t, code, model_output=None
    ):
        """
        Get a term for the variational lower-bound.
        This allows for comparison to other papers.
        :return: a dict with the following keys:
                 - 'output': a shape [N] tensor of NLLs or KLs.
                 - 'pred_xstart': the x_0 predictions.
        """
        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
            x_start=x_start, x_t=x_t, t=t
        )
        out = self.p_mean_variance(
            x_t, t, code, model_output=model_output
        )
        kl = normal_kl(
            true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
        )
        kl = mean_flat(kl) / np.log(2.0)

        decoder_nll = -gaussian_log_likelihood(
            x_start, mean=out["mean"], logvar=out["log_variance"]
        )
        assert decoder_nll.shape == x_start.shape
        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)

        # At the first timestep return the decoder NLL,
        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
        output = torch.where((t == 0), decoder_nll, kl)
        return {"output": output, "pred_xstart": out["pred_xstart"]}
    """
    def training_losses(self, x_start, t, code=None, pointwise_latent=None, noise=None):
        Compute training losses for a single timestep.
        :param model: the model to evaluate loss on.
        :param x_start: the [B x 3 x N] tensor of inputs.
        :param t: a batch of timestep indices.
        :param noise: if specified, the specific Gaussian noise to try to remove.
        :return: a dict with the key "loss" containing a tensor of shape [N].
                 Some mean or variance settings may also have other keys.
        if noise is None:
            noise = torch.randn_like(x_start)
        x_t = self.q_sample(x_start, t, noise=noise)
        if self.include_anchor_latent:
            assert pointwise_latent is not None 
            ctx = pointwise_latent
        elif self.include_global_latent:
            assert code is not None
            ctx = code 
        else:
            ctx = None

        terms = {}

        if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:
            terms["loss"] = self._vb_terms_bpd(
                x_start=x_start,
                x_t=x_t,
                t=t,
                code=code,
                clip_denoised=False,
            )["output"]
            if self.loss_type == LossType.RESCALED_KL:
                terms["loss"] *= self.num_timesteps
        elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:
            model_output = self.model(x_t, self._scale_timesteps(t), ctx)

            target = {
                ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(
                    x_start=x_start, x_t=x_t, t=t
                )[0],
                ModelMeanType.START_X: x_start,
                ModelMeanType.EPSILON: noise,
            }[self.model_mean_type]
            assert model_output.shape == target.shape == x_start.shape
            terms["mse"] = mean_flat((target - model_output) ** 2)
            if "vb" in terms:
                terms["recons_loss"] = terms["mse"] + terms["vb"]
            else:
                terms["recons_loss"] = terms["mse"] if not self.scaled_loss else extract_into_tensor( self.betas * self.betas / (2 * (1. - self.betas) * (1. - self.alphas_cumprod)), t, terms["mse"].shape) * terms["mse"]
        else:
            raise NotImplementedError(self.loss_type)
        
        return terms
    """
    def training_losses(self, x_start, t, anchors=None, code=None, pointwise_latent=None, noise=None):
        """
        Compute training losses for a single timestep.
        :param model: the model to evaluate loss on.
        :param x_start: the [B x 3 x N] tensor of inputs.
        :param t: a batch of timestep indices.
        :param noise: if specified, the specific Gaussian noise to try to remove.
        :return: a dict with the key "loss" containing a tensor of shape [N].
                 Some mean or variance settings may also have other keys.
        """
        if noise is None:
            noise = torch.randn_like(x_start)
        x_t = self.q_sample(x_start, t, noise=noise)
        if self.include_anchors:
            input = torch.cat([x_t, anchors], dim=1)
        else:
            input = x_t
        if self.include_anchor_latent:
            assert pointwise_latent is not None 
            ctx = pointwise_latent
        elif self.include_global_latent:
            assert code is not None
            ctx = code 
        else:
            ctx = None

        model_output = self.model(input, self._scale_timesteps(t), ctx)

        target = noise
        assert model_output.shape == target.shape == x_start.shape
        mse_loss = mean_flat((target - model_output) ** 2).mean()
        if self.scaled_loss:
            mse_loss = mse_loss * extract_into_tensor( self.betas * self.betas / (2 * (1. - self.betas) * (1. - self.alphas_cumprod)), t, mse_loss.shape)
        return dict(diffusion_loss=mse_loss)
python/difffacto/models/diffusions/point_diffusion.py
diffFacto/diffFacto
import torch
import torch.nn as nn
import torch.nn.functional as F
from difffacto.utils.registry import DISCRIMINATORS

@DISCRIMINATORS.register_module()
class Discriminator(nn.Module):

    def __init__(self, inp_dim, use_bn, use_ln, use_sigmoid, dims):
        super().__init__()

        self.inp_dim = inp_dim
        self.use_bn = use_bn
        self.use_ln = use_ln
        self.use_sigmoid = use_sigmoid
        self.dims = dims

        curr_dim = self.inp_dim
        self.layers = []
        self.bns = []
        self.lns = []
        for hid in self.dims:
            self.layers.append(nn.Linear(curr_dim, hid))
            if self.use_bn:
                self.bns.append(nn.BatchNorm1d(hid))
            else:
                self.bns.append(None)
            if self.use_ln:
                self.lns.append(nn.LayerNorm(hid))
            else:
                self.lns.append(None)
            curr_dim = hid
        self.layers = nn.ModuleList(self.layers)
        self.bns = nn.ModuleList(self.bns)
        self.lns = nn.ModuleList(self.lns)
        self.out = nn.Linear(curr_dim, 1)

    def forward(self, z=None, bs=None, return_all=False):
        if z is None:
            assert bs is not None
            z = torch.randn(bs, self.inp_dim).cuda()

        y = z
        for layer, bn, ln in zip(self.layers, self.bns, self.lns):
            y = layer(y)
            if self.use_bn:
                y = bn(y)
            if self.use_ln:
                y = ln(y)
            y = F.leaky_relu(y, 0.2)
        y = self.out(y)

        if self.use_sigmoid:
            y = torch.sigmoid(y)
        if return_all:
            return {
                'x': y
            }
        else:
            return y


@DISCRIMINATORS.register_module()
class DiscriminatorcGan(nn.Module):

    def __init__(self, inp_dim, ctx_dim, inner_dim, n_class):
        super().__init__()

        self.proj_in = nn.Conv1d(inp_dim, inner_dim, 1)
        self.n_class=n_class
        self.seq = nn.Sequential(
            nn.Linear(inner_dim*n_class + ctx_dim*n_class, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 128),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(128, 1),
        )

    def forward(self, z=None, ctx=None, mask=None):
        bs = z.shape[0]
        z = self.proj_in(z)
        if mask is not None:
            z = z * mask.unsqueeze(1)
            ctx = ctx * mask.unsqueeze(1)
        inp = torch.cat([z, ctx], dim=1).transpose(1,2).reshape(bs, -1)
        y = self.seq(inp)

        return y

python/difffacto/models/discriminators/discriminator.py
diffFacto/diffFacto
import torch
import torch.nn as nn
import torch.nn.functional as F
from difffacto.utils.registry import ENCODERS

@ENCODERS.register_module()
class PCN(nn.Module):
    """
    "PCN: Point Cloud Completion Network"
    (https://arxiv.org/pdf/1808.00671.pdf)
    Attributes:
        num_dense:  16384
        latent_dim: 1024
        grid_size:  4
        num_coarse: 1024
    """

    def __init__(
        self, 
        point_dim=3, 
        part_latent_dim=512, 
        num_anchors=4,
        normalize_latent=False
        ):
        super().__init__()

        self.point_dim=point_dim
        self.normalize_latent=normalize_latent
        self.num_anchors=num_anchors
        self.part_latent_dim=part_latent_dim

        self.first_conv = nn.Sequential(
            nn.Conv1d(point_dim, 128, 1),
            nn.BatchNorm1d(128),
            nn.ReLU(inplace=True),
            nn.Conv1d(128, 256, 1)
        )

        self.second_conv = nn.Sequential(
            nn.Conv1d(512, 512, 1),
            nn.BatchNorm1d(512),
            nn.ReLU(inplace=True),
            nn.Conv1d(512, 1024, 1)
        )

        self.latent_mlp = nn.Sequential(
            nn.Linear(1024, 1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, 1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, part_latent_dim * num_anchors)
        )

    def forward(self, xyz=None):
        B, N, C = xyz.shape
        assert C == self.point_dim
        
        # encoder
        feature = self.first_conv(xyz.transpose(2, 1))                                       # (B,  256, N)
        feature_global = torch.max(feature, dim=2, keepdim=True)[0]                          # (B,  256, 1)
        feature = torch.cat([feature_global.expand(-1, -1, N), feature], dim=1)              # (B,  512, N)
        feature = self.second_conv(feature)                                                  # (B, 1024, N)
        feature_global = torch.max(feature,dim=2,keepdim=False)[0]                           # (B, 1024)
        part_latent = self.latent_mlp(feature_global).reshape(B, self.num_anchors, -1)
        if self.normalize_latent:
            part_latent = F.normalize(part_latent.reshape(B*self.num_anchors, -1)).reshape(B, self.num_anchors, -1)
        return part_latent

python/difffacto/models/encoders/PCN.py
diffFacto/diffFacto
from .PCN import PCN
from .pnet import Pnet2Stage
from .pointnet import PointNet, PointNetV2, PointNetV3, PointNetV2VAE, PointNetVAE
from .super_seg_encoder import SupSegsEncoder, PartglotSupSegsEncoderWithCBN
from .language_encoders import LSTM, MultiHeadCrossAttention, PartLanguageSelector, LatentLanguageFuser
from .mean_variance_regressor import MeanVarianceRegressor
from .pointnet2 import PointNet2MSG, PointNet2SSG
from .part_encoders import PartAligner, PartEncoder
from .flow import build_latent_flow
python/difffacto/models/encoders/__init__.py
diffFacto/diffFacto
import types
import torch
import torch.nn as nn
import torch.nn.functional as F


class CouplingLayer(nn.Module):

    def __init__(self, d, intermediate_dim, swap=False):
        nn.Module.__init__(self)
        self.d = d - (d // 2)
        self.swap = swap
        self.net_s_t = nn.Sequential(
            nn.Linear(self.d, intermediate_dim),
            nn.ReLU(inplace=True),
            nn.Linear(intermediate_dim, intermediate_dim),
            nn.ReLU(inplace=True),
            nn.Linear(intermediate_dim, (d - self.d) * 2),
        )

    def forward(self, x, logpx=None, reverse=False):

        if self.swap:
            x = torch.cat([x[:, self.d:], x[:, :self.d]], 1)

        in_dim = self.d
        out_dim = x.shape[1] - self.d

        s_t = self.net_s_t(x[:, :in_dim])
        scale = torch.sigmoid(s_t[:, :out_dim] + 2.)
        shift = s_t[:, out_dim:]

        logdetjac = torch.sum(torch.log(scale).view(scale.shape[0], -1), 1, keepdim=True)

        if not reverse:
            y1 = x[:, self.d:] * scale + shift
            delta_logp = -logdetjac
        else:
            y1 = (x[:, self.d:] - shift) / scale
            delta_logp = logdetjac

        y = torch.cat([x[:, :self.d], y1], 1) if not self.swap else torch.cat([y1, x[:, :self.d]], 1)

        if logpx is None:
            return y
        else:
            return y, logpx + delta_logp


class SequentialFlow(nn.Module):
    """A generalized nn.Sequential container for normalizing flows.
    """

    def __init__(self, layersList):
        super(SequentialFlow, self).__init__()
        self.chain = nn.ModuleList(layersList)

    def forward(self, x, logpx=None, reverse=False, inds=None):
        if inds is None:
            if reverse:
                inds = range(len(self.chain) - 1, -1, -1)
            else:
                inds = range(len(self.chain))

        if logpx is None:
            for i in inds:
                x = self.chain[i](x, reverse=reverse)
            return x
        else:
            for i in inds:
                x, logpx = self.chain[i](x, logpx, reverse=reverse)
            return x, logpx


def build_latent_flow(latent_flow_depth, latent_flow_hidden_dim, latent_dim):
    chain = []
    for i in range(latent_flow_depth):
        chain.append(CouplingLayer(latent_dim, latent_flow_hidden_dim, swap=(i % 2 == 0)))
    return SequentialFlow(chain)
python/difffacto/models/encoders/flow.py
diffFacto/diffFacto
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
from difffacto.utils.registry import ENCODERS
from pointnet2_ops.pointnet2_utils import gather_operation
import math

@ENCODERS.register_module()
class LSTM(nn.Module):
    def __init__(self, text_dim=64, embedding_dim=100, vocab_size=2787, padding_idx=0):
        super().__init__()
        self.text_dim=text_dim
        self.padding_idx = padding_idx
        self.word_embedding = nn.Embedding(
            vocab_size, embedding_dim, padding_idx=padding_idx
        )
        self.rnn = nn.LSTM(embedding_dim, text_dim, batch_first=True)
        self.w_attn = nn.Parameter(torch.Tensor(1, text_dim))
        nn.init.xavier_uniform_(self.w_attn)

    def forward(self, padded_tokens, dropout=0.5):
        
        w_emb = self.word_embedding(padded_tokens)
        w_emb = F.dropout(w_emb, dropout, self.training)
        len_seq = (padded_tokens != self.padding_idx).sum(dim=1).cpu()
        x_packed = pack_padded_sequence(
            w_emb, len_seq, enforce_sorted=False, batch_first=True
        )

        B = padded_tokens.shape[0]

        rnn_out, _ = self.rnn(x_packed)
        rnn_out, dummy = pad_packed_sequence(rnn_out, batch_first=True)
        h = rnn_out[torch.arange(B), len_seq - 1]
        final_feat, attn = self.word_attention(rnn_out, h, len_seq)
        return final_feat, attn
        
        
        # w_emb = self.word_embedding(padded_tokens)
        # w_emb = F.dropout(w_emb, dropout, self.training)
        # len_seq = (padded_tokens != self.padding_idx).sum(dim=1).cpu()
        # x_packed = pack_padded_sequence(
        #     w_emb, len_seq, enforce_sorted=False, batch_first=True
        # )

        # B = padded_tokens.shape[0]

        # rnn_out, _ = self.rnn(x_packed)
        # rnn_out, dummy = pad_packed_sequence(rnn_out, batch_first=True)
        # h = rnn_out[torch.arange(B), len_seq - 1]
        # final_feat, attn = self.word_attention(rnn_out, h, len_seq)
        # if self.regress_part:
        #     regress_first, regress_rest = self.part_regressor[0], self.part_regressor[1:]
        #     part_f_in = regress_first(final_feat).reshape(B, -1, self.num_part)
        #     part_f_out = regress_rest(part_f_in)
        # else:
        #     part_f_out = final_feat.unsqueeze(-1)
        # if self.return_zero:
        #     part_f_out = torch.zeros_like(part_f_out)
        # return part_f_out.transpose(1,2), attn

    def word_attention(self, R, h, len_seq):
        """
        Input:
            R: hidden states of the entire words
            h: the final hidden state after processing the entire words
            len_seq: the length of the sequence
        Output:
            final_feat: the final feature after the bilinear attention
            attn: word attention weights
        """
        B, N, D = R.shape
        device = R.device
        len_seq = len_seq.to(device)

        W_attn = (self.w_attn * torch.eye(D).to(device))[None].repeat(B, 1, 1)
        score = torch.bmm(torch.bmm(R, W_attn), h.unsqueeze(-1))

        mask = torch.arange(N).reshape(1, N, 1).repeat(B, 1, 1).to(device)
        mask = mask < len_seq.reshape(B, 1, 1)

        score = score.masked_fill(mask == 0, -1e9)
        attn = F.softmax(score, 1)
        final_feat = torch.bmm(R.transpose(1, 2), attn).squeeze(-1)

        return final_feat, attn.squeeze(-1)

@ENCODERS.register_module()
class MultiHeadCrossAttention(nn.Module):
    def __init__(self, src_channel, tgt_channel, out_channel, qk_dim=512, nheads=8, qkv_bias=True, scale=True):
        super(MultiHeadCrossAttention, self).__init__()
        assert out_channel % nheads == 0
        feedforward_dim = out_channel * 2
        self.qk_ndim = qk_dim // nheads
        self.v_ndim = out_channel // nheads
        self.nheads=nheads
        self.Wq = nn.Linear(src_channel, qk_dim, bias=qkv_bias)
        self.Wk = nn.Linear(tgt_channel, qk_dim, bias=qkv_bias)
        self.Wv = nn.Linear(tgt_channel, out_channel, bias=qkv_bias)
        self.scale= 1 / math.sqrt(self.qk_ndim) if scale else 1. 

        self.dp1 = nn.Dropout(0.1)
        self.dp2 = nn.Dropout(0.1)

        self.linear1 = nn.Linear(out_channel, feedforward_dim)
        self.linear2 = nn.Linear(feedforward_dim, out_channel)

        self.norm = nn.LayerNorm(out_channel)

    def forward(self, src, tgt, part_indicator):
        B, N, C = src.shape
        _, M, D = tgt.shape
        Q, K, V = self.Wq(src).reshape(B, N, self.nheads, self.qk_ndim), self.Wk(tgt).reshape(B, M, self.nheads, self.qk_ndim), self.Wv(tgt).reshape(B, M, self.nheads, self.v_ndim)
        
        QK = torch.matmul(Q.permute([0, 2, 1, 3]), K.permute([0, 2, 3, 1])) * self.scale # [B, head, N, M]
        QK = F.softmax(QK, dim=3)
        QK_dp = self.dp1(QK)
        QK_slice = QK_dp * part_indicator.reshape(B, 1, -1, 1)
        QK_slice = torch.sum(QK_slice, 2, keepdim=True)
        QK_dp = QK_slice
        message = torch.matmul(QK_dp, V.permute([0, 2, 1, 3])).transpose(1,2).reshape(B, -1, self.nheads * self.v_ndim)
        message2 = self.linear2(F.relu(self.linear1(message)))
        message = message + self.dp2(message2)
        message = self.norm(message)
        
        return message
    
@ENCODERS.register_module()
class LatentLanguageFuser(nn.Module):
    def __init__(
        self, 
        text_dim, 
        part_dim, 
        num_part, 
        residual=True, 
        cat_weights=True, 
        mid_dim=1024,
        conditional=False,
        conditional_dim=24,
        normalize_latent=False,
        regress_weights=False,
        direct_add_text_to_part=False,
        concat_weight=False):
        super().__init__()
        self.residual = residual
        self.normalize_latent=normalize_latent
        self.direct_add_text_to_part=direct_add_text_to_part
        if self.direct_add_text_to_part:
            assert not regress_weights
        self.concat_weight=concat_weight
        self.conditional=conditional
        self.num_part = num_part
        self.conditional_dim=conditional_dim if self.conditional else 0
        self.cat_weights=cat_weights
        self.part_latent_fuser = nn.ModuleList([nn.Linear(part_dim* num_part, part_dim),
                                                nn.Linear(part_dim + self.num_part if cat_weights else part_dim, part_dim),
                                                nn.Linear(part_dim + text_dim + self.conditional_dim, mid_dim),
                                                nn.Linear(mid_dim, mid_dim),
                                                nn.Linear(mid_dim, part_dim)]
                                                )
        self.regress_weights=regress_weights
        if self.regress_weights:
            self.weight_regressor = nn.Sequential(nn.Linear(text_dim, 128), nn.ReLU(inplace=True),
                                                  nn.Linear(128, 128), nn.ReLU(inplace=True),
                                                  nn.Linear(128, self.num_part), nn.Sigmoid())
        if self.direct_add_text_to_part:
            input_dim = text_dim
            if self.concat_weight:
                input_dim += self.num_part
            if self.conditional:
                input_dim += conditional_dim
            self.text_regressor = nn.Sequential(nn.Linear(input_dim, 128), nn.ReLU(inplace=True),
                                                nn.Linear(128, 256), nn.ReLU(inplace=True),
                                                nn.Linear(256, 256), nn.ReLU(inplace=True),
                                                nn.Linear(256, part_dim))
         
    def forward(self, part_latent, part_indicators, text_latent, conditional=None):
        if self.conditional:
            assert conditional is not None and conditional.shape[-1] == self.conditional_dim
        part_id = part_indicators.max(1)[1].to(torch.int)
        B = part_latent.shape[0]
        if self.regress_weights:
            weights = self.weight_regressor(text_latent) #[B, num_part]
        else:
            weights = part_indicators
        if self.direct_add_text_to_part:
            input = torch.cat([text_latent, weights], dim=-1) if self.concat_weight else text_latent
            if self.conditional:
                K = conditional.shape[1]
                input = input.unsqueeze(1).repeat_interleave(K, dim=1)
            else:
                K = 1
                text_latent_out = self.text_regressor(input).unsqueeze(2).repeat_interleave(self.num_part, dim=2)
                
            out_latent = part_latent.unsqueeze(1) + weights.reshape(B, 1, self.num_part, 1) * text_latent_out.reshape(B, K, self.num_part, -1)
            if self.normalize_latent:
                out_latent = F.normalize(out_latent, dim=-1)
            return out_latent
                
        
            
        fused_latent = F.relu(self.part_latent_fuser[0](part_latent.reshape(B, -1)))
        
        if self.cat_weights:
            fused_latent = F.relu(self.part_latent_fuser[1](torch.cat([fused_latent, weights], dim=-1)))
        else:
            fused_latent = F.relu(self.part_latent_fuser[1](fused_latent))
        if self.conditional:
            K = conditional.shape[1]
            fused_latent = torch.cat([text_latent, fused_latent], dim=-1).unsqueeze(1).repeat_interleave(K, dim=1)
            fused_latent = torch.cat([fused_latent, conditional], dim=-1)
        else:
            K = 1
            fused_latent = torch.cat([text_latent, fused_latent], dim=-1).unsqueeze(1)
        
        fused_latent = F.relu(self.part_latent_fuser[2](fused_latent))
        fused_latent = self.part_latent_fuser[4](F.relu(self.part_latent_fuser[3](fused_latent)))
        if self.residual:
            fused_latent=part_latent.unsqueeze(1) + fused_latent
        out_latent = (1 - part_indicators).reshape(B, 1, self.num_part, 1) * part_latent.unsqueeze(1) + part_indicators.reshape(B, 1, self.num_part, 1) * fused_latent.unsqueeze(2)
        if self.normalize_latent:
            out_latent = F.normalize(out_latent, p=2, dim=-1)
        return out_latent
            
        
@ENCODERS.register_module()
class PartLanguageSelector(nn.Module):
    def __init__(self, text_dim, part_dim, out_channel):
        super(PartLanguageSelector, self).__init__()
        self.linears = nn.Sequential(
            nn.Linear(text_dim + part_dim, 512), nn.ReLU(inplace=True),
            nn.Linear(512, 1024), nn.ReLU(inplace=True),
            nn.Linear(1024, 1024), nn.ReLU(inplace=True),
            nn.Linear(1024, 512), nn.ReLU(inplace=True),
            nn.Linear(512, out_channel)
        )

        self.dp = nn.Dropout(0.2)
        self.norm = nn.LayerNorm(out_channel)

    def forward(self, src, tgt, part_indicator):
        B, N, C = src.shape
        _, M, D = tgt.shape
        part_id = part_indicator.max(1)[1].to(torch.int)
        tgt_part = gather_operation(tgt.transpose(1,2).contiguous(), part_id.reshape(B, 1))
        return tgt_part.transpose(1,2)

        
python/difffacto/models/encoders/language_encoders.py
diffFacto/diffFacto
import torch
import torch.nn.functional as F
from torch import nn
from difffacto.utils.registry import ENCODERS

@ENCODERS.register_module()
class MeanVarianceRegressor(nn.Module):
    def __init__(self, latent_dim, num_class):
        super().__init__()
        self.num_class = num_class
        # Mapping to [c], cmean
        self.fc1_m = nn.Linear(latent_dim, 256)
        self.fc2_m = nn.Linear(256, 128)
        self.fc3_m = nn.Linear(128, 256)
        self.fc4_m = nn.Linear(256, latent_dim)
        self.fc_bn1_m = nn.ModuleList([nn.BatchNorm1d(256) for _ in range(num_class)]) 
        self.fc_bn2_m = nn.ModuleList([nn.BatchNorm1d(128) for _ in range(num_class)]) 
        self.fc_bn3_m = nn.ModuleList([nn.BatchNorm1d(256) for _ in range(num_class)]) 
        

        # Mapping to [c], cmean
        self.fc1_v = nn.Linear(latent_dim, 256)
        self.fc2_v = nn.Linear(256, 128)
        self.fc3_v = nn.Linear(128, 256)
        self.fc4_v = nn.Linear(256, latent_dim)
        self.fc_bn1_v = nn.ModuleList([nn.BatchNorm1d(256) for _ in range(num_class)]) 
        self.fc_bn2_v = nn.ModuleList([nn.BatchNorm1d(128) for _ in range(num_class)]) 
        self.fc_bn3_v = nn.ModuleList([nn.BatchNorm1d(256) for _ in range(num_class)]) 
        

    def forward(self, x):
        m = F.relu(torch.stack([self.fc_bn1_m[i](self.fc1_m(x)[:, i]) for i in range(self.num_class)], dim=1))
        m = F.relu(torch.stack([self.fc_bn2_m[i](self.fc2_m(m)[:, i]) for i in range(self.num_class)], dim=1))
        m = F.relu(torch.stack([self.fc_bn3_m[i](self.fc3_m(m)[:, i]) for i in range(self.num_class)], dim=1))
        m = self.fc4_m(m)
        v = F.relu(torch.stack([self.fc_bn1_v[i](self.fc1_v(x)[:, i]) for i in range(self.num_class)], dim=1))
        v = F.relu(torch.stack([self.fc_bn2_v[i](self.fc2_v(v)[:, i]) for i in range(self.num_class)], dim=1))
        v = F.relu(torch.stack([self.fc_bn3_v[i](self.fc3_v(v)[:, i]) for i in range(self.num_class)], dim=1))
        v = self.fc4_v(v)

        # Returns both mean and logvariance, just ignore the latter in deteministic cases.
        return m, v
python/difffacto/models/encoders/mean_variance_regressor.py
diffFacto/diffFacto
import torch.nn as nn
import torch.nn.functional as F
import torch
import math
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
from functools import partial
from difffacto.utils.registry import ENCODERS, DECOMPOSERS, build_from_cfg
from difffacto.models.encoders.pointnet import PointNet, PointNetVAEBase
from pointnet2_ops.pointnet2_utils import gather_operation
from einops import rearrange, repeat
from difffacto.utils.misc import gaussian_log_likelihood, reparameterize_gaussian, gaussian_entropy, standard_normal_logprob
from difffacto.models.diffusions.nets.attention import FeedForward, BasicTransformerBlock, zero_module, MLP
from .flow import build_latent_flow
def init_linear(l, stddev):
    nn.init.normal_(l.weight, std=stddev)
    if l.bias is not None:
        nn.init.constant_(l.bias, 0.0)

@ENCODERS.register_module()
class PartAlignerTransformer(nn.Module):
    """
    Transformer block for image-like data.
    First, project the input (aka embedding)
    and reshape to b, t, d.
    Then apply standard transformer action.
    Finally, reshape to image
    NEW: use_linear for more efficiency instead of the 1x1 convs
    """
    def __init__(self, in_channels, n_heads, d_head, out_channels,
                 depth=1, dropout=0., use_linear=False, n_class=4,
                 use_checkpoint=False, single_attn=False, class_cond=True,
                 mask_out_unreferenced_code=True, cimle=False, noise_dim=32, noise_scale=10, cimle_start_epoch=0,
                 add_class_cond=False, cond_noise_type=0, cond_noise_as_token=False, 
                ):
        super().__init__()
        self.n_class=n_class
        self.cimle=cimle
        self.add_class_cond=add_class_cond
        self.noise_scale=noise_scale
        self.cond_noise_type=cond_noise_type 
        if cond_noise_as_token:
            self.cond_noise_type = 1
            print("cond noise as token Deprecated!! use cond_noise_type=1 instead")
        in_channels = in_channels + int(class_cond and not add_class_cond) * n_class + int(cimle and self.cond_noise_type in [0, 4]) * noise_dim
        self.in_channels = in_channels
        self.class_cond=class_cond
        self.mask_out_unreferenced_code=mask_out_unreferenced_code
        inner_dim = n_heads * d_head
        self.noise_dim=noise_dim if self.cond_noise_type!=1 else inner_dim
        self.cimle_start_epoch=cimle_start_epoch
        self.inner_dim = inner_dim
        if class_cond and add_class_cond:
            self.class_emb = nn.Embedding(n_class, inner_dim)
        self.pre_norm = nn.LayerNorm(inner_dim)
        self.post_norm = nn.LayerNorm(inner_dim)
        if self.cond_noise_type == 3:
            self.emb_layer = nn.Sequential(
                nn.SiLU(),
                nn.Linear(
                    self.noise_dim,
                    2 * inner_dim,
                ),
            )
        if self.cond_noise_type == 4:
            self.emb_layer = MLP([self.noise_dim, 4*self.noise_dim,4*self.noise_dim, self.noise_dim], use_linear=True)
        if not use_linear:
            self.proj_in = nn.Conv1d(in_channels,
                                     inner_dim,
                                     kernel_size=1,
                                     stride=1,
                                     padding=0)
        else:
            self.proj_in = nn.Linear(in_channels, inner_dim)

        self.transformer_blocks = nn.ModuleList(
            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=None,
                                   checkpoint=use_checkpoint, single_attn=single_attn, adaln=(self.cond_noise_type==2), y_dim=self.noise_dim)
                for d in range(depth)]
        )
        if not use_linear:
            self.proj_out = nn.Conv1d(inner_dim, out_channels, kernel_size=1, stride=1, padding=0)
        else:
            self.proj_out = nn.Linear(inner_dim, out_channels)
        self.use_linear = use_linear

    def forward(self, x: torch.Tensor, mask: torch.Tensor=None, noise: torch.Tensor=None):
        assert x.shape[-1] == self.n_class
        if self.class_cond and not self.add_class_cond:
            class_embed = torch.eye(self.n_class).to(x).unsqueeze(0).repeat_interleave(x.shape[0], dim=0) # B, n_ctx, n_ctx
            x = torch.cat([x, class_embed], dim=1)
        if self.cimle and self.cond_noise_type != 1:
            assert noise is not None
            if noise.shape[1] != self.noise_dim:
                noise = torch.zeros(x.shape[0], self.noise_dim).to(x)
            noise = noise * self.noise_scale
            if self.cond_noise_type == 4:
                noise = self.emb_layer(noise)
            if self.cond_noise_type in [0, 4]:
                x = torch.cat([x, noise.unsqueeze(-1).expand(-1, -1, self.n_class)], dim=1)
        assert self.in_channels == x.shape[1]
        h =  self._forward_attn(x, mask=mask if self.mask_out_unreferenced_code else None, noise=noise if self.cimle else None)
        mean, logvar = torch.split(h, 3, dim=1)
        return mean, logvar
    
    def _forward_attn(self, x, mask=None, noise=None):
        # note: if no context is given, cross-attention defaults to self-attention
        if not self.use_linear:
            x = self.proj_in(x)
        x = rearrange(x, 'b c n -> b n c').contiguous()
        if self.use_linear:
            x = self.proj_in(x)
        if self.class_cond and self.add_class_cond:
            class_emb = self.class_emb(torch.arange(self.n_class, device=x.device)).unsqueeze(0)
            x = x + class_emb
        if self.cimle:
            if self.cond_noise_type==1:
                new_mask = torch.zeros(x.shape[0], self.n_class + 1).to(mask)
                new_mask[:, 1:] = mask 
                if noise.shape[1] != self.noise_dim:
                    noise = torch.randn(x.shape[0], self.noise_dim).to(x)
                else:
                    new_mask[:, 0] = 1.0
                noise = noise * self.noise_scale
                x = torch.cat([noise.unsqueeze(1), x], dim=1)
                mask = new_mask
                x = self.pre_norm(x)
            elif self.cond_noise_type == 3:
                shift, scale = torch.chunk(self.emb_layer(noise), 2, dim=-1)
                x = self.pre_norm(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
        else:
            x = self.pre_norm(x)    
        
        for i, block in enumerate(self.transformer_blocks):
            x = block(x, context=None, mask=mask, y=noise if self.cond_noise_type==2 else None)
        x = self.post_norm(x)
        if self.cimle and self.cond_noise_type==1:
            x = x[:, 1:]
        if self.use_linear:
            x = self.proj_out(x)
        x = rearrange(x, 'b n c -> b c n').contiguous()
        if not self.use_linear:
            x = self.proj_out(x)
        return  x


class AdaInstanceNorm1d(nn.Module):
    def __init__(self, dim, ctx_dim):
        super().__init__()
        self.norm = nn.InstanceNorm1d(dim)
        self.proj = nn.Sequential(
            nn.Linear(ctx_dim, dim * 4), 
            nn.SiLU(),
            nn.Linear(dim * 4, 2* dim)
        )
        self.dim=dim
        with torch.no_grad():
            self.proj[0].weight.zero_()
            self.proj[0].bias.zero_()
            self.proj[2].weight.zero_()
            self.proj[2].bias.zero_()
    def forward(self, x, ctx):
        x = self.norm(x)
        shift, scale = torch.split(self.proj(ctx), self.dim, dim=1)
        return x * (1 + scale) + shift

@ENCODERS.register_module()
class PartAlignerCIMLE(nn.Module):
    def __init__(
        self,
        width: int,
        n_class: int,
        param_dim:int=3,
        noise_dim:int=32,
        norm: str='bn',
        noise_encoder_type: int=0
    ):
        super().__init__()
        self.width=width 
        self.n_class = n_class
        self.noise_dim=noise_dim
        self.norm = norm
        if norm == 'bn':
            enc_norm = nn.BatchNorm1d 
            norm_layer = nn.BatchNorm1d 
        elif norm == 'in':
            enc_norm = nn.InstanceNorm1d 
            norm_layer = nn.InstanceNorm1d 
        elif norm == 'adain':
            enc_norm = nn.InstanceNorm1d 
            norm_layer = partial(AdaInstanceNorm1d, ctx_dim=noise_dim)
        self.noise_encoder_type=noise_encoder_type
        if noise_encoder_type == 0:
            self.noise_encoder=nn.Sequential(
                nn.Linear(width*n_class + noise_dim, 512),
                enc_norm(512),
                nn.ReLU(False),
                nn.Linear(512, 128),
                enc_norm(128),
                nn.ReLU(False),
                nn.Linear(128, noise_dim),
            )
        self.fc1_m = nn.Linear(width*n_class+ noise_dim if norm != 'adain' else width*n_class, 512)
        self.fc2_m = nn.Linear(512, 256)
        self.fc3_m = nn.Linear(256, 64)
        self.fc4_m = nn.Linear(64, param_dim*n_class)
        self.fc_bn1_m = norm_layer(512)
        self.fc_bn2_m = norm_layer(256)
        self.fc_bn3_m = norm_layer(64)
        
        
        self.fc1_v = nn.Linear(width*n_class+ noise_dim if norm != 'adain' else width*n_class, 512)
        self.fc2_v = nn.Linear(512, 256)
        self.fc3_v = nn.Linear(256, 64)
        self.fc4_v = nn.Linear(64, param_dim*n_class)
        self.fc_bn1_v = norm_layer(512)
        self.fc_bn2_v = norm_layer(256)
        self.fc_bn3_v = norm_layer(64)
        for i in range(1, 5):
            for k in ['m', 'v']:
                init_linear(getattr(self, f"fc{i}_{k}"), 0.25)
    
    
    
    def forward(self, part_code, valid_id, noise):
        if noise.shape[1] != self.noise_dim:
            noise = torch.randn(part_code.shape[0], self.noise_dim).to(part_code.device)
        # print(noise.shape)
        assert noise.shape[1] == self.noise_dim
        assert part_code.shape[1:] == (self.width, self.n_class)
        part_code = part_code * valid_id.unsqueeze(1)
        B = part_code.shape[0]
        part_code = part_code.reshape(B, -1)
        if self.noise_encoder_type == 0:
            noise = self.noise_encoder(torch.cat([part_code, noise], dim=-1))
        if self.norm in ['bn', 'in']:
            noised_part_code = torch.cat([part_code, noise], dim=-1)
            m = F.relu(self.fc_bn1_m(self.fc1_m(noised_part_code)))
            m = F.relu(self.fc_bn2_m(self.fc2_m(m)))
            m = F.relu(self.fc_bn3_m(self.fc3_m(m)))
            m = self.fc4_m(m)
            v = F.relu(self.fc_bn1_v(self.fc1_v(noised_part_code)))
            v = F.relu(self.fc_bn2_v(self.fc2_v(v)))
            v = F.relu(self.fc_bn3_v(self.fc3_v(v)))
            v = self.fc4_v(v)
        elif self.norm in ['adain']:
            m = F.relu(self.fc_bn1_m(self.fc1_m(part_code), noise))
            m = F.relu(self.fc_bn2_m(self.fc2_m(m), noise))
            m = F.relu(self.fc_bn3_m(self.fc3_m(m), noise))
            m = self.fc4_m(m)
            v = F.relu(self.fc_bn1_v(self.fc1_v(part_code), noise))
            v = F.relu(self.fc_bn2_v(self.fc2_v(v), noise))
            v = F.relu(self.fc_bn3_v(self.fc3_v(v), noise))
            v = self.fc4_v(v)
        return m.reshape(B, 3, self.n_class), v.reshape(B, 3, self.n_class)

@ENCODERS.register_module()
class PartAligner(nn.Module):
    def __init__(
        self, 
        n_class: int, 
        width: int):
        super().__init__()

        self.width=width 
        self.n_class = n_class
        self.fc1_m = nn.Linear(width*n_class, 512)
        self.fc2_m = nn.Linear(512, 256)
        self.fc3_m = nn.Linear(256, 64)
        self.fc4_m = nn.Linear(64, 3*n_class)
        self.fc_bn1_m = nn.BatchNorm1d(512)
        self.fc_bn2_m = nn.BatchNorm1d(256)
        self.fc_bn3_m = nn.BatchNorm1d(64)
        
        
        self.fc1_v = nn.Linear(width*n_class, 512)
        self.fc2_v = nn.Linear(512, 256)
        self.fc3_v = nn.Linear(256, 64)
        self.fc4_v = nn.Linear(64, 3*n_class)
        self.fc_bn1_v = nn.BatchNorm1d(512)
        self.fc_bn2_v = nn.BatchNorm1d(256)
        self.fc_bn3_v = nn.BatchNorm1d(64)
        for i in range(1, 5):
            for k in ['m', 'v']:
                init_linear(getattr(self, f"fc{i}_{k}"), 0.25)
        
    def forward(self, part_code, valid_id=None, noise=None):
        r"""
            Forward pass of the network

            Parameters
            ----------
            pointcloud: Variable(torch.cuda.FloatTensor)
                (B, N, 3 + input_channels) tensor
                Point cloud to run predicts on
                Each point in the point-cloud MUST
                be formated as (x, y, z, features...)
        """
        assert part_code.shape[1:] == (self.width, self.n_class)
        part_code = part_code * valid_id.unsqueeze(1)
        B = part_code.shape[0]
        part_code = part_code.reshape(B, -1)
        m = F.relu(self.fc_bn1_m(self.fc1_m(part_code)))
        m = F.relu(self.fc_bn2_m(self.fc2_m(m)))
        m = F.relu(self.fc_bn3_m(self.fc3_m(m)))
        m = self.fc4_m(m)
        v = F.relu(self.fc_bn1_v(self.fc1_v(part_code)))
        v = F.relu(self.fc_bn2_v(self.fc2_v(v)))
        v = F.relu(self.fc_bn3_v(self.fc3_v(v)))
        v = self.fc4_v(v)
        return m.reshape(B, 3, self.n_class), v.reshape(B, 3, self.n_class)

@ENCODERS.register_module()
class PartEncoder(nn.Module):
    def __init__(
        self, 
        encoder=None, 
        n_class=4, 
        part_aligner=None, 
        fit_loss_weight=1.0, 
        include_z=True,
        include_part_code=False,
        include_params=False,
        use_gt_params=False,
        encode_ref=False,
        scale_var=1.0,
        fit_loss_type=0,
        origin_scale=False,
        kl_weight=0.001, 
        use_flow=False,
        latent_flow_depth=14, 
        latent_flow_hidden_dim=256, 
        use_gt_params_in_training=False,
        gen=False,
        gt_param_annealing=False,
        gt_param_annealing_start_epoch=500,
        gt_param_annealing_end_epoch=1000,
        kl_weight_annealing=False,
        min_kl_weight=1e-7,
        include_class_label=False,
        kl_weight_annealing_end_epoch=3000,
        normalize_part_code=False,
        detach_params_in_ctx=False,
        prior_var=1.0,
        per_part_encoder=False,
        selective_noise_sampling=False,
        selective_noise_sampling_global=False,
        ):
        super().__init__()
        self.per_part_encoder=per_part_encoder
        if not per_part_encoder:
            self.encoder = build_from_cfg(encoder, ENCODERS, num_anchors=n_class)
            self.zdim = self.encoder.zdim
        else:
            self.encoder = nn.ModuleList([build_from_cfg(encoder, ENCODERS, num_anchors=1) for _ in range(n_class)])
            self.zdim = self.encoder[0].zdim
        self.part_aligner = build_from_cfg(part_aligner, ENCODERS)
        self.include_z=include_z
        self.include_part_code=include_part_code
        self.include_params=include_params
        self.detach_params_in_ctx=detach_params_in_ctx
        self.n_class = n_class
        self.encode_ref=encode_ref
        self.prior_var=prior_var
        if encode_ref:
            self.ref_encoder = PointNet(zdim=self.encoder.zdim, point_dim=3, num_anchors=1)
        self.log_scale_var=math.log(scale_var)
        self.fit_loss_weight=fit_loss_weight
        self.use_gt_params=use_gt_params
        self.fit_loss_type = fit_loss_type
        self.selective_noise_sampling=selective_noise_sampling
        self.selective_noise_sampling_global=selective_noise_sampling_global
        self.origin_scale=origin_scale
        self.gen=gen 
        self.include_class_label=include_class_label
        self.normalize_part_code=normalize_part_code
        self.gt_param_annealing=gt_param_annealing
        self.gt_param_annealing_start_epoch=gt_param_annealing_start_epoch
        self.gt_param_annealing_end_epoch=gt_param_annealing_end_epoch
        self.use_gt_params_in_training=use_gt_params_in_training
        if gen:
            self.kl_weight=kl_weight
            self.kl_weight_annealing=kl_weight_annealing
            self.min_kl_weight=min_kl_weight
            self.kl_weight_annealing_end_epoch=kl_weight_annealing_end_epoch
            self.use_flow=use_flow
            if use_flow:
                self.flow = nn.ModuleList([build_latent_flow(latent_flow_depth, latent_flow_hidden_dim, self.zdim) for _ in range(self.n_class)])
    
    def sample_noise(self, pcds, device, num):
        input = pcds['input'].to(device)
        valid_id = pcds['present'].to(device)
        ref = pcds['ref'].to(device).transpose(1,2)
        seg_mask = pcds['ref_seg_mask'].to(device).to(torch.int32)
        seg_flag = pcds['attn_map'].to(device)
        gt_shift = pcds.get('part_shift', torch.zeros(ref.shape[0], 3, self.n_class)).to(device)
        gt_var = pcds.get('part_scale', torch.ones(ref.shape[0], 3, self.n_class)).to(device)
        if not self.origin_scale:
            gt_var = gt_var ** 2
        B, N, C = input.shape
        part_code_means, part_code_logvars = self.get_part_code(input, seg_flag)
        assert part_code_means.shape[1] == part_code_logvars.shape[1] == self.n_class
        if self.gen:
            part_code = reparameterize_gaussian(mean=part_code_means, logvar=part_code_logvars).transpose(1,2)  # (B, F)
        else:
            part_code = part_code_means.transpose(1,2)
        noise = torch.randn(B*num, self.part_aligner.noise_dim).to(device)
        part_code, valid_id, seg_mask, ref, gt_shift, gt_var = map(lambda t: t.repeat_interleave(num, dim=0), [part_code, valid_id, seg_mask, ref, gt_shift, gt_var])
        mean, logvar = self.get_params_from_part_code(part_code, valid_id, noise=noise, gt_mean=gt_shift, gt_var=gt_var, ref=ref)
        assert part_code.shape[-1] == self.n_class
        mean_per_point, logvar_per_point, flag_per_point = self.gather_all(seg_mask, anchors=mean, variances=logvar, valid_id=valid_id)
        assert mean_per_point.shape == logvar_per_point.shape == ref.shape
        fit_loss = self.get_fit_loss(ref, mean, logvar, valid_id, mean_per_point, logvar_per_point, flag_per_point, gt_shift, gt_var, seg_mask)
        # print(part_code.mean(), mean.mean(), logvar.mean(), fit_loss.mean(), gt_shift.mean(), gt_var.mean())
        id = fit_loss.reshape(B, num).min(1)[1]
        return noise.reshape(B, num, -1), id
    
    def gather_all(self, anchor_assignments, anchors=None, variances=None, valid_id=None):
        B, N = anchor_assignments.shape
        anchor_per_point = torch.zeros(B, 3, N).to(anchor_assignments.device) 
        variance_per_point = torch.zeros(B, 3, N).to(anchor_assignments.device) 
        flag_per_point = torch.ones(B, 1, N).to(anchor_assignments.device) 
        if anchors is not None:
            anchor_per_point = gather_operation(anchors.contiguous(), anchor_assignments).reshape(B, 3, N)  
        if variances is not None:
            variance_per_point = gather_operation(variances.contiguous(), anchor_assignments).reshape(B, 3, N) 
        if valid_id is not None:
            flag_per_point = gather_operation(valid_id.unsqueeze(1).contiguous(), anchor_assignments).reshape(B, 1, N) 
        return anchor_per_point, variance_per_point, flag_per_point
    def get_part_code(self, input, seg_flag):
        B = input.shape[0]
        if self.per_part_encoder:
            m, v = [], []
            for i in range(self.n_class):
                idx = seg_flag[..., i] != 1
                if not torch.all(idx):
                    _m, _v = self.encoder[i](input, idx)
                    m.append(_m)
                    v.append(_v)
                else:
                    m.append(torch.zeros(B, self.zdim, 1).cuda())
                    v.append(torch.zeros(B, self.zdim, 1).cuda())
            m = torch.stack(m, dim=1)
            v = torch.stack(v, dim=1)
            return m, v
        return self.encoder(input, seg_flag)
            
    def get_params_from_part_code(self, part_code, valid_id, gt_mean=None, gt_var=None, ref=None, noise=None, **kwargs):
        if self.part_aligner is not None:
            if self.encode_ref:
                global_feat = self.ref_encoder(ref.transpose(1,2)).expand(-1, self.n_class, -1).transpose(1,2)
                mean, logvar = self.part_aligner(global_feat, valid_id, noise=noise, **kwargs)
            else:
                mean, logvar = self.part_aligner(part_code, valid_id, noise=noise, **kwargs)
        else:
            mean = logvar = None
        if self.use_gt_params:
            logvar = torch.log(gt_var) # this is std, we need convert it to var
            mean = gt_mean
        return  mean, logvar


    def prepare_ctx(self, part_code, mean, logvar, anchor_assignments=None, z=None):
        ctx = []
        B = part_code.shape[0]
        npoint = anchor_assignments.shape[1]
            
        if self.include_z:
            ctx.append(part_code.reshape(B, -1).unsqueeze(1).repeat_interleave(npoint, dim=1))
        if self.include_part_code:
            latent_per_point = gather_operation(part_code.contiguous(), anchor_assignments).reshape(B, part_code.shape[1], npoint)
            ctx.append(latent_per_point.transpose(1,2))
        if self.include_class_label:
            seg_flag = gather_operation(torch.eye(self.n_class).cuda().unsqueeze(0).repeat_interleave(B, dim=0).contiguous(), anchor_assignments).reshape(B, self.n_class, npoint)
            ctx.append(seg_flag.transpose(1,2))
        if self.include_params:                    
            if mean is not None:
                mean_per_point = gather_operation(mean.contiguous(), anchor_assignments).reshape(B, 3, npoint)
                if self.detach_params_in_ctx:
                    mean_per_point = mean_per_point.detach()
                ctx.append(mean_per_point.transpose(1,2))
                
            if logvar is not None:
                logvar_per_point = gather_operation(logvar.contiguous(), anchor_assignments).reshape(B, 3, npoint)
                if self.detach_params_in_ctx:
                    logvar_per_point = logvar_per_point.detach()
                ctx.append(torch.exp(logvar_per_point + self.log_scale_var).transpose(1,2))
        return ctx
    
    def get_fit_loss(self, ref, mean, logvar, valid_id, mean_per_point, logvar_per_point, flag_per_point, gt_shift, gt_var, seg_mask):
        if self.part_aligner is not None:
            if self.fit_loss_type == 0:
                fit_loss = -gaussian_log_likelihood(ref, mean_per_point, logvar_per_point).mean(1, keepdim=True) * flag_per_point
                fit_loss = fit_loss.sum(dim=(-1, -2)) / flag_per_point.sum(dim=(-1, -2))
            elif self.fit_loss_type == 1:
                gt_mean = gt_shift
                gt_var = gt_var
                fit_loss = F.mse_loss(torch.cat([mean, torch.exp(logvar)], dim=1), torch.cat([gt_mean, gt_var], dim=1), reduction='none')
                fit_loss = fit_loss * valid_id.unsqueeze(1)
                fit_loss = fit_loss.sum(dim=(-1, -2)) / valid_id.sum(dim=-1)
            elif self.fit_loss_type == 2:
                gt_logvar = torch.log(gt_var)
                prior_samples = reparameterize_gaussian(mean, logvar)
                p_z = -gaussian_log_likelihood(prior_samples, gt_shift, gt_logvar, dim=3).mean(1)
                entropy = -gaussian_entropy(logvar, dim=1)
                fit_loss = p_z + entropy
                fit_loss = fit_loss * valid_id
                fit_loss = fit_loss.sum(dim=1) / valid_id.sum(dim=1)
            elif self.fit_loss_type == 3:
                gt_mean_per_point = gather_operation(gt_shift, seg_mask)
                gt_var_per_point = gather_operation(gt_var, seg_mask)
                samples = torch.randn_like(gt_var_per_point) * torch.sqrt(gt_var_per_point) + gt_mean_per_point
                fit_loss = -gaussian_log_likelihood(samples, mean_per_point, logvar_per_point).mean(1, keepdim=True) * flag_per_point
                fit_loss = fit_loss.sum(dim=(-1, -2)) / flag_per_point.sum(dim=(-1, -2))
            elif self.fit_loss_type == 4:
                gt_mean = gt_shift
                gt_logvar = torch.log(gt_var)
                fit_loss = F.mse_loss(torch.cat([mean, logvar], dim=1), torch.cat([gt_mean, gt_logvar], dim=1), reduction='none')
                fit_loss = fit_loss * valid_id.unsqueeze(1)
                fit_loss = fit_loss.sum(dim=(-1, -2)) / valid_id.sum(dim=-1)
        else:
            fit_loss = torch.zeros([1]).to(ref)
        return fit_loss
    
    # def sample_latents(self, sample_num, sample_points, device, fixed_id=None, valid_id=None, **kwargs):
    #     part_code = torch.randn(sample_num, self.zdim, self.n_class).to(device)
    #     fixed_codes = part_code[0].unsqueeze(0)
    #     part_code = part_code * (1 - fixed_id).unsqueeze(0).unsqueeze(0) + fixed_id.unsqueeze(0).unsqueeze(0) * fixed_codes
    #     if self.use_flow:
    #         part_code_ = []
    #         for i in range(self.n_class):
    #             _part_code = part_code[..., i]
    #             part_code_.append(self.flow[i](_part_code, reverse=True).view(sample_num, self.zdim))
    #         part_code = torch.stack(part_code_, dim=-1)
    #     if valid_id is None:
    #         valid_id = torch.ones(sample_num, self.n_class).to(part_code)

    #     mean, logvar = self.get_params_from_part_code(part_code, valid_id, **kwargs)
        
    #     assert part_code.shape[-1] == self.n_class
    #     ids = torch.arange(self.n_class).to(device).unsqueeze(0) * valid_id
    #     seg_mask = ids.to(torch.int32).unsqueeze(-1).expand(-1, -1, sample_points // self.n_class).reshape(sample_num, sample_points)
    #     mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=mean, variances=logvar + self.log_scale_var)
    #     ctx = self.prepare_ctx(part_code, mean, logvar + self.log_scale_var, anchor_assignments=seg_mask)
    #     return ctx, mean_per_point, logvar_per_point, seg_mask, valid_id
    def subsample_params(self, mean, logvar, valid_id, num, return_ids=False):
        print('Selective Down Sampling----------')
        B = mean.shape[0]
        K = mean.shape[1]
        param = torch.cat([mean, logvar], dim=2)
        param_score = torch.cat([mean, logvar], dim=2).reshape(B*K, 6, 4)
        outs = []
        idss = []
        expand_valid_id = valid_id.repeat_interleave(K, dim=0)
        for i in range(B*K):
            points = torch.randn(512, 3, self.n_class).cuda() * torch.exp(0.5 * logvar.reshape(B*K, 3, 4)[i][None]) + mean.reshape(B*K, 3, 4)[i][None]
            _points = points[..., expand_valid_id[i].bool()].transpose(1,2).reshape(-1, 3)
            shift = (_points.max(0)[0] + _points.min(0)[0]) / 2
            scale = (_points.max(0)[0] - _points.min(0)[0]).max() / 2
            points = (points - shift.reshape(1, 3, 1)) / scale.reshape(1, 1, 1)
            param_score[i] = torch.cat([points.mean(0), 2.0 * torch.log(points.std(0))], dim=0)
        param_score = param_score.reshape(B, K, 6, 4)
        for k in range(B):
            _param = param[k]
            _param_score = param_score[k]
            _valid_id = valid_id[k]
            out = _param[0][None]
            out_score = _param_score[0][None]
            selected = [0]
            while out.shape[0] < num:
                dists = []
                for i in range(_param.shape[0]):
                    if i in selected: continue
                    p = _param_score[i]
                    d = F.mse_loss(p[None].expand(out.shape[0], -1, -1), out_score, reduction='none')
                    d = d * _valid_id.unsqueeze(0).unsqueeze(0)
                    d = d.sum(dim=(-1, -2)) / _valid_id.sum(dim=-1)
                    dists.append((i, d.min()))
                ids, ds = zip(*dists)
                selected_id = ids[torch.argmax(torch.tensor(ds))]
                selected.append(selected_id)
                out = torch.cat([out, _param[selected_id][None]], dim=0)
                out_score = torch.cat([out_score, _param_score[selected_id][None]], dim=0)
                # print(out.shape)
            outs.append(out)
            idss.append(torch.tensor(selected).cuda())
        print(len(outs), len(idss))
        if return_ids:
            return torch.chunk(torch.stack(outs, dim=0).reshape(-1, 6, self.n_class), 2, dim=1), idss
        return torch.chunk(torch.stack(outs, dim=0).reshape(-1, 6, self.n_class), 2, dim=1)

    def subsample_params_global(self, mean, logvar, valid_id, num):
        B = mean.shape[0]
        param = torch.cat([mean, logvar], dim=1)
        param_score = torch.cat([mean, logvar], dim=1)
        for i in range(B):
            points = torch.randn(512, 3, self.n_class).cuda() * torch.exp(0.5 * logvar[i][None]) + mean[i][None]
            _points = points[..., valid_id[i].bool()].transpose(1,2).reshape(-1, 3)
            shift = (_points.max(0)[0] + _points.min(0)[0]) / 2
            scale = (_points.max(0)[0] - _points.min(0)[0]).max() / 2
            points = (points - shift.reshape(1, 3, 1)) / scale.reshape(1, 1, 1)
            param_score[i] = torch.cat([points.mean(0), 2.0 * torch.log(points.std(0))], dim=0)
        out = param[0][None]
        out_score = param_score[0][None]
        selected_ids = [0]
        selected = torch.ones(B).cuda()
        selected_valid = valid_id[0][None]
        selected[0] = 0
        while out.shape[0] < num:
            d = F.mse_loss(param_score.unsqueeze(1).expand(-1, out_score.shape[0], -1, -1), out_score[None].expand(B, -1, -1, -1), reduction='none')
            # d[:, :, :3] = 0
            d = d * valid_id.unsqueeze(1).unsqueeze(1) * selected_valid.unsqueeze(1).unsqueeze(0)
            d = d.sum(dim=(-1, -2)) / (valid_id.unsqueeze(1) * selected_valid.unsqueeze(0)).sum(-1)
            d = d.min(1)[0]
            d = d * selected.reshape(B)
            selected_id = torch.argmax(d)
            selected[selected_id] = 0
            out = torch.cat([out, param[selected_id][None]], dim=0)
            selected_ids.append(selected_id)
            selected_valid = torch.cat([selected_valid, valid_id[selected_id][None]], dim=0)
            # print(out.shape)
        return torch.chunk(out, 2, dim=1), torch.tensor(selected_ids).long()
    
    def sample_with_fixed_latents(
        self, 
        codes, 
        valid_id, 
        gt_mean, 
        gt_logvar, 
        seg_mask, 
        sample_part_id, 
        how_many_each, 
        fix_size, 
        param_sample_num, 
        selective_param_sample):
        # import pickle
        # data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp655_5e-4_debugged_adamax_n50/val/sample_airplane_n50_encode_1000.pkl", "rb"))
        # data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp624_debugged_lamp_1e-3_n10_/val/sample_lamp_n10_4000.pkl", "rb"))
        # data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp624_n100_debugged/val/sample_chairs_from_lion_1250.pkl", "rb"))
        # codes = torch.from_numpy(data['part_latents']).cuda()
        # valid_id = torch.from_numpy(data['valid_id']).cuda()
        # gt_mean = torch.from_numpy(data['sample 0 mean']).cuda()
        # gt_logvar = torch.from_numpy(data['sample 0 logvar']).cuda()
        # sample_points=8192
        # sample_part_id = 0
        # start=0
        # how_many_each = 10
        # ids = list(range(64))
        sample_points = seg_mask.shape[1]
        how_many = codes.shape[0]
        codes = codes.unsqueeze(1).repeat_interleave(how_many_each, dim=1)
        # valid_id = valid_id[ids]
        # gt_mean = gt_mean[ids]
        # gt_logvar = gt_logvar[ids]
        print(codes.shape, valid_id.shape, gt_mean.shape, gt_logvar.shape)
        w = torch.randn(how_many*how_many_each, self.zdim).cuda()
        if self.use_flow:
            part_code = self.flow[sample_part_id](w, reverse=True).view(how_many, how_many_each, self.zdim)
        else:
            part_code = w.view(how_many, how_many_each, self.zdim)
        codes[..., sample_part_id] = part_code
        if fix_size:
            param_sample_num = 1
            selective_param_sample = False
        K = 100 
        min_noises = []
        min_means = []
        min_logvars = []
        for i in range(how_many):
            noise_latent = torch.randn(how_many_each* K, self.part_aligner.noise_dim).cuda()
            _part_code = codes[i][:, None].expand(-1, K, -1, -1).reshape(K*how_many_each, self.zdim, self.n_class)
            _gt_mean = gt_mean[i].reshape(1, 1, 3, self.n_class)
            _gt_logvar = gt_logvar[i].reshape(1, 1, 3, self.n_class)
            _valid_id = valid_id[i].reshape(1, self.n_class).expand(how_many_each * K, -1)
            mean, logvar = self.get_params_from_part_code(_part_code, _valid_id, noise=noise_latent)
            print(mean.shape, logvar.shape)
            mean, logvar = mean.reshape(how_many_each, K, 3, self.n_class), logvar.reshape(how_many_each, K, 3, self.n_class)
            if fix_size:
                fit_loss = (torch.cat([mean, logvar], dim=2) - torch.cat([_gt_mean, _gt_logvar], dim=2)) ** 2
                fit_loss = fit_loss.sum(2) * _valid_id.reshape(how_many_each, K, self.n_class)
                fit_loss[..., sample_part_id] = 0
                min_noise_idx = fit_loss.sum(-1).argmin(dim=1).int()
                min_noise_idx = min_noise_idx.unsqueeze(-1)
            else:
                if selective_param_sample:
                    (mean, logvar), min_noise_idx = self.subsample_params(mean, logvar, valid_id[:, 0], num=param_sample_num, return_ids=True)
                else:
                    min_noise_idx = torch.arange(K).unsqueeze(1).cuda() < param_sample_num
                    
            print(min_noise_idx.shape)
            min_noise = gather_operation(noise_latent.reshape(how_many_each, K, self.part_aligner.noise_dim).transpose(1,2).contiguous(), min_noise_idx).transpose(1,2).reshape(how_many_each, param_sample_num, self.part_aligner.noise_dim)
            min_mean = gather_operation(mean.reshape(how_many_each, K, -1).transpose(1,2).contiguous(), min_noise_idx).transpose(1,2).reshape(how_many_each, param_sample_num, 3, self.n_class)
            min_logvar = gather_operation(logvar.reshape(how_many_each, K, -1).transpose(1,2).contiguous(), min_noise_idx).transpose(1,2).reshape(how_many_each, param_sample_num, 3, self.n_class)
            min_noises.append(min_noise)
            min_means.append(min_mean)
            min_logvars.append(min_logvar)
        min_noises = torch.stack(min_noises, dim=0)
        min_means = torch.stack(min_means, dim=0)
        min_logvars = torch.stack(min_logvars, dim=0)
        print(min_noises.shape, min_means.shape, min_logvars.shape)
        assert codes.shape[-1] == self.n_class
        codes = codes.reshape(how_many * how_many_each, 1, self.zdim, self.n_class).repeat_interleave(param_sample_num, dim=1).reshape(how_many * how_many_each*param_sample_num, self.zdim, self.n_class)
        min_noises = min_noises.reshape(how_many * how_many_each * param_sample_num, self.part_aligner.noise_dim)
        min_means = min_means.reshape(how_many * how_many_each * param_sample_num, 3, self.n_class)
        min_logvars = min_logvars.reshape(how_many * how_many_each * param_sample_num, 3, self.n_class)
        valid_id = valid_id.reshape(how_many, 1, self.n_class).repeat_interleave(how_many_each*param_sample_num, dim=1).reshape(how_many_each * how_many*param_sample_num, self.n_class)
        seg_mask = seg_mask.reshape(how_many, 1, sample_points).expand(-1, how_many_each * param_sample_num, -1).reshape(-1, sample_points)
        mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=min_means, variances=min_logvars)
        ctx = self.prepare_ctx(codes, min_means, min_logvars, anchor_assignments=seg_mask)
        print(mean_per_point.shape, logvar_per_point.shape, seg_mask.shape, valid_id.shape)
        return ctx, mean_per_point, logvar_per_point, seg_mask, valid_id, [codes, min_noises, min_means, min_logvars]
    
    def sample_with_one_fixed_latents(self):
        import pickle
        # data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp655_5e-4_debugged_adamax_n50/val/sample_airplane_n50_encode_1000.pkl", "rb"))
        # data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp624_n100_debugged/val/sample_624_train_encode_1250.pkl", "rb"))
        # data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp624_debugged_lamp_1e-3_n10_/val/sample_lamp_n10_4000.pkl", "rb"))
        data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp624_n100_debugged/val/sample_chairs_from_lion_1250.pkl", "rb"))
        
        codes = torch.from_numpy(data['part_latents']).cuda()
        valid_id = torch.from_numpy(data['valid_id']).cuda()
        gt_mean = torch.from_numpy(data['sample 0 mean']).cuda()
        gt_logvar = torch.from_numpy(data['sample 0 logvar']).cuda()
        sample_points=2048
        fixed_part_id = 0
        # sample_ids = [9, 0, 20, 37, 53, 75, 81, 94]
        sample_ids = [34]
        how_many = len(sample_ids)
        how_many_each = 10
        codes = codes[sample_ids].reshape(how_many, 1, self.zdim, self.n_class).expand(-1, how_many_each, -1, -1).clone()
        gt_mean = gt_mean[sample_ids]
        gt_logvar = gt_logvar[sample_ids]
        valid_ids = torch.ones([how_many, how_many_each, self.n_class]).cuda()
        # valid_ids = torch.tensor([0, 1, 1, 1]).reshape(1, 1, self.n_class).cuda().repeat_interleave(how_many_each, dim=1)
        #setting arm to none
        valid_ids[:, :, 3] = 0
        # setting back to none
        # valid_ids[:, 27:30, 0] = 0
        print(codes.shape, gt_mean.shape, gt_logvar.shape)
        w = torch.randn(how_many*how_many_each, self.zdim).cuda()
        for i in range(self.n_class):
            if i == fixed_part_id: 
                continue
            codes[..., i] = self.flow[i](w, reverse=True).view(how_many, how_many_each, self.zdim)
        K = 100 
        keep_K = 10
        min_noises = []
        min_means = []
        min_logvars = []
        for i, s_id in enumerate(sample_ids):
            noise_latent = torch.randn(how_many_each* K, self.part_aligner.noise_dim).cuda()
            _part_code = codes[i][:, None].expand(-1, K, -1, -1).reshape(K*how_many_each, self.zdim, self.n_class)
            _gt_mean = gt_mean[i].reshape(1, 1, 3, self.n_class)
            _gt_logvar = gt_logvar[i].reshape(1, 1, 3, self.n_class)
            _valid_id = valid_ids[i].reshape(how_many_each, 1, self.n_class).expand(-1, K, -1).reshape(how_many_each*K, self.n_class)
            mean, logvar = self.get_params_from_part_code(_part_code, _valid_id, noise=noise_latent)
            print(mean.shape, logvar.shape)
            mean, logvar = mean.reshape(how_many_each, K, 3, self.n_class), logvar.reshape(how_many_each, K, 3, self.n_class)
            fit_loss = (torch.cat([mean, logvar], dim=2) - torch.cat([_gt_mean, _gt_logvar], dim=2)) ** 2
            fit_loss = fit_loss[:, :, 3:].sum(2)[..., fixed_part_id]
            print(fit_loss.shape)
            min_noise_idx = fit_loss.topk(keep_K, dim=1, largest=False)[1].int()
            print(min_noise_idx.shape)
            min_noise = gather_operation(noise_latent.reshape(how_many_each, K, self.part_aligner.noise_dim).transpose(1,2).contiguous(), min_noise_idx).transpose(1,2)
            min_mean = gather_operation(mean.reshape(how_many_each, K, -1).transpose(1,2).contiguous(), min_noise_idx).transpose(1,2).reshape(how_many_each, keep_K, 3, self.n_class)
            min_logvar = gather_operation(logvar.reshape(how_many_each, K, -1).transpose(1,2).contiguous(), min_noise_idx).transpose(1,2).reshape(how_many_each, keep_K, 3, self.n_class)
            min_noises.append(min_noise)
            min_means.append(min_mean)
            min_logvars.append(min_logvar)
        min_noises = torch.stack(min_noises, dim=0)
        min_means = torch.stack(min_means, dim=0)
        min_logvars = torch.stack(min_logvars, dim=0)
        print(min_noises.shape, min_means.shape, min_logvars.shape)
        assert codes.shape[-1] == self.n_class
        codes = codes.reshape(how_many * how_many_each, 1, self.zdim, self.n_class).expand(-1, keep_K, -1, -1).reshape(-1, self.zdim, self.n_class)
        min_noises = min_noises.reshape(how_many * how_many_each * keep_K, self.part_aligner.noise_dim)
        min_means = min_means.reshape(how_many * how_many_each * keep_K, 3, self.n_class)
        min_logvars = min_logvars.reshape(how_many * how_many_each * keep_K, 3, self.n_class)
        valid_ids = valid_ids.reshape(how_many, how_many_each,1 , self.n_class).repeat_interleave(keep_K, dim=2).reshape(how_many_each * how_many*keep_K, self.n_class)
        ids = torch.arange(self.n_class).cuda().unsqueeze(0) * valid_ids + valid_ids.argmax(dim=1, keepdim=True) * (1 - valid_ids)
        seg_mask = ids.to(torch.int32).unsqueeze(-1).expand(-1, -1, sample_points // self.n_class).reshape(how_many_each * how_many * keep_K, sample_points)
        mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=min_means, variances=min_logvars)
        ctx = self.prepare_ctx(codes, min_means, min_logvars, anchor_assignments=seg_mask)
        print(mean_per_point.shape, logvar_per_point.shape, seg_mask.shape, valid_ids.shape)
        return ctx, mean_per_point, logvar_per_point, seg_mask, valid_ids, [codes, min_noises, min_means, min_logvars]
    
    def sample_with_given_latents(self):
        import pickle
        
        # data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp624_n100_debugged/val/part_sample_back_chairs_lion_1250.pkl", "rb"))
        # ref_data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp657_noise50/val/sample_car_1000.pkl", "rb"))
        # data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp624_debugged_lamp_1e-3_n10_/val/lamp_fix_cap0_lamp_n10_4000.pkl", "rb"))
        # data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp655_5e-4_debugged_adamax_n50/val/tail_2322_airplane_n50_1000.pkl", "rb"))
        data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp657_noise50/val/param_variation_lamp_car_1000.pkl", "rb"))
        codes = torch.from_numpy(data['part_codes']).cuda()
        noises = torch.from_numpy(data['noises']).cuda()
        gt_mean = torch.from_numpy(data['means']).cuda()
        gt_logvar = torch.from_numpy(data['logvars']).cuda()
        seg_mask = torch.from_numpy(data['pred_seg_mask']).cuda()
        valid_ids = torch.zeros(codes.shape[0], self.n_class).cuda()
        for i in range(self.n_class):
            _valid = seg_mask == i 
            _valid = torch.any(_valid, dim=1).int()
            valid_ids[:, i] = _valid
        # ref_code = torch.from_numpy(ref_data['part_latents']).cuda()[33]
        sample_points=8192
        # sample_ids = [175]
        sample_ids = list(range(60, 70))
        # sample_ids = []
        # for i in ids:
        #     sample_ids = sample_ids + [i * 3, i * 3 + 1, i * 3 + 2]
        codes = codes[sample_ids]
        # codes[..., 0] = ref_code[..., 0]
        noises = noises[sample_ids]
        gt_mean = gt_mean[sample_ids]
        gt_logvar = gt_logvar[sample_ids]
        valid_ids = valid_ids[sample_ids]

        assert codes.shape[-1] == self.n_class
        seg_mask = torch.arange(self.n_class).cuda().repeat_interleave(torch.tensor([512, 512, 512, 8192 - 3 * 512]).cuda(), dim=0)[None].repeat_interleave(valid_ids.shape[0], dim=0).int()
        print(seg_mask.shape)
        # seg_mask = ids.to(torch.int32).unsqueeze(-1).expand(-1, -1, sample_points // self.n_class).reshape(-1, sample_points)
        mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=gt_mean, variances=gt_logvar)
        ctx = self.prepare_ctx(codes, gt_mean, gt_logvar, anchor_assignments=seg_mask)
        print(mean_per_point.shape, logvar_per_point.shape, seg_mask.shape, valid_ids.shape)
        return ctx, mean_per_point, logvar_per_point, seg_mask, valid_ids, [codes, noises, gt_mean, gt_logvar]
    def sample_variation_with_fixed_latents(self):
        import pickle
        data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp655_5e-4_debugged_adamax_n50/val/sample_airplane_n50_encode_1000.pkl", "rb"))
        codes = torch.from_numpy(data['part_latents']).cuda()
        # gt_mean = torch.from_numpy(data['means']).cuda()
        # gt_logvar = torch.from_numpy(data['logvars']).cuda()
        seg_mask = torch.from_numpy(data['pred_seg_mask']).cuda()
        valid_ids = torch.zeros(codes.shape[0], self.n_class).cuda()
        for i in range(self.n_class):
            _valid = seg_mask == i 
            _valid = torch.any(_valid, dim=1).int()
            valid_ids[:, i] = _valid
        
        sample_points=8192
        sample_ids = [6, 22]
        print(codes.shape, seg_mask.shape)
        means = []
        logvars = []
        noises = []
        for s_id in sample_ids:
            K = 1000 
            code = codes[s_id][None].repeat_interleave(K, dim=0)
            valid_id = valid_ids[s_id][None].repeat_interleave(K, dim=0)
            print(code.shape, valid_id.shape)
            noise_latent = torch.randn(K, self.part_aligner.noise_dim).cuda()
            mean, logvar = self.get_params_from_part_code(code, valid_id, noise=noise_latent)
            (mean, logvar), ids = self.subsample_params(mean.reshape(1, K, 3, self.n_class), logvar.reshape(1, K, 3, self.n_class), valid_id.reshape(1, K, self.n_class)[:, 0], num=50, return_ids=True)
            ids = ids[0]
            print(ids.shape)
            selected_noise = noise_latent[ids]
            print(mean.shape, logvar.shape, selected_noise.shape)
            means.append(mean)
            logvars.append(logvar)
            noises.append(selected_noise)
        means = torch.cat(means, dim=0)
        logvars = torch.cat(logvars, dim=0)
        noises = torch.cat(noises, dim=0)
        codes = codes[sample_ids].reshape(-1, 1, self.zdim, self.n_class).repeat_interleave(50, dim=1).reshape(-1, self.zdim, self.n_class)
        valid_ids = valid_ids[sample_ids].reshape(-1, 1, self.n_class).repeat_interleave(50, dim=1).reshape(-1, self.n_class)
        print(codes.shape, valid_ids.shape, means.shape, logvars.shape, noises.shape)
            
        ids = torch.arange(self.n_class).cuda().unsqueeze(0) * valid_ids + torch.argmax(valid_ids, dim=1, keepdim=True) * (1 - valid_ids)
        seg_mask = ids.to(torch.int32).unsqueeze(-1).expand(-1, -1, sample_points // self.n_class).reshape(-1, sample_points)
        print(seg_mask.shape)
        mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=means, variances=logvars)
        ctx = self.prepare_ctx(codes, means, logvars, anchor_assignments=seg_mask)
        print(mean_per_point.shape, logvar_per_point.shape, seg_mask.shape, valid_ids.shape)
        return ctx, mean_per_point, logvar_per_point, seg_mask, valid_ids, [codes, noises, means, logvars]
    
    def part_swapping_with_fixed_latents(self):
        import pickle
        data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp624_n100_debugged/val/encoded_variation_10_624_train_encode_1250.pkl", "rb"))
        ref_data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp624_n100_debugged/val/sample_624_train_encode_1250.pkl", "rb"))
        ref_codes = torch.from_numpy(ref_data['part_latents']).cuda()
        codes = torch.from_numpy(data['part_codes']).cuda()
        noise_latents = torch.from_numpy(data['noises']).cuda()
        gt_mean = torch.from_numpy(data['means']).cuda()
        gt_logvar = torch.from_numpy(data['logvars']).cuda()
        seg_mask = torch.from_numpy(data['pred_seg_mask']).cuda()
        valid_ids = torch.zeros(codes.shape[0], self.n_class).cuda()
        for i in range(self.n_class):
            _valid = seg_mask == i 
            _valid = torch.any(_valid, dim=1).int()
            valid_ids[:, i] = _valid
        
        sample_points=2048
        start=0
        swap_ids = torch.tensor([0, 1, 0, 1]).cuda()
        ref_ids = [10, 16]
        codes = codes[start: start+50]
        noise_latents = noise_latents[start: start+50]
        gt_mean = gt_mean[start: start+50]
        gt_logvar = gt_logvar[start: start+50]
        valid_ids = valid_ids[start: start+50]
        print(codes.shape, gt_mean.shape, gt_logvar.shape, seg_mask.shape)
        means = []
        logvars = []
        noises = []
        swapped_codes = []
        swapped_valid_ids = []
        for s_id in ref_ids:
            K = 1000 
            ref_code = ref_codes[s_id][None]
            code = ref_code * swap_ids.unsqueeze(0).unsqueeze(0) + codes * (1 - swap_ids.unsqueeze(0).unsqueeze(0))
            valid_id = torch.ones_like(code)[:, 0]
            print(code.shape, valid_id.shape)
            mean, logvar = self.get_params_from_part_code(code, valid_id, noise=noise_latents)
            # (mean, logvar), ids = self.subsample_params(mean.reshape(1, K, 3, self.n_class), logvar.reshape(1, K, 3, self.n_class), valid_id.reshape(1, K, self.n_class)[:, 0], num=50, return_ids=True)
            # ids = ids[0]
            # print(ids.shape)
            # selected_noise = noise_latent[ids]
            # print(mean.shape, logvar.shape, selected_noise.shape)
            means.append(mean)
            logvars.append(logvar)
            noises.append(noise_latents)
            swapped_valid_ids.append(valid_id)
            swapped_codes.append(code)
        means = torch.cat(means, dim=0)
        logvars = torch.cat(logvars, dim=0)
        noises = torch.cat(noises, dim=0)
        swapped_valid_ids = torch.cat(swapped_valid_ids, dim=0)
        swapped_codes = torch.cat(swapped_codes, dim=0)
        print(swapped_codes.shape, swapped_valid_ids.shape, means.shape, logvars.shape, noises.shape)
            
        ids = torch.arange(self.n_class).cuda().unsqueeze(0) * swapped_valid_ids + torch.argmax(swapped_valid_ids, dim=1, keepdim=True) * (1 - swapped_valid_ids)
        seg_mask = ids.to(torch.int32).unsqueeze(-1).expand(-1, -1, sample_points // self.n_class).reshape(-1, sample_points)
        print(seg_mask.shape)
        mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=means, variances=logvars)
        ctx = self.prepare_ctx(swapped_codes, means, logvars, anchor_assignments=seg_mask)
        print(mean_per_point.shape, logvar_per_point.shape, seg_mask.shape, swapped_valid_ids.shape)
        return ctx, mean_per_point, logvar_per_point, seg_mask, swapped_valid_ids, [swapped_codes, noises, means, logvars]
    
    def part_interpolatioin_with_fixed_latents(self):
        import pickle
        data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp624_n100_debugged/val/part_swapping_406_624_train_encode_1250.pkl", "rb"))
        ref_data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/work_dirs/anchordiff_exp624_n100_debugged/val/sample_624_train_encode_1250.pkl", "rb"))
        ref_codes = torch.from_numpy(ref_data['part_latents']).cuda()
        codes = torch.from_numpy(data['part_codes']).cuda()
        noise_latents = torch.from_numpy(data['noises']).cuda()
        gt_mean = torch.from_numpy(data['means']).cuda()
        gt_logvar = torch.from_numpy(data['logvars']).cuda()
        seg_mask = torch.from_numpy(data['pred_seg_mask']).cuda()
        valid_ids = torch.zeros(codes.shape[0], self.n_class).cuda()
        for i in range(self.n_class):
            _valid = seg_mask == i 
            _valid = torch.any(_valid, dim=1).int()
            valid_ids[:, i] = _valid
        
        sample_points=2048
        interpolate_part_id = 0
        param_ids = [4, 54]
        ref_ids = [5, 6, 12, 63, 115]
        codes = codes[param_ids]
        noise_latents = noise_latents[param_ids]
        gt_mean = gt_mean[param_ids]
        gt_logvar = gt_logvar[param_ids]
        valid_ids = valid_ids[param_ids]
        print(codes.shape, gt_mean.shape, gt_logvar.shape, seg_mask.shape)
        means = []
        logvars = []
        noises = []
        interpolated_codes = []
        interpolated_valid_ids = []
        for s_id in ref_ids:
            dx = torch.arange(0, 1.0, 0.1).cuda()
            ref_interpolate_part_code = ref_codes[s_id, :, interpolate_part_id]
            interpolate_part_code = codes[..., interpolate_part_id]
            part_code_diff = ref_interpolate_part_code[None] - interpolate_part_code
            interpolated_code = codes[:, None].expand(-1, 10, -1, -1).clone()
            interpolated_noise_latents = noise_latents[:, None].repeat_interleave(10, dim=1)
            interpolated_code[..., interpolate_part_id] = interpolate_part_code.reshape(-1, 1, self.zdim) + part_code_diff.reshape(-1, 1, self.zdim) * dx.reshape(1, 10, 1)
            interpolated_code = interpolated_code.reshape(-1, self.zdim, self.n_class)
            interpolated_noise_latents = interpolated_noise_latents.reshape(-1, self.part_aligner.noise_dim)
            valid_id = torch.ones_like(interpolated_code)[:, 0]
            
            print(interpolated_code.shape, valid_id.shape, interpolated_noise_latents.shape)
            mean, logvar = self.get_params_from_part_code(interpolated_code, valid_id, noise=interpolated_noise_latents)
            # (mean, logvar), ids = self.subsample_params(mean.reshape(1, K, 3, self.n_class), logvar.reshape(1, K, 3, self.n_class), valid_id.reshape(1, K, self.n_class)[:, 0], num=50, return_ids=True)
            # ids = ids[0]
            # print(ids.shape)
            # selected_noise = noise_latent[ids]
            # print(mean.shape, logvar.shape, selected_noise.shape)
            means.append(mean)
            logvars.append(logvar)
            noises.append(interpolated_noise_latents)
            interpolated_valid_ids.append(valid_id)
            interpolated_codes.append(interpolated_code)
        means = torch.cat(means, dim=0)
        logvars = torch.cat(logvars, dim=0)
        noises = torch.cat(noises, dim=0)
        interpolated_valid_ids = torch.cat(interpolated_valid_ids, dim=0)
        interpolated_codes = torch.cat(interpolated_codes, dim=0)
        print(interpolated_codes.shape, interpolated_valid_ids.shape, means.shape, logvars.shape, noises.shape)
            
        ids = torch.arange(self.n_class).cuda().unsqueeze(0) * interpolated_valid_ids + torch.argmax(interpolated_valid_ids, dim=1, keepdim=True) * (1 - interpolated_valid_ids)
        seg_mask = ids.to(torch.int32).unsqueeze(-1).expand(-1, -1, sample_points // self.n_class).reshape(-1, sample_points)
        print(seg_mask.shape)
        mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=means, variances=logvars)
        ctx = self.prepare_ctx(interpolated_codes, means, logvars, anchor_assignments=seg_mask)
        print(mean_per_point.shape, logvar_per_point.shape, seg_mask.shape, interpolated_valid_ids.shape)
        return ctx, mean_per_point, logvar_per_point, seg_mask, interpolated_valid_ids, [interpolated_codes, noises, means, logvars]
    
    @torch.no_grad()
    def interpolate_two_shape(self, xyz1, mask1, xyz2, mask2, sample_points, part_id, param_shift, param_scale, valid_id, param, mid_num=10):
        B = xyz1.shape[0]
        one_hot_mask1 = F.one_hot(mask1, num_classes=self.n_class)
        one_hot_mask2 = F.one_hot(mask2, num_classes=self.n_class)
        part_code_means1, part_code_logvars1 = self.get_part_code(xyz1, one_hot_mask1)
        part_code_means2, part_code_logvars2 = self.get_part_code(xyz2, one_hot_mask2)
        part_code1 = reparameterize_gaussian(mean=part_code_means1, logvar=part_code_logvars1).transpose(1,2)
        part_code2 = reparameterize_gaussian(mean=part_code_means2, logvar=part_code_logvars2).transpose(1,2)[..., part_id]
        noise = torch.randn(B, 1000, self.part_aligner.noise_dim).cuda()
        mean, logvar = self.get_params_from_part_code(part_code1.repeat_interleave(1000, dim=0), valid_id.repeat_interleave(1000, dim=0), noise=noise)
        print(mean.shape, logvar.shape)
        mean = mean.reshape(B, 1000, 3, 4) 
        logvar  = logvar.reshape(B, 1000, 3, 4) 
        print(torch.cat([mean, logvar], dim=2).shape)
        print(torch.cat([param_shift, 2 * torch.log(param_scale)], dim=1).shape)
        fit_loss = (torch.cat([mean, logvar], dim=2) - torch.cat([param_shift, 2 * torch.log(param_scale)], dim=1).unsqueeze(1)) ** 2
        print(fit_loss.shape, valid_id.shape)
        fit_loss = fit_loss.sum(1) * valid_id.unsqueeze(1) # [b K 4]
        fit_loss = fit_loss.sum(-1) / torch.sum(valid_id, dim=-1, keepdim=True)
        min_noise_idx = torch.argmin(fit_loss, dim=1)
        print(min_noise_idx.shape)
        noise = gather_operation(noise.transpose(1,2).contiguous(), min_noise_idx[:, None].int()).reshape(B, 1, self.part_aligner.noise_dim).expand(-1, mid_num, -1)
        K = mid_num
        dx = torch.linspace(0, 1, steps=K).reshape(1, -1, 1).cuda()
        interpolated_part_code = part_code1[..., part_id].unsqueeze(1) + (part_code2.unsqueeze(1) - part_code1[..., part_id].unsqueeze(1)) * dx
        print(interpolated_part_code.shape)
        part_code1 = part_code1.unsqueeze(1).repeat_interleave(K, dim=1)
        print(part_code1.shape)
        # noise = noise.unsqueeze(1).repeat_interleave(K, dim=1)
        valid_id = valid_id.unsqueeze(1).repeat_interleave(K, dim=1)
        part_code1[..., part_id] = interpolated_part_code
        part_code1 = part_code1.reshape(B*K, -1, 4)
        noise = noise.reshape(B*K,-1)
        valid_id = valid_id.reshape(B*K, 4)
        mean, logvar = self.get_params_from_part_code(part_code1, valid_id, noise=noise)
        mask1 = mask1.reshape(-1, mask1.shape[1], 1).expand(-1, -1, sample_points // mask1.shape[1]).reshape(-1, sample_points)
        anchor_assignments = mask1.repeat_interleave(K, dim=0).int()
        print(anchor_assignments.shape)
        anchor_per_point, logvar_per_point, _ = self.gather_all(anchor_assignments, anchors=mean, variances=logvar)
        variance_per_point = torch.exp(logvar_per_point)
        ctx = self.prepare_ctx(part_code1, mean, logvar, anchor_assignments=anchor_assignments)
        return anchor_per_point, ctx, variance_per_point, anchor_assignments, valid_id
        
    def sample_latents(self, sample_num, sample_points, device, fixed_id=None, valid_id=None, epoch=0, K=None, part_code=None, **kwargs):
        if part_code is None:
            part_code = torch.randn(sample_num, self.zdim, self.n_class).to(device) * math.sqrt(self.prior_var)
            if self.use_flow:
                part_code_ = []
                for i in range(self.n_class):
                    _part_code = part_code[..., i]
                    part_code_.append(self.flow[i](_part_code, reverse=True).view(sample_num, self.zdim))
                part_code = torch.stack(part_code_, dim=-1)
        if self.part_aligner and self.part_aligner.cimle:
            K = 10 if K is None else K
            if self.selective_noise_sampling or self.selective_noise_sampling_global:
                K = 100
            noise = torch.randn(sample_num*K, self.part_aligner.noise_dim).to(device)
            if self.part_aligner.cimle_start_epoch > epoch:
                noise = torch.zeros_like(noise)
        else:
            K=1
            noise=None
        fixed_codes = part_code[0].unsqueeze(0)
        fixed_valid_id = valid_id[0].unsqueeze(0)
        fixed_valid_id = (fixed_valid_id + fixed_id[None]).clamp(min=0, max=1)
        part_code = part_code * (1 - fixed_id).unsqueeze(0).unsqueeze(0) + fixed_id.unsqueeze(0).unsqueeze(0) * fixed_codes
        valid_id = valid_id * (1 - fixed_id).unsqueeze(0) + fixed_id.unsqueeze(0) * fixed_valid_id
        if noise is not None and torch.any(fixed_id == 1):
            noise = noise.reshape(sample_num, K, self.part_aligner.noise_dim)
            print(noise.shape)
            fixed_noise = noise[0].unsqueeze(0)
            noise = fixed_noise.expand(sample_num, -1, -1)
            noise = noise.reshape(sample_num* K, self.part_aligner.noise_dim)
        part_code = part_code.repeat_interleave(K, dim=0)
        
        if valid_id is None:
            valid_id = torch.ones(sample_num, self.n_class).to(part_code)
            
        valid_id = valid_id.repeat_interleave(K, dim=0)
        mean, logvar = self.get_params_from_part_code(part_code, valid_id, noise=noise, **kwargs)
        if self.selective_noise_sampling:
            mean, logvar = self.subsample_params(mean.reshape(sample_num, K, 3, self.n_class), logvar.reshape(sample_num, K, 3, self.n_class), valid_id.reshape(sample_num, K, self.n_class)[:, 0], num=10)
            print(mean.shape)
            K = 10
            valid_id = valid_id.reshape(sample_num, 100, self.n_class)[:, :K].reshape(sample_num*K, self.n_class)
            part_code = part_code.reshape(sample_num, 100, -1, self.n_class)[:, :K].reshape(sample_num*K, -1, self.n_class)
        if self.selective_noise_sampling_global:
            (mean, logvar), selected_ids = self.subsample_params_global(mean, logvar, valid_id, num=sample_num*10)
            print(mean.shape, selected_ids.shape)
            valid_id = valid_id[selected_ids]
            part_code = part_code[selected_ids]
            K = 10
        print(K)
        # print(logvar.reshape(sample_num, 10, 3, self.n_class))
        assert part_code.shape[-1] == self.n_class
        # seg_mask = torch.arange(self.n_class).to(device).unsqueeze(0).repeat_interleave(torch.tensor([1024, 512, 512, 8192-2048]).cuda(), dim=1).repeat_interleave(sample_num*K, dim=0).int()
        ids = torch.arange(self.n_class).to(device).unsqueeze(0) * valid_id + torch.argmax(valid_id, dim=1).unsqueeze(1) * (1 - valid_id)
        seg_mask = ids.to(torch.int32).unsqueeze(-1).expand(-1, -1, sample_points // self.n_class).reshape(sample_num*K, sample_points)
        mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=mean, variances=logvar + self.log_scale_var if logvar is not None else logvar)
        ctx = self.prepare_ctx(part_code, mean, logvar + self.log_scale_var if logvar is not None else logvar, anchor_assignments=seg_mask)
        print(mean_per_point.shape, logvar_per_point.shape, seg_mask.shape, valid_id.shape)
        return ctx, mean_per_point, logvar_per_point, seg_mask, valid_id, [part_code, mean, logvar, noise]
    
    
    def get_prior_loss_part(self, part_code, mean, logvar, i, epoch=-1):
        B, C= part_code.shape 
        assert mean.shape[0] == logvar.shape[0]  == B
        entropy = gaussian_entropy(logvar=logvar.reshape(B, -1), dim=1)
        if self.use_flow:
            w, delta_log_pw = self.flow[i](part_code, torch.zeros([B, 1]).to(part_code), reverse=False)
            log_pw = gaussian_log_likelihood(w, 0, math.log(self.prior_var), dim=w.shape[1]).view(B, -1).sum(dim=1)   # (B, 1)
            log_p_part = log_pw - delta_log_pw.view(B)  # (B, 1)
        else:
            log_p_part = gaussian_log_likelihood(part_code, 0, math.log(self.prior_var), dim=C).sum(dim=-1)
        
        log_p_part = log_p_part.view(B)
        entropy = entropy.view(B)
        loss_prior = (- log_p_part - entropy)
        mmean = mean.mean(1)
        mlogvar = logvar.mean(1)
        mlog_p_part = log_p_part.mean()
        mentropy = entropy.mean()
        loss_dict = dict()
        if self.kl_weight_annealing and self.kl_weight_annealing_end_epoch > epoch:
            kl_weight = self.min_kl_weight + (self.kl_weight - self.min_kl_weight) * epoch / self.kl_weight_annealing_end_epoch
        else:
            kl_weight = self.kl_weight
        loss_dict['prior_loss'] = kl_weight * loss_prior.mean()
        loss_dict['kl_weight'] = torch.ones(1).cuda() * kl_weight
        loss_dict[f'log_p_part'] = mlog_p_part
        loss_dict[f'entropy'] = mentropy
        loss_dict[f'mean'] = mmean
        loss_dict[f'logvar'] = mlogvar
        return loss_dict
    def get_prior_loss(self, part_code, mean, logvar, valid_id, epoch=-1):
        B, C, M = part_code.shape 
        assert mean.shape[0] == logvar.shape[0] == valid_id.shape[0] == B
        assert mean.shape[1] == logvar.shape[1] == valid_id.shape[1] == M
        entropy = gaussian_entropy(logvar=logvar.reshape(B* self.n_class, -1), dim=1)
        if self.use_flow:
            log_p_part = torch.zeros(B, self.n_class).to(part_code)
            for i in range(self.n_class):
                _id = valid_id[:, i] == 1
                if _id.any():
                    _part_code = part_code[_id, :, i]
                    b = _part_code.shape[0]
                    # P(z), Prior probability, parameterized by the flow: z -> w.
                    # print(_part_code.abs().mean(1))
                    w, delta_log_pw = self.flow[i](_part_code, torch.zeros([b, 1]).to(part_code), reverse=False)
                    log_pw = gaussian_log_likelihood(w, 0, math.log(self.prior_var), dim=w.shape[1]).view(b, -1).sum(dim=1)   # (B, 1)
                    log_p_part[_id, i] = log_pw - delta_log_pw.view(b)  # (B, 1)
        else:
            log_p_part = gaussian_log_likelihood(part_code.transpose(1,2), 0, math.log(self.prior_var), dim=C).sum(dim=-1)
        
        log_p_part = log_p_part.view(B, self.n_class)
        entropy = entropy.view(B, self.n_class)
        loss_prior = ((- log_p_part - entropy) * valid_id).sum(1) / valid_id.sum(1)
        mmean = mean.mean(2).sum(0) / valid_id.sum(0)
        mlogvar = logvar.mean(2).sum(0) / valid_id.sum(0)
        mlog_p_part = (log_p_part * valid_id).sum(0) / valid_id.sum(0)
        mentropy = (entropy * valid_id).sum(0) / valid_id.sum(0)
        loss_dict = dict()
        if self.kl_weight_annealing and self.kl_weight_annealing_end_epoch > epoch:
            kl_weight = self.min_kl_weight + (self.kl_weight - self.min_kl_weight) * epoch / self.kl_weight_annealing_end_epoch
        else:
            kl_weight = self.kl_weight
        loss_dict['prior_loss'] = kl_weight * loss_prior.mean()
        loss_dict['kl_weight'] = torch.ones(1).cuda() * kl_weight
        for i in range(self.n_class):
            loss_dict[f'log_p_part_{i}'] = mlog_p_part[i]
            loss_dict[f'entropy_{i}'] = mentropy[i]
            # print(mlog_p_part[i])
            loss_dict[f'part_{i}_mean'] = mmean[i]
            loss_dict[f'part_{i}_logvar'] = mlogvar[i]
        return loss_dict
    
    def forward(self, pcds, device, noise=None, epoch=-1):
        r"""
            Forward pass of the network

            Parameters
            ----------
            pointcloud: Variable(torch.cuda.FloatTensor)
                (B, N, 3 + input_channels) tensor
                Point cloud to run predicts on
                Each point in the point-cloud MUST
                be formated as (x, y, z, features...)
        """
        input = pcds['input'].to(device)
        valid_id = pcds['present'].to(device)
        ref = pcds['ref'].to(device).transpose(1,2)
        seg_mask = pcds['ref_seg_mask'].to(device).to(torch.int32)
        seg_flag = pcds['ref_attn_map'].to(device)
        gt_shift = pcds.get('part_shift', torch.zeros(ref.shape[0], 3, self.n_class)).to(device)
        gt_var = pcds.get('part_scale', torch.ones(ref.shape[0], 3, self.n_class)).to(device)
        if noise is None:
            noise=pcds['noise'].to(device).unsqueeze(1)
        if not self.origin_scale:
            gt_var = gt_var ** 2
        loss_dict = dict()
        B, C, N = ref.shape
        part_code_means, part_code_logvars = self.get_part_code(input, seg_flag)
        assert part_code_means.shape[1] == part_code_logvars.shape[1] == self.n_class
        if self.gen:
            part_code = reparameterize_gaussian(mean=part_code_means, logvar=part_code_logvars).transpose(1,2)  # (B, F, M)
            prior_dict = self.get_prior_loss(part_code, part_code_means, part_code_logvars, valid_id, epoch=epoch)
            loss_dict.update(prior_dict)
        else:
            part_code = part_code_means.transpose(1,2)
        if self.normalize_part_code:
            part_code = F.normalize(part_code, p=2, dim=1)
        num_sample = noise.shape[1]
        noise = noise.reshape(B*num_sample, -1)
        if num_sample > 1:
            part_code, valid_id, seg_mask, ref, gt_shift, gt_var = map(lambda t: t.repeat_interleave(num_sample, dim=0), [part_code, valid_id, seg_mask, ref, gt_shift, gt_var])
        mean, logvar = self.get_params_from_part_code(part_code, valid_id, gt_mean=gt_shift, gt_var=gt_var, ref=ref, noise=noise)

        assert part_code.shape[-1]  == self.n_class
        
        mean_per_point, logvar_per_point, flag_per_point = self.gather_all(seg_mask, anchors=mean, variances=logvar, valid_id=valid_id)
        assert mean_per_point.shape == logvar_per_point.shape == ref.shape
        fit_loss = self.get_fit_loss(ref, mean, logvar, valid_id, mean_per_point, logvar_per_point, flag_per_point, gt_shift, gt_var, seg_mask)
        loss_dict['fit_loss'] = fit_loss
        if self.gt_param_annealing:
            assert epoch >=0, "must provide epoch as parameter when doing annealing"
            if epoch < self.gt_param_annealing_start_epoch:
                mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=gt_shift, variances=torch.log(gt_var), valid_id=None)
                mean = gt_shift
                logvar = torch.log(gt_var)
            elif epoch < self.gt_param_annealing_end_epoch:
                gt_mean_per_point, gt_logvar_per_point, _ = self.gather_all(seg_mask, anchors=gt_shift, variances=torch.log(gt_var), valid_id=None)
                gt_mean = gt_shift
                gt_logvar = torch.log(gt_var)
                gt_prob = epoch / (self.gt_param_annealing_end_epoch - self.gt_param_annealing_start_epoch)
                gt_idx = (torch.randn(B*num_sample) >= gt_prob).to(part_code).reshape(-1, 1, 1)
                mean_per_point = mean_per_point * (1 - gt_idx) + gt_idx * gt_mean_per_point
                logvar_per_point = logvar_per_point * (1 - gt_idx) + gt_idx * gt_logvar_per_point
                mean = mean * (1 - gt_idx) + gt_idx * gt_mean
                logvar =  logvar * (1 - gt_idx) + gt_idx * gt_logvar
        if self.use_gt_params_in_training:
            mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=gt_shift, variances=torch.log(gt_var), valid_id=None)
            mean = gt_shift
            logvar = torch.log(gt_var)
        ctx = self.prepare_ctx( part_code, mean, logvar, anchor_assignments=seg_mask, z=None)
        
        return ctx, mean_per_point, logvar_per_point + self.log_scale_var, flag_per_point, loss_dict, [part_code, mean, logvar, noise]



@ENCODERS.register_module()
class PartEncoderForPartnet(PartEncoder):
    def __init__(
        self, 
        **kwargs
        ):
        super().__init__(**kwargs)


    def forward(self, pcds, device):
        r"""
            Forward pass of the network

            Parameters
            ----------
            pointcloud: Variable(torch.cuda.FloatTensor)
                (B, N, 3 + input_channels) tensor
                Point cloud to run predicts on
                Each point in the point-cloud MUST
                be formated as (x, y, z, features...)
        """
        input = pcds['input'].to(device)
        valid_id = pcds['present'].to(device)
        ref = pcds['ref'].to(device).transpose(1,2)
        seg_mask = pcds['ref_seg_mask'].to(device).to(torch.int32)
        seg_flag = pcds['attn_map'].to(device)
        gt_shift = pcds.get('part_shift', torch.zeros(ref.shape[0], 3, self.n_class)).to(device)
        gt_var = pcds.get('part_scale', torch.ones(ref.shape[0], 3, self.n_class)).to(device)
        if not self.origin_scale:
            gt_var = gt_var ** 2
        loss_dict = dict()
        B, C, N = ref.shape
        part_code = self.get_part_code(input, seg_flag)
        z, part_code, mean, logvar = self.get_params_from_part_code(part_code, valid_id, gt_mean=gt_shift, gt_var=gt_var, ref=ref)
        
        assert part_code.shape[-1]  == self.n_class
        
        mean_per_point, logvar_per_point, flag_per_point = self.gather_all(seg_mask, anchors=mean, variances=logvar, valid_id=None)
        assert mean_per_point.shape == logvar_per_point.shape == ref.shape
        fit_loss = self.get_fit_loss(ref, mean, logvar, valid_id, mean_per_point, logvar_per_point, flag_per_point, gt_shift, gt_var, seg_mask)
        loss_dict['fit_loss'] = self.fit_loss_weight*fit_loss

        ctx = self.prepare_ctx(part_code, mean, logvar, anchor_assignments=seg_mask, z=z)
        
        return ctx, mean_per_point, logvar_per_point + self.log_scale_var, flag_per_point, loss_dict





@ENCODERS.register_module()
class PartEncoderForTransformerDecoder(PartEncoder):
    def __init__(
        self, 
        **kwargs
        ):
        super().__init__(**kwargs)


    def prepare_ctx(self, part_code, mean, logvar, **kwargs):
        ctx = []
        if self.include_part_code:
            ctx.append(part_code)
        if self.include_params:
            params = torch.cat([mean, torch.exp(logvar + self.log_scale_var)], dim=1)
            if self.detach_params_in_ctx:
                params = params.detach()
            ctx.append(params)
        return ctx


@ENCODERS.register_module()
class PartEncodercVAE(PartEncoderForTransformerDecoder):
    def __init__(
        self, 
        cvae_kl_weight=0.1,
        **kwargs
        ):
        super().__init__(**kwargs)
        self.cvae_encoder=MLP([self.zdim*(1 + self.n_class), self.zdim, self.zdim, self.part_aligner.noise_dim*2], use_linear=True)
        assert self.encode_ref
        self.cvae_kl_weight=cvae_kl_weight
        
    def get_params_from_part_code(self, part_code, valid_id, gt_mean=None, gt_var=None, ref=None, noise=None, **kwargs):
        B = part_code.shape[0]
        if self.training:
            global_feat = self.ref_encoder(ref.transpose(1,2)).reshape(B, -1)
            cond = torch.cat([part_code.reshape(B, -1), global_feat], dim=-1)
            noise_mean, noise_logvar = torch.chunk(self.cvae_encoder(cond), 2, dim=-1)
            noise = reparameterize_gaussian(mean=noise_mean, logvar=noise_logvar)
            entropy = gaussian_entropy(logvar=noise_logvar, dim=1)
            log_p_part = gaussian_log_likelihood(noise, 0, 1.0, dim=self.part_aligner.noise_dim).sum(dim=-1)
            loss_prior_cvae = (- log_p_part - entropy).mean()
        else:
            loss_prior_cvae = torch.zeros(1).to(part_code)
        
        mean, logvar = self.part_aligner(part_code, valid_id, noise=noise, **kwargs)
        return  mean, logvar, loss_prior_cvae
    def sample_noise(self, pcds, device, num):
        input = pcds['input'].to(device)
        B, N, C = input.shape
        return torch.randn(B, num, self.part_aligner.noise_dim).to(device), None
    
    def forward(self, pcds, device, noise=None, epoch=-1):
        r"""
            Forward pass of the network

            Parameters
            ----------
            pointcloud: Variable(torch.cuda.FloatTensor)
                (B, N, 3 + input_channels) tensor
                Point cloud to run predicts on
                Each point in the point-cloud MUST
                be formated as (x, y, z, features...)
        """
        input = pcds['input'].to(device)
        valid_id = pcds['present'].to(device)
        ref = pcds['ref'].to(device).transpose(1,2)
        seg_mask = pcds['ref_seg_mask'].to(device).to(torch.int32)
        seg_flag = pcds['ref_attn_map'].to(device)
        gt_shift = pcds.get('part_shift', torch.zeros(ref.shape[0], 3, self.n_class)).to(device)
        gt_var = pcds.get('part_scale', torch.ones(ref.shape[0], 3, self.n_class)).to(device)
        if noise is None:
            noise=pcds['noise'].to(device).unsqueeze(1)
        if not self.origin_scale:
            gt_var = gt_var ** 2
        loss_dict = dict()
        B, C, N = ref.shape
        part_code_means, part_code_logvars = self.get_part_code(input, seg_flag)
        assert part_code_means.shape[1] == part_code_logvars.shape[1] == self.n_class
        if self.gen:
            part_code = reparameterize_gaussian(mean=part_code_means, logvar=part_code_logvars).transpose(1,2)  # (B, F, M)
            prior_dict = self.get_prior_loss(part_code, part_code_means, part_code_logvars, valid_id, epoch=epoch)
            loss_dict.update(prior_dict)
        else:
            part_code = part_code_means.transpose(1,2)
        if self.normalize_part_code:
            part_code = F.normalize(part_code, p=2, dim=1)
        num_sample = noise.shape[1]
        noise = noise.reshape(B*num_sample, -1)
        if num_sample > 1:
            part_code, valid_id, seg_mask, ref, gt_shift, gt_var = map(lambda t: t.repeat_interleave(num_sample, dim=0), [part_code, valid_id, seg_mask, ref, gt_shift, gt_var])
        mean, logvar, cvae_loss = self.get_params_from_part_code(part_code, valid_id, gt_mean=gt_shift, gt_var=gt_var, ref=ref, noise=noise)
        assert part_code.shape[-1]  == self.n_class
        
        mean_per_point, logvar_per_point, flag_per_point = self.gather_all(seg_mask, anchors=mean, variances=logvar, valid_id=valid_id)
        assert mean_per_point.shape == logvar_per_point.shape == ref.shape
        fit_loss = self.get_fit_loss(ref, mean, logvar, valid_id, mean_per_point, logvar_per_point, flag_per_point, gt_shift, gt_var, seg_mask)
        loss_dict['fit_loss'] = fit_loss
        loss_dict['cvae_loss'] = cvae_loss * self.cvae_kl_weight
        if self.gt_param_annealing:
            assert epoch >=0, "must provide epoch as parameter when doing annealing"
            if epoch < self.gt_param_annealing_start_epoch:
                mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=gt_shift, variances=torch.log(gt_var), valid_id=None)
                mean = gt_shift
                logvar = torch.log(gt_var)
            elif epoch < self.gt_param_annealing_end_epoch:
                gt_mean_per_point, gt_logvar_per_point, _ = self.gather_all(seg_mask, anchors=gt_shift, variances=torch.log(gt_var), valid_id=None)
                gt_mean = gt_shift
                gt_logvar = torch.log(gt_var)
                gt_prob = epoch / (self.gt_param_annealing_end_epoch - self.gt_param_annealing_start_epoch)
                gt_idx = (torch.randn(B*num_sample) >= gt_prob).to(part_code).reshape(-1, 1, 1)
                mean_per_point = mean_per_point * (1 - gt_idx) + gt_idx * gt_mean_per_point
                logvar_per_point = logvar_per_point * (1 - gt_idx) + gt_idx * gt_logvar_per_point
                mean = mean * (1 - gt_idx) + gt_idx * gt_mean
                logvar =  logvar * (1 - gt_idx) + gt_idx * gt_logvar
        if self.use_gt_params_in_training:
            mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=gt_shift, variances=torch.log(gt_var), valid_id=None)
            mean = gt_shift
            logvar = torch.log(gt_var)
        ctx = self.prepare_ctx( part_code, mean, logvar, anchor_assignments=seg_mask, z=None)
        
        return ctx, mean_per_point, logvar_per_point + self.log_scale_var, flag_per_point, loss_dict
    
@ENCODERS.register_module()
class PartEncodercVAE2(PartEncoderForTransformerDecoder):
    def __init__(
        self, 
        cvae_kl_weight=0.1,
        **kwargs
        ):
        super().__init__(**kwargs)
        self.cvae_encoder=MLP([self.zdim*(self.n_class), self.zdim, self.zdim, self.part_aligner.noise_dim*2], use_linear=True)
        self.cvae_kl_weight=cvae_kl_weight
        
    def get_params_from_part_code(self, part_code, valid_id, gt_mean=None, gt_var=None, ref=None, noise=None, **kwargs):
        B = part_code.shape[0]
        if self.training:
            noise_mean, noise_logvar = torch.chunk(self.cvae_encoder(part_code.reshape(B, -1)), 2, dim=-1)
            noise = reparameterize_gaussian(mean=noise_mean, logvar=noise_logvar)
            entropy = gaussian_entropy(logvar=noise_logvar, dim=1)
            log_p_part = gaussian_log_likelihood(noise, 0, 1.0, dim=self.part_aligner.noise_dim).sum(dim=-1)
            loss_prior_cvae = (- log_p_part - entropy).mean()
        else:
            loss_prior_cvae = torch.zeros(1).to(part_code)
        
        mean, logvar = self.part_aligner(part_code, valid_id, noise=noise, **kwargs)
        return  mean, logvar, loss_prior_cvae
    def sample_noise(self, pcds, device, num):
        input = pcds['input'].to(device)
        B, N, C = input.shape
        return torch.randn(B, num, self.part_aligner.noise_dim).to(device), None
    
    def forward(self, pcds, device, noise=None, epoch=-1):
        r"""
            Forward pass of the network

            Parameters
            ----------
            pointcloud: Variable(torch.cuda.FloatTensor)
                (B, N, 3 + input_channels) tensor
                Point cloud to run predicts on
                Each point in the point-cloud MUST
                be formated as (x, y, z, features...)
        """
        input = pcds['input'].to(device)
        valid_id = pcds['present'].to(device)
        ref = pcds['ref'].to(device).transpose(1,2)
        seg_mask = pcds['ref_seg_mask'].to(device).to(torch.int32)
        seg_flag = pcds['ref_attn_map'].to(device)
        gt_shift = pcds.get('part_shift', torch.zeros(ref.shape[0], 3, self.n_class)).to(device)
        gt_var = pcds.get('part_scale', torch.ones(ref.shape[0], 3, self.n_class)).to(device)
        if noise is None:
            noise=pcds['noise'].to(device).unsqueeze(1)
        if not self.origin_scale:
            gt_var = gt_var ** 2
        loss_dict = dict()
        B, C, N = ref.shape
        part_code_means, part_code_logvars = self.get_part_code(input, seg_flag)
        assert part_code_means.shape[1] == part_code_logvars.shape[1] == self.n_class
        if self.gen:
            part_code = reparameterize_gaussian(mean=part_code_means, logvar=part_code_logvars).transpose(1,2)  # (B, F, M)
            prior_dict = self.get_prior_loss(part_code, part_code_means, part_code_logvars, valid_id, epoch=epoch)
            loss_dict.update(prior_dict)
        else:
            part_code = part_code_means.transpose(1,2)
        if self.normalize_part_code:
            part_code = F.normalize(part_code, p=2, dim=1)
        num_sample = noise.shape[1]
        noise = noise.reshape(B*num_sample, -1)
        if num_sample > 1:
            part_code, valid_id, seg_mask, ref, gt_shift, gt_var = map(lambda t: t.repeat_interleave(num_sample, dim=0), [part_code, valid_id, seg_mask, ref, gt_shift, gt_var])
        mean, logvar, cvae_loss = self.get_params_from_part_code(part_code, valid_id, gt_mean=gt_shift, gt_var=gt_var, ref=ref, noise=noise)
        assert part_code.shape[-1]  == self.n_class
        
        mean_per_point, logvar_per_point, flag_per_point = self.gather_all(seg_mask, anchors=mean, variances=logvar, valid_id=valid_id)
        assert mean_per_point.shape == logvar_per_point.shape == ref.shape
        fit_loss = self.get_fit_loss(ref, mean, logvar, valid_id, mean_per_point, logvar_per_point, flag_per_point, gt_shift, gt_var, seg_mask)
        loss_dict['fit_loss'] = fit_loss
        loss_dict['cvae_loss'] = cvae_loss * self.cvae_kl_weight
        if self.gt_param_annealing:
            assert epoch >=0, "must provide epoch as parameter when doing annealing"
            if epoch < self.gt_param_annealing_start_epoch:
                mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=gt_shift, variances=torch.log(gt_var), valid_id=None)
                mean = gt_shift
                logvar = torch.log(gt_var)
            elif epoch < self.gt_param_annealing_end_epoch:
                gt_mean_per_point, gt_logvar_per_point, _ = self.gather_all(seg_mask, anchors=gt_shift, variances=torch.log(gt_var), valid_id=None)
                gt_mean = gt_shift
                gt_logvar = torch.log(gt_var)
                gt_prob = epoch / (self.gt_param_annealing_end_epoch - self.gt_param_annealing_start_epoch)
                gt_idx = (torch.randn(B*num_sample) >= gt_prob).to(part_code).reshape(-1, 1, 1)
                mean_per_point = mean_per_point * (1 - gt_idx) + gt_idx * gt_mean_per_point
                logvar_per_point = logvar_per_point * (1 - gt_idx) + gt_idx * gt_logvar_per_point
                mean = mean * (1 - gt_idx) + gt_idx * gt_mean
                logvar =  logvar * (1 - gt_idx) + gt_idx * gt_logvar
        if self.use_gt_params_in_training:
            mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=gt_shift, variances=torch.log(gt_var), valid_id=None)
            mean = gt_shift
            logvar = torch.log(gt_var)
        ctx = self.prepare_ctx( part_code, mean, logvar, anchor_assignments=seg_mask, z=None)
        
        return ctx, mean_per_point, logvar_per_point + self.log_scale_var, flag_per_point, loss_dict
    
@ENCODERS.register_module()
class PartEncoderVAE(PartEncodercVAE):
    def __init__(
        self, 
        **kwargs
        ):
        super().__init__(**kwargs)
        self.ref_encoder = PointNetVAEBase(zdim=self.part_aligner.noise_dim)        
        
    def get_params_from_part_code(self, part_code, valid_id, gt_mean=None, gt_var=None, ref=None, noise=None, **kwargs):
        B = part_code.shape[0]
        if self.training:
            noise_mean, noise_logvar = self.ref_encoder(ref.transpose(1,2))
            noise = reparameterize_gaussian(mean=noise_mean, logvar=noise_logvar)
            entropy = gaussian_entropy(logvar=noise_logvar, dim=1)
            log_p_part = gaussian_log_likelihood(noise, 0, 1.0, dim=self.part_aligner.noise_dim).sum(dim=-1)
            loss_prior_cvae = (- log_p_part - entropy).mean()
        else:
            loss_prior_cvae = torch.zeros(1).to(part_code)
        
        mean, logvar = self.part_aligner(part_code, valid_id, noise=noise, **kwargs)

        return  mean, logvar, loss_prior_cvae
    
    

@ENCODERS.register_module()
class PartEncoderForPartnetAndTransformerDecoder(PartEncoderForPartnet):
    def __init__(
        self, 
        **kwargs
        ):
        super().__init__(**kwargs)
    
    def prepare_ctx(self, part_code, mean, logvar, **kwargs):
        ctx = []
        if self.include_part_code:
            ctx.append(part_code)
        if self.include_params:
            ctx.append(torch.cat([mean, torch.exp(logvar + self.log_scale_var)], dim=1))
        return ctx
    
    
@ENCODERS.register_module()
class PartEncoderCIMLE(PartEncoder):
    def __init__(self,  **kwargs):
        super().__init__(**kwargs)
    
    def get_params_from_part_code(self, part_code, valid_id, noise):
        mean, logvar = self.part_aligner(part_code, valid_id, noise)
        l = [mean]
        if self.include_var:
            l.append(logvar)
        else:
            logvar = None
        z, part_code = self.mixer(part_code, l, valid_id)
        return z, part_code, mean, logvar
    
    def sample_noise(self, pcds, device, num):
        input = pcds['input'].to(device)
        valid_id = pcds['present'].to(device)
        ref = pcds['ref'].to(device).transpose(1,2)
        seg_mask = pcds['ref_seg_mask'].to(device).to(torch.int32)
        seg_flag = pcds['attn_map'].to(device)
        B, N, C = input.shape
        part_code = self.get_part_code(torch.cat([input, seg_flag], dim=-1))
        noise = torch.randn(B*num, self.part_aligner.noise_dim).to(device)
        part_code = part_code.repeat_interleave(num, dim=0)
        valid_id = valid_id.repeat_interleave(num, dim=0)
        seg_mask = seg_mask.repeat_interleave(num, dim=0)
        ref = ref.repeat_interleave(num, dim=0)
        z, part_code, mean, logvar = self.get_params_from_part_code(part_code, valid_id, noise)
        assert part_code.shape[-1] == self.n_class
        mean_per_point, logvar_per_point = self.gather_all(seg_mask, anchors=mean, variances=logvar)
        assert mean_per_point.shape == logvar_per_point.shape == ref.shape
        fit_loss = gaussian_log_likelihood(ref, mean_per_point, logvar_per_point).mean(dim=(1,2))
        id = fit_loss.reshape(B, num).max(1)[1]
        return noise.reshape(B, num, -1), id
    def forward(self, pcds, device, noise=None):
        r"""
            Forward pass of the network

            Parameters
            ----------
            pointcloud: Variable(torch.cuda.FloatTensor)
                (B, N, 3 + input_channels) tensor
                Point cloud to run predicts on
                Each point in the point-cloud MUST
                be formated as (x, y, z, features...)
        """
        
        input = pcds['input'].to(device)
        valid_id = pcds['present'].to(device)
        ref = pcds['ref'].to(device).transpose(1,2)
        seg_mask = pcds['ref_seg_mask'].to(device).to(torch.int32)
        seg_flag = pcds['attn_map'].to(device)
        if noise is None:
            noise=pcds['noise'].to(device).unsqueeze(1)
        B, C, N = ref.shape
        part_code = self.get_part_code(torch.cat([input, seg_flag], dim=-1))
        num_sample=noise.shape[1]
        noise = noise.reshape(B*num_sample, -1)
        part_code = part_code.repeat_interleave(num_sample, dim=0)
        valid_id = valid_id.repeat_interleave(num_sample, dim=0)
        seg_mask = seg_mask.repeat_interleave(num_sample, dim=0)
        ref = ref.repeat_interleave(num_sample, dim=0)
        z, part_code, mean, logvar = self.get_params_from_part_code(part_code, valid_id, noise)
        assert part_code.shape[-1] == self.n_class
        
        mean_per_point, logvar_per_point = self.gather_all(seg_mask, anchors=mean, variances=logvar)
        assert mean_per_point.shape == logvar_per_point.shape == ref.shape
        if self.training and self.part_aligner is not None:
            fit_loss = -gaussian_log_likelihood(ref, mean_per_point, logvar_per_point).mean()
        else:
            fit_loss = torch.zeros([1]).to(device)

        ctx = self.prepare_ctx(part_code, mean, logvar, anchor_assignments=seg_mask, z=z)
        
        return ctx, mean_per_point, logvar_per_point, dict(fit_loss=self.fit_loss_weight*fit_loss)


@ENCODERS.register_module()
class PartEncoderForTransformerDecoderCIMLE(PartEncoderForTransformerDecoder):
    def __init__(
        self, 
        **kwargs
        ):
        super().__init__(**kwargs)

    def get_params_from_part_code(self, part_code, valid_id, noise):
        mean, logvar = self.part_aligner(part_code, valid_id, noise)
        l = [mean]
        if self.include_var:
            l.append(logvar)
        else:
            logvar = None
        z, part_code = self.mixer(part_code, l, valid_id)
        return z, part_code, mean, logvar
    
    def sample_noise(self, pcds, device, num):
        input = pcds['input'].to(device)
        valid_id = pcds['present'].to(device)
        ref = pcds['ref'].to(device).transpose(1,2)
        seg_mask = pcds['ref_seg_mask'].to(device).to(torch.int32)
        seg_flag = pcds['attn_map'].to(device)
        B, N, C = input.shape
        part_code = self.get_part_code(torch.cat([input, seg_flag], dim=-1))
        noise = torch.randn(B*num, self.part_aligner.noise_dim).to(device)
        part_code = part_code.repeat_interleave(num, dim=0)
        valid_id = valid_id.repeat_interleave(num, dim=0)
        seg_mask = seg_mask.repeat_interleave(num, dim=0)
        ref = ref.repeat_interleave(num, dim=0)
        z, part_code, mean, logvar = self.get_params_from_part_code(part_code, valid_id, noise)
        assert part_code.shape[-1] == self.n_class
        mean_per_point, logvar_per_point = self.gather_all(seg_mask, anchors=mean, variances=logvar)
        assert mean_per_point.shape == logvar_per_point.shape == ref.shape
        fit_loss = gaussian_log_likelihood(ref, mean_per_point, logvar_per_point).mean(dim=(1,2))
        id = fit_loss.reshape(B, num).max(1)[1]
        return noise.reshape(B, num, -1), id
    def forward(self, pcds, device, noise=None):
        r"""
            Forward pass of the network

            Parameters
            ----------
            pointcloud: Variable(torch.cuda.FloatTensor)
                (B, N, 3 + input_channels) tensor
                Point cloud to run predicts on
                Each point in the point-cloud MUST
                be formated as (x, y, z, features...)
        """
        
        input = pcds['input'].to(device)
        valid_id = pcds['present'].to(device)
        ref = pcds['ref'].to(device).transpose(1,2)
        seg_mask = pcds['ref_seg_mask'].to(device).to(torch.int32)
        seg_flag = pcds['attn_map'].to(device)
        if noise is None:
            noise=pcds['noise'].to(device).unsqueeze(1)
        B, C, N = ref.shape
        part_code = self.get_part_code(torch.cat([input, seg_flag], dim=-1))
        num_sample=noise.shape[1]
        noise = noise.reshape(B*num_sample, -1)
        part_code = part_code.repeat_interleave(num_sample, dim=0)
        valid_id = valid_id.repeat_interleave(num_sample, dim=0)
        seg_mask = seg_mask.repeat_interleave(num_sample, dim=0)
        ref = ref.repeat_interleave(num_sample, dim=0)
        z, part_code, mean, logvar = self.get_params_from_part_code(part_code, valid_id, noise)
        assert part_code.shape[-1] == self.n_class
        
        mean_per_point, logvar_per_point = self.gather_all(seg_mask, anchors=mean, variances=logvar)
        assert mean_per_point.shape == logvar_per_point.shape == ref.shape
        if self.training and self.part_aligner is not None:
            fit_loss = -gaussian_log_likelihood(ref, mean_per_point, logvar_per_point).mean()
        else:
            fit_loss = torch.zeros([1]).to(device)

        ctx = self.prepare_ctx(part_code, mean, logvar)
        
        return ctx, mean_per_point, logvar_per_point, dict(fit_loss=self.fit_loss_weight*fit_loss)
    
@ENCODERS.register_module()
class PartEncoderWithKLLoss(PartEncoder):
    def __init__(self, 
                 kl_weight=0.001, 
                 use_flow=False,
                 latent_flow_depth=14, 
                 latent_flow_hidden_dim=256, 
                 use_gt_params_in_training=False,
                 **kwargs):
        super().__init__(**kwargs)
        self.use_gt_params_in_training=use_gt_params_in_training
        self.kl_weight=kl_weight
        self.use_flow=use_flow 
        if use_flow:
            self.flow = nn.ModuleList([build_latent_flow(latent_flow_depth, latent_flow_hidden_dim, self.encoder.zdim) for _ in range(self.n_class)])
        self.zdim = self.encoder.zdim

    def get_part_code(self, input, seg_flag):
        return self.encoder(torch.cat([input, seg_flag], dim=-1))
    
    def sample_noise(self, pcds, device, num):
        input = pcds['input'].to(device)
        valid_id = pcds['present'].to(device)
        ref = pcds['ref'].to(device).transpose(1,2)
        seg_mask = pcds['ref_seg_mask'].to(device).to(torch.int32)
        seg_flag = pcds['attn_map'].to(device)
        gt_shift = pcds.get('part_shift', torch.zeros(ref.shape[0], 3, self.n_class)).to(device)
        gt_var = pcds.get('part_scale', torch.ones(ref.shape[0], 3, self.n_class)).to(device)
        if not self.origin_scale:
            gt_var = gt_var ** 2
        B, N, C = input.shape
        
        part_code_means, part_code_logvars = self.get_part_code(input, seg_flag)
        assert part_code_means.shape[1] == part_code_logvars.shape[1] == self.n_class
        part_code = reparameterize_gaussian(mean=part_code_means.reshape(B* self.n_class, -1), logvar=part_code_logvars.reshape(B* self.n_class, -1)).reshape(B, self.n_class, -1).transpose(1,2)  # (B, F)
        
        noise = torch.randn(B*num, self.part_aligner.noise_dim).to(device)
        part_code, valid_id, seg_mask, ref, gt_shift, gt_var = map(lambda t: t.repeat_interleave(num, dim=0), [part_code, valid_id, seg_mask, ref, gt_shift, gt_var])
        z, part_code, mean, logvar = self.get_params_from_part_code(part_code, valid_id, noise=noise, gt_mean=gt_shift, gt_var=gt_var, ref=ref)
        assert part_code.shape[-1] == self.n_class
        mean_per_point, logvar_per_point, flag_per_point = self.gather_all(seg_mask, anchors=mean, variances=logvar, valid_id=valid_id)
        assert mean_per_point.shape == logvar_per_point.shape == ref.shape
        fit_loss = self.get_fit_loss(ref, mean, logvar, valid_id, mean_per_point, logvar_per_point, flag_per_point, gt_shift, gt_var, seg_mask)
        id = fit_loss.reshape(B, num).min(1)[1]
        return noise.reshape(B, num, -1), id
    
    def cimle_sample_latents(self, sample_num, sample_points, device, fixed_id=None, valid_id=None, **kwargs):
        part_codes = torch.randn(sample_num, self.zdim, self.n_class).to(device)
        noise = torch.randn(sample_num*10, self.part_aligner.noise_dim).to(device)
        fixed_codes = part_codes[0].unsqueeze(0)
        part_codes = part_codes * (1 - fixed_id).unsqueeze(0).unsqueeze(0) + fixed_id.unsqueeze(0).unsqueeze(0) * fixed_codes
        part_codes = part_codes.repeat_interleave(10, dim=0)
        if self.use_flow:
            part_codes_ = []
            for i in range(self.n_class):
                _part_code = part_codes[..., i]
                part_codes_.append(self.flow[i](_part_code, reverse=True).view(sample_num, self.n_class, self.zdim))
            part_codes = torch.stack(_part_code, dim=-1)
        if valid_id is None:
            valid_id = torch.ones(sample_num, self.n_class).to(part_codes)
        valid_id = valid_id.repeat_interleave(10, dim=0)
        z, part_code, mean, logvar = self.get_params_from_part_code(part_codes, valid_id, noise=noise, **kwargs)
        # print(logvar.reshape(sample_num, 10, 3, self.n_class))
        assert part_code.shape[-1] == self.n_class
        ids = torch.arange(self.n_class).to(device).unsqueeze(0) * valid_id
        seg_mask = ids.to(torch.int32).unsqueeze(-1).expand(-1, -1, sample_points // self.n_class).reshape(sample_num*10, sample_points)
        mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=mean, variances=logvar + self.log_scale_var)
        ctx = self.prepare_ctx(part_codes, mean, logvar + self.log_scale_var)
        print(mean_per_point.shape, logvar_per_point.shape, seg_mask.shape, valid_id.shape)
        return ctx, mean_per_point, logvar_per_point, seg_mask, valid_id
    
    
    def sample_latents(self, sample_num, sample_points, device, fixed_id=None, valid_id=None, **kwargs):
        part_codes = torch.randn(sample_num, self.zdim, self.n_class).to(device)
        fixed_codes = part_codes[0].unsqueeze(0)
        part_codes = part_codes * (1 - fixed_id).unsqueeze(0).unsqueeze(0) + fixed_id.unsqueeze(0).unsqueeze(0) * fixed_codes
        if self.use_flow:
            part_codes_ = []
            for i in range(self.n_class):
                _part_code = part_codes[..., i]
                part_codes_.append(self.flow[i](_part_code, reverse=True).view(sample_num, self.zdim))
            part_codes = torch.stack(part_codes_, dim=-1)
        if valid_id is None:
            valid_id = torch.ones(sample_num, self.n_class).to(part_codes)

        z, part_code, mean, logvar = self.get_params_from_part_code(part_codes, valid_id, **kwargs)
        
        assert part_code.shape[-1] == self.n_class
        ids = torch.arange(self.n_class).to(device).unsqueeze(0) * valid_id
        seg_mask = ids.to(torch.int32).unsqueeze(-1).expand(-1, -1, sample_points // self.n_class).reshape(sample_num, sample_points)
        mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=mean, variances=logvar + self.log_scale_var)
        ctx = self.prepare_ctx(part_codes, mean, logvar + self.log_scale_var, anchor_assignments=seg_mask, z=z)
        return ctx, mean_per_point, logvar_per_point, seg_mask, valid_id
    
    def get_prior_loss(self, part_code, mean, logvar, valid_id):
        B, C, M = part_code.shape 
        assert mean.shape[0] == logvar.shape[0] == valid_id.shape[0] == B
        assert mean.shape[1] == logvar.shape[1] == valid_id.shape[1] == M
        entropy = gaussian_entropy(logvar=logvar.reshape(B* self.n_class, -1), dim=1)
        if self.use_flow:
            log_p_part = []
            for i in range(self.n_class):
                _part_code = part_code[..., i]
                # P(z), Prior probability, parameterized by the flow: z -> w.
                w, delta_log_pw = self.flow[i](_part_code, torch.zeros([B, 1]).to(part_code), reverse=False)
                log_pw = standard_normal_logprob(w).view(B, -1).sum(dim=1)   # (B, 1)
                log_p_part.append(log_pw - delta_log_pw.view(B))  # (B, 1)
            log_p_part = torch.stack(log_p_part, dim=1)
        else:
            log_p_part = standard_normal_logprob(part_code.transpose(1,2).reshape(B*self.n_class, -1)).sum(dim=1)

        log_p_part = log_p_part.view(B, self.n_class)
        entropy = entropy.view(B, self.n_class)
        loss_prior = ((- log_p_part - entropy) * valid_id).sum(1) / valid_id.sum(1)
        mmean = mean.mean(2).sum(0) / valid_id.sum(0)
        mlogvar = logvar.mean(2).sum(0) / valid_id.sum(0)
        return loss_prior, mmean, mlogvar
    
    def forward(self, pcds, device, noise=None, epoch=-1):
        r"""
            Forward pass of the network

            Parameters
            ----------
            pointcloud: Variable(torch.cuda.FloatTensor)
                (B, N, 3 + input_channels) tensor
                Point cloud to run predicts on
                Each point in the point-cloud MUST
                be formated as (x, y, z, features...)
        """
        input = pcds['input'].to(device)
        valid_id = pcds['present'].to(device)
        ref = pcds['ref'].to(device).transpose(1,2)
        seg_mask = pcds['ref_seg_mask'].to(device).to(torch.int32)
        seg_flag = pcds['attn_map'].to(device)
        gt_shift = pcds.get('part_shift', torch.zeros(ref.shape[0], 3, self.n_class)).to(device)
        gt_var = pcds.get('part_scale', torch.ones(ref.shape[0], 3, self.n_class)).to(device)
        if noise is None:
            noise=pcds['noise'].to(device).unsqueeze(1)
        if not self.origin_scale:
            gt_var = gt_var ** 2
        loss_dict = dict()
        B, C, N = ref.shape
        part_code_means, part_code_logvars = self.get_part_code(input, seg_flag)
        assert part_code_means.shape[1] == part_code_logvars.shape[1] == self.n_class
        part_code = reparameterize_gaussian(mean=part_code_means.reshape(B* self.n_class, -1), logvar=part_code_logvars.reshape(B* self.n_class, -1)).reshape(B, self.n_class, -1).transpose(1,2)  # (B, F)
        
        loss_prior, mmean, mlogvar = self.get_prior_loss(part_code, part_code_means, part_code_logvars, valid_id)
        loss_dict['prior_loss'] = self.kl_weight * loss_prior.mean()
        for i in range(self.n_class):
            loss_dict[f'part_{i}_mean'] = mmean[i]
            loss_dict[f'part_{i}_logvar'] = mlogvar[i]
        
        num_sample=noise.shape[1]
        noise = noise.reshape(B*num_sample, -1)
        if num_sample > 1:
            part_code, valid_id, seg_mask, ref, gt_shift, gt_var = map(lambda t: t.repeat_interleave(num_sample, dim=0), [part_code, valid_id, seg_mask, ref, gt_shift, gt_var])
        
        z, part_code, mean, logvar = self.get_params_from_part_code(part_code, valid_id, gt_mean=gt_shift, gt_var=gt_var, ref=ref, noise=noise)
        assert part_code.shape[-1] == self.n_class
        
        mean_per_point, logvar_per_point, flag_per_point = self.gather_all(seg_mask, anchors=mean, variances=logvar, valid_id=valid_id)
        assert mean_per_point.shape == logvar_per_point.shape == ref.shape
        fit_loss = self.get_fit_loss(ref, mean, logvar, valid_id, mean_per_point, logvar_per_point, flag_per_point, gt_shift, gt_var, seg_mask)
        loss_dict['fit_loss'] = self.fit_loss_weight * fit_loss
        if self.use_gt_params_in_training:
            mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=gt_shift, variances=torch.log(gt_var), valid_id=None)
            mean = gt_shift
            logvar = torch.log(gt_var)
        ctx = self.prepare_ctx(part_code, mean, logvar + self.log_scale_var, anchor_assignments=seg_mask, z=z)
        # print(loss_dict)
        
        return ctx, mean_per_point, logvar_per_point + self.log_scale_var, flag_per_point, loss_dict
    

@ENCODERS.register_module()
class PartEncoderForTransformerDecoderKLLoss(PartEncoderWithKLLoss):
    def __init__(
        self, 
        include_std=False,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.include_std=include_std
    
    def prepare_ctx(self, part_code, mean, logvar, **kwargs):
        ctx = []
        if self.include_part_code:
            ctx.append(part_code)
        if self.include_params:
            if not self.include_std:
                ctx.append(torch.cat([mean, torch.exp(logvar)], dim=1))
            else:
                ctx.append(torch.cat([mean, torch.exp(0.5 * logvar)], dim=1))
                
        return ctx


@ENCODERS.register_module()
class PartEncoderForTransformerDecoderKLLossCIMLE(PartEncoderForTransformerDecoderKLLoss):
    def __init__(
        self, 
        **kwargs
    ):
        super().__init__(**kwargs)


    def sample_noise(self, pcds, device, num):
        input = pcds['input'].to(device)
        valid_id = pcds['present'].to(device)
        ref = pcds['ref'].to(device).transpose(1,2)
        seg_mask = pcds['ref_seg_mask'].to(device).to(torch.int32)
        seg_flag = pcds['attn_map'].to(device)
        gt_shift = pcds.get('part_shift', torch.zeros(ref.shape[0], 3, self.n_class)).to(device)
        gt_var = pcds.get('part_scale', torch.ones(ref.shape[0], 3, self.n_class)).to(device)
        if not self.origin_scale:
            gt_var = gt_var ** 2
        B, N, C = input.shape
        part_code_means, part_code_logvars = self.get_part_code(input, seg_flag)
        assert part_code_means.shape[1] == part_code_logvars.shape[1] == self.n_class
        part_code = reparameterize_gaussian(mean=part_code_means.reshape(B* self.n_class, -1), logvar=part_code_logvars.reshape(B* self.n_class, -1)).reshape(B, self.n_class, -1).transpose(1,2)  # (B, F)
        noise = torch.randn(B*num, self.part_aligner.noise_dim).to(device)
        part_code = part_code.repeat_interleave(num, dim=0)
        valid_id = valid_id.repeat_interleave(num, dim=0)
        seg_mask = seg_mask.repeat_interleave(num, dim=0)
        ref = ref.repeat_interleave(num, dim=0)
        gt_shift = gt_shift.repeat_interleave(num, dim=0)
        gt_var = gt_var.repeat_interleave(num, dim=0)
        z, part_code, mean, logvar = self.get_params_from_part_code(part_code, valid_id, noise=noise)
        assert part_code.shape[-1] == self.n_class
        mean_per_point, logvar_per_point, flag_per_point = self.gather_all(seg_mask, anchors=mean, variances=logvar, valid_id=valid_id)
        assert mean_per_point.shape == logvar_per_point.shape == ref.shape
        fit_loss = self.get_fit_loss(ref, mean, logvar, valid_id, mean_per_point, logvar_per_point, flag_per_point, gt_shift, gt_var, seg_mask)
        id = fit_loss.reshape(B, num).min(1)[1]
        return noise.reshape(B, num, -1), id
    
    def cimle_sample_latents(self, sample_num, sample_points, device, fixed_id=None, valid_id=None, **kwargs):
        part_codes = torch.randn(sample_num, self.zdim, self.n_class).to(device)
        noise = torch.randn(sample_num*10, self.part_aligner.noise_dim).to(device)
        fixed_codes = part_codes[0].unsqueeze(0)
        part_codes = part_codes * (1 - fixed_id).unsqueeze(0).unsqueeze(0) + fixed_id.unsqueeze(0).unsqueeze(0) * fixed_codes
        part_codes = part_codes.repeat_interleave(10, dim=0)
        if self.use_flow:
            part_codes_ = []
            for i in range(self.n_class):
                _part_code = part_codes[..., i]
                part_codes_.append(self.flow[i](_part_code, reverse=True).view(sample_num, self.n_class, self.zdim))
            part_codes = torch.stack(_part_code, dim=-1)
        if valid_id is None:
            valid_id = torch.ones(sample_num, self.n_class).to(part_codes)
        valid_id = valid_id.repeat_interleave(10, dim=0)
        z, part_code, mean, logvar = self.get_params_from_part_code(part_codes, valid_id, noise=noise, **kwargs)
        # print(logvar.reshape(sample_num, 10, 3, self.n_class))
        assert part_code.shape[-1] == self.n_class
        ids = torch.arange(self.n_class).to(device).unsqueeze(0) * valid_id
        seg_mask = ids.to(torch.int32).unsqueeze(-1).expand(-1, -1, sample_points // self.n_class).reshape(sample_num*10, sample_points)
        mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=mean, variances=logvar + self.log_scale_var)
        ctx = self.prepare_ctx(part_codes, mean, logvar + self.log_scale_var)
        print(mean_per_point.shape, logvar_per_point.shape, seg_mask.shape, valid_id.shape)
        return ctx, mean_per_point, logvar_per_point, seg_mask, valid_id
    
    
    def forward(self, pcds, device, noise=None):
        r"""
            Forward pass of the network

            Parameters
            ----------
            pointcloud: Variable(torch.cuda.FloatTensor)
                (B, N, 3 + input_channels) tensor
                Point cloud to run predicts on
                Each point in the point-cloud MUST
                be formated as (x, y, z, features...)
        """
        input = pcds['input'].to(device)
        valid_id = pcds['present'].to(device)
        ref = pcds['ref'].to(device).transpose(1,2)
        seg_mask = pcds['ref_seg_mask'].to(device).to(torch.int32)
        seg_flag = pcds['ref_attn_map'].to(device)
        gt_shift = pcds.get('part_shift', torch.zeros(ref.shape[0], 3, self.n_class)).to(device)
        gt_var = pcds.get('part_scale', torch.ones(ref.shape[0], 3, self.n_class)).to(device)
        if not self.origin_scale:
            gt_var = gt_var ** 2
        loss_dict = dict()
        B, C, N = ref.shape
        if noise is None:
            noise=pcds['noise'].to(device).unsqueeze(1)
            
        part_code_means, part_code_logvars = self.get_part_code(input, seg_flag)
        assert part_code_means.shape[1] == part_code_logvars.shape[1] == self.n_class
        part_code = reparameterize_gaussian(mean=part_code_means.reshape(B* self.n_class, -1), logvar=part_code_logvars.reshape(B* self.n_class, -1)).reshape(B, self.n_class, -1).transpose(1,2)  # (B, F)
        
        num_sample=noise.shape[1]
        noise = noise.reshape(B*num_sample, -1)
        if num_sample > 1:
            part_code_means = part_code_means.repeat_interleave(num_sample, dim=0)
            part_code_logvars = part_code_logvars.repeat_interleave(num_sample, dim=0)
            part_code = part_code.repeat_interleave(num_sample, dim=0)
            valid_id = valid_id.repeat_interleave(num_sample, dim=0)
            seg_mask = seg_mask.repeat_interleave(num_sample, dim=0)
            ref = ref.repeat_interleave(num_sample, dim=0)
            gt_shift = gt_shift.repeat_interleave(num_sample, dim=0)
            gt_var = gt_var.repeat_interleave(num_sample, dim=0)
            
        loss_prior = self.get_prior_loss(part_code, part_code_means, part_code_logvars, valid_id)
        loss_dict['prior_loss'] = self.kl_weight * loss_prior
        
        z, part_code, mean, logvar = self.get_params_from_part_code(part_code, valid_id, noise=noise)
        assert part_code.shape[-1] == self.n_class
        
        mean_per_point, logvar_per_point, flag_per_point = self.gather_all(seg_mask, anchors=mean, variances=logvar, valid_id=valid_id)
        assert mean_per_point.shape == logvar_per_point.shape == ref.shape
        fit_loss = self.get_fit_loss(ref, mean, logvar, valid_id, mean_per_point, logvar_per_point, flag_per_point, gt_shift, gt_var, seg_mask)
        loss_dict['fit_loss'] = self.fit_loss_weight * fit_loss
        ctx = self.prepare_ctx(part_code, mean, logvar + self.log_scale_var)
        
        return ctx, mean_per_point, logvar_per_point + self.log_scale_var, flag_per_point, loss_dict

@ENCODERS.register_module()
class EncoderWithKLLoss(PartEncoder):
    def __init__(self,  kl_weight=0.001,
                use_flow=False,
                latent_flow_depth=14,
                latent_flow_hidden_dim=256,
                **kwargs,):
        super().__init__(**kwargs)
        self.zdim = self.encoder.zdim
        self.kl_weight=kl_weight
        self.use_flow=use_flow 
        if use_flow:
            self.flow = build_latent_flow(latent_flow_depth, latent_flow_hidden_dim, self.encoder.zdim)
        self.zdim = self.encoder.zdim
    def get_part_code_params(self, input):
        return self.encoder(input)
    
    def sample_latents(self, sample_num, sample_points, device, fixed_id=None, valid_id=None, **kwargs):
        z = torch.randn(sample_num, self.zdim).to(device)
        if self.use_flow:
            z = self.flow(z, reverse=True).view(sample_num, self.zdim)
        
        part_code = z.unsqueeze(1).repeat_interleave(sample_points, dim=1)
        mean, logvar = self.get_params_from_part_code(part_code, valid_id)
        
        ids = torch.arange(self.n_class).to(device).unsqueeze(0) * valid_id
        seg_mask = ids.to(torch.int32).unsqueeze(-1).expand(-1, -1, sample_points // self.n_class).reshape(sample_num, -1)
        print(seg_mask.shape)
        mean_per_point, logvar_per_point, _ = self.gather_all(seg_mask, anchors=mean, variances=logvar)
        return [part_code], mean_per_point, logvar_per_point, seg_mask, valid_id

    def forward(self, pcds, device, **kwargs):
        r"""
            Forward pass of the network

            Parameters
            ----------
            pointcloud: Variable(torch.cuda.FloatTensor)
                (B, N, 3 + input_channels) tensor
                Point cloud to run predicts on
                Each point in the point-cloud MUST
                be formated as (x, y, z, features...)
        """
        ref = pcds['ref'].to(device).transpose(1,2)
        valid_id = pcds['present'].to(device)
        seg_mask = pcds['ref_seg_mask'].to(device).to(torch.int32)
        loss_dict = dict()
        B = ref.shape[0]
        mean, logvar = self.get_part_code_params(ref)
        mean, logvar = mean.reshape(B, self.zdim), logvar.reshape(B, self.zdim)
        z = reparameterize_gaussian(mean=mean, logvar=logvar).reshape(B, -1)  # (B, F)
        if self.use_flow:
            # P(z), Prior probability, parameterized by the flow: z -> w.
            w, delta_log_pw = self.flow(z.reshape(B, self.zdim), torch.zeros([B, 1]).to(z), reverse=False)
            log_pw = standard_normal_logprob(w).view(B, -1).sum(dim=1, keepdim=True)   # (B, 1)
            log_pz = log_pw - delta_log_pw.view(B, 1)  # (B, 1)
        else:
            log_pz = standard_normal_logprob(z).sum(dim=1)
        entropy = gaussian_entropy(logvar=logvar)
        loss_prior = (- log_pz - entropy).mean()
        loss_dict['prior_loss'] = self.kl_weight * loss_prior
        
        part_code = z.unsqueeze(-1).expand(-1, -1, self.n_class)
        mean, logvar = self.get_params_from_part_code(part_code, valid_id)
        
        mean_per_point, logvar_per_point, flag_per_point = self.gather_all(seg_mask, anchors=mean, variances=logvar, valid_id=valid_id)
        assert mean_per_point.shape == logvar_per_point.shape == ref.shape
        if self.training and self.part_aligner is not None:
            fit_loss = -gaussian_log_likelihood(ref, mean_per_point, logvar_per_point).mean(1, keepdim=True) * flag_per_point
            loss_dict['fit_loss'] = self.fit_loss_weight * fit_loss.sum(dim=(-1, -2)) / flag_per_point.sum(dim=(-1, -2))
        else:
            fit_loss = torch.zeros([1]).to(device)
        ctx = self.prepare_ctx( part_code, mean, logvar, anchor_assignments=seg_mask, z=z)
        return ctx, mean_per_point, logvar_per_point, flag_per_point, loss_dict
    
@ENCODERS.register_module()
class PartEncoderNoMixer(PartEncoder):
    def __init__(self, encoder, n_class, part_aligner=None, include_var=False, fit_loss_weight=1.0):
        super().__init__(encoder, None, n_class, part_aligner=part_aligner, include_var=include_var, fit_loss_weight=fit_loss_weight)
    
    def get_params_from_part_code(self, part_code, valid_id):
        B, C, N = part_code.shape
        if self.part_aligner is not None:
            mean, logvar = self.part_aligner(part_code, valid_id)
            l = [mean]
            if self.include_var:
                l.append(logvar)
            else:
                logvar = None
        else:
            mean = logvar = None
        return torch.zeros(B, C).to(part_code.device), part_code, mean, logvar
        
        
        
    
python/difffacto/models/encoders/part_encoders.py
diffFacto/diffFacto
import torch
import torch.nn as nn
from difffacto.utils.registry import ENCODERS

@ENCODERS.register_module()
class Pnet2Stage(nn.Module):
    def __init__(self, latent_dim=1024, point_dim=3):
        super().__init__()
        self.latent_dim = latent_dim
        self.first_conv = nn.Sequential(
            nn.Conv1d(point_dim, 128, 1),
            nn.BatchNorm1d(128),
            nn.ReLU(inplace=True),
            nn.Conv1d(128, 256, 1)
        )

        self.second_conv = nn.Sequential(
            nn.Conv1d(512, 512, 1),
            nn.BatchNorm1d(512),
            nn.ReLU(inplace=True),
            nn.Conv1d(512, self.latent_dim, 1)
        )

    def forward(self, xyz):
        B, N, C = xyz.shape
        
        # encoder
        feature = self.first_conv(xyz.transpose(2, 1))                                       # (B,  256, N)
        feature_global = torch.max(feature, dim=2, keepdim=True)[0]                          # (B,  256, 1)
        feature = torch.cat([feature_global.expand(-1, -1, N), feature], dim=1)              # (B,  512, N)
        feature = self.second_conv(feature)                                                  # (B, 1024, N)
        feature_global = torch.max(feature,dim=2,keepdim=False)[0]
        
        return feature_global
python/difffacto/models/encoders/pnet.py
diffFacto/diffFacto
import torch 
from torch import nn
import torch.nn.functional as F
from difffacto.utils.registry import ENCODERS

@ENCODERS.register_module()
class PointNetVAEBase(nn.Module):
    def __init__(self, point_dim=3,  zdim=1024, **kwargs):
        super().__init__()
        self.zdim = zdim
        self.conv1 = nn.Conv1d(point_dim, 128 , 1)
        self.conv2 = nn.Conv1d(128, 128, 1)
        self.conv3 = nn.Conv1d(128, 256, 1)
        self.conv4 = nn.Conv1d(256, 512, 1)
        self.bn1 = nn.BatchNorm1d(128)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(256)
        self.bn4 = nn.BatchNorm1d(512)

        # Mapping to [c], cmean
        self.fc1_m = nn.Linear(512, 256)
        self.fc2_m = nn.Linear(256, 128)
        self.fc3_m = nn.Linear(128, zdim)
        self.fc_bn1_m = nn.BatchNorm1d(256)
        self.fc_bn2_m = nn.BatchNorm1d(128)
        
        # Mapping to [c], cmean
        self.fc1_v = nn.Linear(512, 256)
        self.fc2_v = nn.Linear(256, 128)
        self.fc3_v = nn.Linear(128, zdim)
        self.fc_bn1_v = nn.BatchNorm1d(256)
        self.fc_bn2_v = nn.BatchNorm1d(128)

    def forward(self, x, mask=None):
        B = x.shape[0]
        x = x.transpose(1, 2)
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = self.bn4(self.conv4(x))
        if mask is not None:
            INF = float('inf')
            _mask = torch.zeros_like(x)
            _mask.masked_fill_(mask[:, None], -1.0 * float('inf'))
            x = x + _mask
        x = torch.max(x, 2, keepdim=True)[0]
        x = x.view(-1, 512)
        x = torch.nan_to_num(x, neginf=0.0)
        m = F.relu(self.fc_bn1_m(self.fc1_m(x)))
        m = F.relu(self.fc_bn2_m(self.fc2_m(m)))
        m = self.fc3_m(m)
        
        v = F.relu(self.fc_bn1_v(self.fc1_v(x)))
        v = F.relu(self.fc_bn2_v(self.fc2_v(v)))
        v = self.fc3_v(v)

        # Returns both mean and logvariance, just ignore the latter in deteministic cases.
        return m, v


@ENCODERS.register_module()
class PointNetVAE(nn.Module):
    def __init__(self, point_dim=7,  zdim=1024, num_anchors=4):
        super().__init__()
        self.zdim = zdim
        self.num_anchors=num_anchors
        self.conv1 = nn.Conv1d(point_dim, 128, 1)
        self.conv2 = nn.Conv1d(128, 128, 1)
        self.conv3 = nn.Conv1d(128, 256, 1)
        self.conv4 = nn.Conv1d(256, 512, 1)
        self.bn1 = nn.BatchNorm1d(128)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(256)
        self.bn4 = nn.BatchNorm1d(512)

        # Mapping to [c], cmean
        self.fc1_m = nn.Linear(512, 256)
        self.fc2_m = nn.Linear(256, 128)
        self.fc3_m = nn.Linear(128, zdim*num_anchors)
        self.fc_bn1_m = nn.BatchNorm1d(256)
        self.fc_bn2_m = nn.BatchNorm1d(128)
        
        # Mapping to [c], cmean
        self.fc1_v = nn.Linear(512, 256)
        self.fc2_v = nn.Linear(256, 128)
        self.fc3_v = nn.Linear(128, zdim*num_anchors)
        self.fc_bn1_v = nn.BatchNorm1d(256)
        self.fc_bn2_v = nn.BatchNorm1d(128)

    def forward(self, x):
        B = x.shape[0]
        x = x.transpose(1, 2)
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = self.bn4(self.conv4(x))
        x = torch.max(x, 2, keepdim=True)[0]
        x = x.view(B, -1)

        m = F.relu(self.fc_bn1_m(self.fc1_m(x)))
        m = F.relu(self.fc_bn2_m(self.fc2_m(m)))
        m = self.fc3_m(m)
        
        v = F.relu(self.fc_bn1_v(self.fc1_v(x)))
        v = F.relu(self.fc_bn2_v(self.fc2_v(v)))
        v = self.fc3_v(v)

        # Returns both mean and logvariance, just ignore the latter in deteministic cases.
        return m.reshape(B, self.num_anchors, self.zdim), v.reshape(B, self.num_anchors, self.zdim)

@ENCODERS.register_module()
class PointNet(PointNetVAE):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def forward(self, x):
        B = x.shape[0]
        m, v = super().forward(x)
        return m.reshape(B, self.num_anchors, self.zdim)
    

    
@ENCODERS.register_module()
class PointNetV2(nn.Module):
    def __init__(self, point_dim=3, zdim=1024, num_anchors=4, reweight_by_anchor=True, use_ln=False, per_part_mlp=False):
        super().__init__()
        self.reweight_by_anchor=reweight_by_anchor
        self.per_part_mlp=per_part_mlp
        self.zdim = zdim
        self.num_anchors=num_anchors
        self.use_ln=use_ln
        self.conv1 = nn.Conv1d(point_dim, 128, 1)
        self.conv2 = nn.Conv1d(128, 128, 1)
        self.conv3 = nn.Conv1d(128, 256, 1)
        self.conv4 = nn.Conv1d(256, 512, 1)
        self.bn1 = nn.BatchNorm1d(128)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(256)
        self.bn4 = nn.BatchNorm1d(512)

        if use_ln:
            self.mlp_m = nn.Sequential(nn.Linear(512, 256), 
                                     nn.LayerNorm(256),
                                     nn.ReLU(),
                                     nn.Linear(256, 128), 
                                     nn.LayerNorm(128),
                                     nn.ReLU(),
                                     nn.Linear(128, zdim))
            self.mlp_v = nn.Sequential(nn.Linear(512, 256), 
                                     nn.LayerNorm(256),
                                     nn.ReLU(),
                                     nn.Linear(256, 128), 
                                     nn.LayerNorm(128),
                                     nn.ReLU(),
                                     nn.Linear(128, zdim))
        else:
            self.mlp_m = nn.Sequential(nn.Conv1d(512, 256, 1), 
                                     nn.BatchNorm1d(256),
                                     nn.ReLU(),
                                     nn.Conv1d(256, 128, 1), 
                                     nn.BatchNorm1d(128),
                                     nn.ReLU(),
                                     nn.Conv1d(128, zdim, 1))
            self.mlp_v = nn.Sequential(nn.Conv1d(512, 256, 1), 
                                     nn.BatchNorm1d(256),
                                     nn.ReLU(),
                                     nn.Conv1d(256, 128, 1), 
                                     nn.BatchNorm1d(128),
                                     nn.ReLU(),
                                     nn.Conv1d(128, zdim, 1))
        if self.per_part_mlp:
            self.mlp_m = nn.Sequential(nn.Conv1d(512*self.num_anchors, 256*self.num_anchors, 1, groups=self.num_anchors), 
                                     nn.BatchNorm1d(256*self.num_anchors),
                                     nn.ReLU(),
                                     nn.Conv1d(256*self.num_anchors, 128*self.num_anchors, 1, groups=self.num_anchors), 
                                     nn.BatchNorm1d(128*self.num_anchors),
                                     nn.ReLU(),
                                     nn.Conv1d(128*self.num_anchors, zdim*self.num_anchors, 1, groups=self.num_anchors))
            self.mlp_v = nn.Sequential(nn.Conv1d(512*self.num_anchors, 256*self.num_anchors, 1, groups=self.num_anchors), 
                                     nn.BatchNorm1d(256*self.num_anchors),
                                     nn.ReLU(),
                                     nn.Conv1d(256*self.num_anchors, 128*self.num_anchors, 1, groups=self.num_anchors), 
                                     nn.BatchNorm1d(128*self.num_anchors),
                                     nn.ReLU(),
                                     nn.Conv1d(128*self.num_anchors, zdim*self.num_anchors, 1, groups=self.num_anchors))

    def forward(self, x, attn_weight):
        B = x.shape[0]
        x = x.transpose(1, 2)
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = self.bn4(self.conv4(x))
        weighted_x = x.unsqueeze(-1) * attn_weight.unsqueeze(1)
        if self.reweight_by_anchor:
            weighted_x = weighted_x * self.num_anchors
        x = torch.max(weighted_x, 2, keepdim=True)[0]
        x = x.view(B, 512, self.num_anchors)
        
        if self.per_part_mlp:
            x = x.transpose(1,2).reshape(B, -1, 1)
            m = self.mlp_m(x).reshape(B, self.num_anchors, -1)
            v = self.mlp_v(x).reshape(B, self.num_anchors, -1)
        else:
            if self.use_ln:
                x = x.transpose(1,2)
                m = self.mlp_m(x)
                v = self.mlp_v(x)
            else:
                m = self.mlp_m(x).transpose(1,2)
                v = self.mlp_v(x).transpose(1,2)
        # Returns both mean and logvariance, just ignore the latter in deteministic cases.
        return m, v
    
@ENCODERS.register_module()
class PointNetV2Variant(nn.Module):
    def __init__(self, point_dim=3, zdim=1024, num_anchors=4, per_part_mlp=False):
        super().__init__()
        self.per_part_mlp=per_part_mlp
        self.zdim = zdim
        self.num_anchors=num_anchors
        self.pnet = nn.Sequential(nn.Conv1d(point_dim, 128, 1), 
                                  nn.BatchNorm1d(128), 
                                  nn.ReLU(),
                                  nn.Conv1d(128, 128, 1),
                                  nn.BatchNorm1d(128),
                                  nn.ReLU(),
                                  nn.Conv1d(128, 256, 1),
                                  nn.BatchNorm1d(256),
                                  nn.ReLU(),
                                  nn.Conv1d(256, 512, 1),
                                  nn.BatchNorm1d(512))

        self.mlp_m = nn.ModuleList([nn.Sequential(nn.Conv1d(512, 256, 1), 
                                    nn.BatchNorm1d(256),
                                    nn.ReLU(),
                                    nn.Conv1d(256, 128, 1), 
                                    nn.BatchNorm1d(128),
                                    nn.ReLU(),
                                    nn.Conv1d(128, zdim, 1)) for _ in range(self.num_anchors)])
        self.mlp_v = nn.ModuleList([nn.Sequential(nn.Conv1d(512, 256, 1), 
                                    nn.BatchNorm1d(256),
                                    nn.ReLU(),
                                    nn.Conv1d(256, 128, 1), 
                                    nn.BatchNorm1d(128),
                                    nn.ReLU(),
                                    nn.Conv1d(128, zdim, 1)) for _ in range(self.num_anchors)])
        if not self.per_part_mlp:
            self.mlp_m = self.mlp_m[0]
            self.mlp_v = self.mlp_v[0]
    def get_part_parameters(self, part_id):
        assert 0 <= part_id < self.num_anchors
        return list(self.pnet.parameters) + list(self.mlp_m[part_id].parameters()) + list(self.mlp_v[part_id].parameters())
    def forward(self, x, seg_flag, part_id=-1):
        assert x.shape[-1] == 3 + self.num_anchors
        B = x.shape[0]
        x = x.transpose(1, 2)
        x = self.pnet(x)
        weighted_x = x.unsqueeze(-1) * seg_flag.unsqueeze(1) - (1 - seg_flag.unsqueeze(1)) * float('inf')
        x = torch.max(weighted_x, 2, keepdim=True)[0]
        x = torch.nan_to_num(x, 0, 0, 0)
        x = x.view(B, 512, self.num_anchors)
        if part_id >= 0:
            return self.mlp_m[part_id](x[..., part_id]), self.mlp_v[part_id](x[..., part_id])
        if self.per_part_mlp:
            m, v = [], []
            for i in range(self.num_anchors):
                mm = self.mlp_m[i](x[..., i])
                vv = self.mlp_v[i](x[..., i])
                m.append(mm)
                v.append(vv)
            m = torch.stack(m, dim=1)
            v = torch.stack(v, dim=1)
        else:
            m = self.mlp_m(x).transpose(1,2)
            v = self.mlp_v(x).transpose(1,2)
        # Returns both mean and logvariance, just ignore the latter in deteministic cases.
        return m, v

@ENCODERS.register_module()
class PointNetV2VAE(nn.Module):
    def __init__(self, point_dim=3, zdim=1024, num_anchors=4, reweight_by_anchor=True, use_ln=False, per_part_mlp=False, deterministic=False):
        super().__init__()
        self.reweight_by_anchor=reweight_by_anchor
        self.per_part_mlp=per_part_mlp
        self.zdim = zdim
        self.num_anchors=num_anchors
        self.use_ln=use_ln
        self.conv1 = nn.Conv1d(point_dim, 128, 1)
        self.conv2 = nn.Conv1d(128, 128, 1)
        self.conv3 = nn.Conv1d(128, 256, 1)
        self.conv4 = nn.Conv1d(256, 512, 1)
        self.bn1 = nn.BatchNorm1d(128)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(256)
        self.bn4 = nn.BatchNorm1d(512)
        self.deterministic=deterministic

        if use_ln:
            self.mlp_m = nn.Sequential(nn.Linear(512, 256), 
                                     nn.LayerNorm(256),
                                     nn.ReLU(),
                                     nn.Linear(256, 128), 
                                     nn.LayerNorm(128),
                                     nn.ReLU(),
                                     nn.Linear(128, zdim))
            self.mlp_v = nn.Sequential(nn.Linear(512, 256), 
                                     nn.LayerNorm(256),
                                     nn.ReLU(),
                                     nn.Linear(256, 128), 
                                     nn.LayerNorm(128),
                                     nn.ReLU(),
                                     nn.Linear(128, zdim))
        else:
            self.mlp_m = nn.Sequential(nn.Conv1d(512, 256, 1), 
                                     nn.BatchNorm1d(256),
                                     nn.ReLU(),
                                     nn.Conv1d(256, 128, 1), 
                                     nn.BatchNorm1d(128),
                                     nn.ReLU(),
                                     nn.Conv1d(128, zdim, 1))
            self.mlp_v = nn.Sequential(nn.Conv1d(512, 256, 1), 
                                     nn.BatchNorm1d(256),
                                     nn.ReLU(),
                                     nn.Conv1d(256, 128, 1), 
                                     nn.BatchNorm1d(128),
                                     nn.ReLU(),
                                     nn.Conv1d(128, zdim, 1))
        if self.per_part_mlp:
            self.mlp_m = nn.Sequential(nn.Conv1d(512*self.num_anchors, 256*self.num_anchors, 1, groups=self.num_anchors), 
                                     nn.BatchNorm1d(256*self.num_anchors),
                                     nn.ReLU(),
                                     nn.Conv1d(256*self.num_anchors, 128*self.num_anchors, 1, groups=self.num_anchors), 
                                     nn.BatchNorm1d(128*self.num_anchors),
                                     nn.ReLU(),
                                     nn.Conv1d(128*self.num_anchors, zdim*self.num_anchors, 1, groups=self.num_anchors))
            self.mlp_v = nn.Sequential(nn.Conv1d(512*self.num_anchors, 256*self.num_anchors, 1, groups=self.num_anchors), 
                                     nn.BatchNorm1d(256*self.num_anchors),
                                     nn.ReLU(),
                                     nn.Conv1d(256*self.num_anchors, 128*self.num_anchors, 1, groups=self.num_anchors), 
                                     nn.BatchNorm1d(128*self.num_anchors),
                                     nn.ReLU(),
                                     nn.Conv1d(128*self.num_anchors, zdim*self.num_anchors, 1, groups=self.num_anchors))

    def forward(self, x):
        assert x.shape[-1] == 3 + self.num_anchors
        attn_weight = x[..., 3:]
        x = x[..., :3]
        B = x.shape[0]
        x = x.transpose(1, 2)
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = self.bn4(self.conv4(x))
        weighted_x = x.unsqueeze(-1) * attn_weight.unsqueeze(1)
        if self.reweight_by_anchor:
            weighted_x = weighted_x * self.num_anchors
        x = torch.max(weighted_x, 2, keepdim=True)[0]
        x = x.view(B, 512, self.num_anchors)
        
        if self.per_part_mlp:
            x = x.transpose(1,2).reshape(B, -1, 1)
            m = self.mlp_m(x).reshape(B, self.num_anchors, -1)
            v = self.mlp_v(x).reshape(B, self.num_anchors, -1)
        else:
            if self.use_ln:
                x = x.transpose(1,2)
                m = self.mlp_m(x)
                v = self.mlp_v(x)
            else:
                m = self.mlp_m(x).transpose(1,2)
                v = self.mlp_v(x).transpose(1,2)
        # Returns both mean and logvariance, just ignore the latter in deteministic cases.
        if self.deterministic:
            return m
        return m, v


@ENCODERS.register_module()
class PointNetV3(nn.Module):
    def __init__(self, point_dim=3, zdim=1024, num_anchors=4, mult=1, reweight_by_anchor=True):
        super().__init__()
        self.mult = mult
        self.reweight_by_anchor=reweight_by_anchor
        self.zdim = zdim
        self.num_anchors=num_anchors
        self.conv1 = nn.Conv1d(point_dim, 128 * mult, 1)
        self.conv2 = nn.Conv1d(128* mult, 128* mult, 1)
        self.conv3 = nn.Conv1d(128* mult, 256* mult, 1)
        self.conv4 = nn.Conv1d(256* mult, 256* mult, 1)
        self.bn1 = nn.BatchNorm1d(128* mult)
        self.bn2 = nn.BatchNorm1d(128* mult)
        self.bn3 = nn.BatchNorm1d(256* mult)
        self.bn4 = nn.BatchNorm1d(256* mult)

        # Mapping to [c], cmean
        self.fc1_m = nn.Conv1d(512* mult, 256* mult, 1)
        self.fc2_m = nn.Conv1d(256* mult, 128* mult, 1)
        self.fc3_m = nn.Conv1d(128* mult, zdim, 1)
        self.fc_bn1_m = nn.BatchNorm1d(256* mult)
        self.fc_bn2_m = nn.BatchNorm1d(128* mult)

    def forward(self, x):
        if x.shape[-1] == 3 + self.num_anchors:
            attn_weight = x[..., 3:]
            x = x[..., :3]
        B = x.shape[0]
        x = x.transpose(1, 2)
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = self.bn4(self.conv4(x))
        weighted_x = x.unsqueeze(-1) * attn_weight.unsqueeze(1)
        if self.reweight_by_anchor:
            weighted_x = weighted_x * self.num_anchors
        weighted_x = torch.max(weighted_x, 2, keepdim=True)[0]
        x = torch.max(x, 2, keepdim=True)[0].unsqueeze(-1).expand(-1, -1, -1, self.num_anchors)
        x = torch.cat([x, weighted_x], dim=1).view(B, 512, self.num_anchors)

        m = F.relu(self.fc_bn1_m(self.fc1_m(x)))
        m = F.relu(self.fc_bn2_m(self.fc2_m(m)))
        m = self.fc3_m(m)

        # Returns both mean and logvariance, just ignore the latter in deteministic cases.
        return m.transpose(1,2)
python/difffacto/models/encoders/pointnet.py
diffFacto/diffFacto
import torch.nn as nn
import torch.nn.functional as F
from pointnet2_ops.pointnet2_modules import PointnetSAModuleMSG, PointnetSAModule
from difffacto.utils.registry import ENCODERS


@ENCODERS.register_module()
class PointNet2SSG(nn.Module):
    def __init__(self, additioinal_dim=4, zdim=256, num_anchors=4):
        super().__init__()

        self.additioinal_dim = additioinal_dim
        self.zdim=zdim 
        self.num_anchors=num_anchors

        self._build_model()

    def _build_model(self):
        self.SA_modules = nn.ModuleList()
        self.SA_modules.append(
            PointnetSAModule(
                npoint=512,
                radius=0.2,
                nsample=64,
                mlp=[self.additioinal_dim, 64, 64, 128],
                use_xyz=True,
            )
        )
        self.SA_modules.append(
            PointnetSAModule(
                npoint=128,
                radius=0.4,
                nsample=64,
                mlp=[128, 128, 128, 256],
                use_xyz=True,
            )
        )
        self.SA_modules.append(
            PointnetSAModule(
                mlp=[256, 256, 512, 1024], use_xyz=True
            )
        )

        self.fc_layer = nn.Sequential(
            nn.Linear(1024, 512, bias=False),
            nn.BatchNorm1d(512),
            nn.ReLU(True),
            nn.Linear(512, 256, bias=False),
            nn.BatchNorm1d(256),
            nn.ReLU(True),
            nn.Dropout(0.5),
            nn.Linear(256, self.zdim*self.num_anchors),
        )

    def _break_up_pc(self, pc):
        xyz = pc[..., 0:3].contiguous()
        features = pc[..., 3:].transpose(1, 2).contiguous() if pc.size(-1) > 3 else None

        return xyz, features

    def forward(self, pointcloud):
        r"""
            Forward pass of the network

            Parameters
            ----------
            pointcloud: Variable(torch.cuda.FloatTensor)
                (B, N, 3 + input_channels) tensor
                Point cloud to run predicts on
                Each point in the point-cloud MUST
                be formated as (x, y, z, features...)
        """
        B = pointcloud.shape[0]
        xyz, features = self._break_up_pc(pointcloud)

        for module in self.SA_modules:
            xyz, features = module(xyz, features)

        return self.fc_layer(features.squeeze(-1)).reshape(B, self.num_anchors, self.zdim)

@ENCODERS.register_module()
class PointNet2MSG(PointNet2SSG):
    def _build_model(self):
        super()._build_model()

        self.SA_modules = nn.ModuleList()
        self.SA_modules.append(
            PointnetSAModuleMSG(
                npoint=512,
                radii=[0.1, 0.2, 0.4],
                nsamples=[16, 32, 128],
                mlps=[[self.additioinal_dim, 32, 32, 64], [self.additioinal_dim, 64, 64, 128], [self.additioinal_dim, 64, 96, 128]],
                use_xyz=True,
            )
        )

        input_channels = 64 + 128 + 128
        self.SA_modules.append(
            PointnetSAModuleMSG(
                npoint=128,
                radii=[0.2, 0.4, 0.8],
                nsamples=[32, 64, 128],
                mlps=[
                    [input_channels, 64, 64, 128],
                    [input_channels, 128, 128, 256],
                    [input_channels, 128, 128, 256],
                ],
                use_xyz=True,
            )
        )
        self.SA_modules.append(
            PointnetSAModule(
                mlp=[128 + 256 + 256, 256, 512, 1024],
                use_xyz=True,
            )
        )
python/difffacto/models/encoders/pointnet2.py
diffFacto/diffFacto
import torch
import torch.nn as nn
import torch.nn.functional as F
from time import time
import numpy as np

def timeit(tag, t):
    print("{}: {}s".format(tag, time() - t))
    return time()

def pc_normalize(pc):
    l = pc.shape[0]
    centroid = np.mean(pc, axis=0)
    pc = pc - centroid
    m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
    pc = pc / m
    return pc

def square_distance(src, dst):
    """
    Calculate Euclid distance between each two points.
    src^T * dst = xn * xm + yn * ym + zn * zm；
    sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
    sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
    dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
         = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
    Input:
        src: source points, [B, N, C]
        dst: target points, [B, M, C]
    Output:
        dist: per-point square distance, [B, N, M]
    """
    B, N, _ = src.shape
    _, M, _ = dst.shape
    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
    dist += torch.sum(src ** 2, -1).view(B, N, 1)
    dist += torch.sum(dst ** 2, -1).view(B, 1, M)
    return dist


def index_points(points, idx):
    """
    Input:
        points: input points data, [B, N, C]
        idx: sample index data, [B, S]
    Return:
        new_points:, indexed points data, [B, S, C]
    """
    device = points.device
    B = points.shape[0]
    view_shape = list(idx.shape)
    view_shape[1:] = [1] * (len(view_shape) - 1)
    repeat_shape = list(idx.shape)
    repeat_shape[0] = 1
    batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
    new_points = points[batch_indices, idx, :]
    return new_points


def farthest_point_sample(xyz, npoint):
    """
    Input:
        xyz: pointcloud data, [B, N, 3]
        npoint: number of samples
    Return:
        centroids: sampled pointcloud index, [B, npoint]
    """
    device = xyz.device
    B, N, C = xyz.shape
    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
    distance = torch.ones(B, N).to(device) * 1e10
    farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
    batch_indices = torch.arange(B, dtype=torch.long).to(device)
    for i in range(npoint):
        centroids[:, i] = farthest
        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
        dist = torch.sum((xyz - centroid) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = torch.max(distance, -1)[1]
    return centroids


def query_ball_point(radius, nsample, xyz, new_xyz):
    """
    Input:
        radius: local region radius
        nsample: max sample number in local region
        xyz: all points, [B, N, 3]
        new_xyz: query points, [B, S, 3]
    Return:
        group_idx: grouped points index, [B, S, nsample]
    """
    device = xyz.device
    B, N, C = xyz.shape
    _, S, _ = new_xyz.shape
    group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1])
    sqrdists = square_distance(new_xyz, xyz)
    group_idx[sqrdists > radius ** 2] = N
    group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
    group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample])
    mask = group_idx == N
    group_idx[mask] = group_first[mask]
    return group_idx


def sample_and_group(npoint, radius, nsample, xyz, points, returnfps=False):
    """
    Input:
        npoint:
        radius:
        nsample:
        xyz: input points position data, [B, N, 3]
        points: input points data, [B, N, D]
    Return:
        new_xyz: sampled points position data, [B, npoint, nsample, 3]
        new_points: sampled points data, [B, npoint, nsample, 3+D]
    """
    B, N, C = xyz.shape
    S = npoint
    fps_idx = farthest_point_sample(xyz, npoint) # [B, npoint, C]
    new_xyz = index_points(xyz, fps_idx)
    idx = query_ball_point(radius, nsample, xyz, new_xyz)
    grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C]
    grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)

    if points is not None:
        grouped_points = index_points(points, idx)
        new_points = torch.cat([grouped_xyz_norm, grouped_points], dim=-1) # [B, npoint, nsample, C+D]
    else:
        new_points = grouped_xyz_norm
    if returnfps:
        return new_xyz, new_points, grouped_xyz, fps_idx
    else:
        return new_xyz, new_points


def sample_and_group_all(xyz, points):
    """
    Input:
        xyz: input points position data, [B, N, 3]
        points: input points data, [B, N, D]
    Return:
        new_xyz: sampled points position data, [B, 1, 3]
        new_points: sampled points data, [B, 1, N, 3+D]
    """
    device = xyz.device
    B, N, C = xyz.shape
    new_xyz = torch.zeros(B, 1, C).to(device)
    grouped_xyz = xyz.view(B, 1, N, C)
    if points is not None:
        new_points = torch.cat([grouped_xyz, points.view(B, 1, N, -1)], dim=-1)
    else:
        new_points = grouped_xyz
    return new_xyz, new_points


class PointNetSetAbstraction(nn.Module):
    def __init__(self, npoint, radius, nsample, in_channel, mlp, group_all):
        super(PointNetSetAbstraction, self).__init__()
        self.npoint = npoint
        self.radius = radius
        self.nsample = nsample
        self.mlp_convs = nn.ModuleList()
        self.mlp_bns = nn.ModuleList()
        last_channel = in_channel
        for out_channel in mlp:
            self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
            last_channel = out_channel
        self.group_all = group_all

    def forward(self, xyz, points):
        """
        Input:
            xyz: input points position data, [B, C, N]
            points: input points data, [B, D, N]
        Return:
            new_xyz: sampled points position data, [B, C, S]
            new_points_concat: sample points feature data, [B, D', S]
        """
        xyz = xyz.permute(0, 2, 1)
        if points is not None:
            points = points.permute(0, 2, 1)

        if self.group_all:
            new_xyz, new_points = sample_and_group_all(xyz, points)
        else:
            new_xyz, new_points = sample_and_group(self.npoint, self.radius, self.nsample, xyz, points)
        # new_xyz: sampled points position data, [B, npoint, C]
        # new_points: sampled points data, [B, npoint, nsample, C+D]
        new_points = new_points.permute(0, 3, 2, 1) # [B, C+D, nsample,npoint]
        for i, conv in enumerate(self.mlp_convs):
            bn = self.mlp_bns[i]
            new_points =  F.relu(bn(conv(new_points)))

        new_points = torch.max(new_points, 2)[0]
        new_xyz = new_xyz.permute(0, 2, 1)
        return new_xyz, new_points


class PointNetSetAbstractionMsg(nn.Module):
    def __init__(self, npoint, radius_list, nsample_list, in_channel, mlp_list):
        super(PointNetSetAbstractionMsg, self).__init__()
        self.npoint = npoint
        self.radius_list = radius_list
        self.nsample_list = nsample_list
        self.conv_blocks = nn.ModuleList()
        self.bn_blocks = nn.ModuleList()
        for i in range(len(mlp_list)):
            convs = nn.ModuleList()
            bns = nn.ModuleList()
            last_channel = in_channel + 3
            for out_channel in mlp_list[i]:
                convs.append(nn.Conv2d(last_channel, out_channel, 1))
                bns.append(nn.BatchNorm2d(out_channel))
                last_channel = out_channel
            self.conv_blocks.append(convs)
            self.bn_blocks.append(bns)

    def forward(self, xyz, points):
        """
        Input:
            xyz: input points position data, [B, C, N]
            points: input points data, [B, D, N]
        Return:
            new_xyz: sampled points position data, [B, C, S]
            new_points_concat: sample points feature data, [B, D', S]
        """
        xyz = xyz.permute(0, 2, 1)
        if points is not None:
            points = points.permute(0, 2, 1)

        B, N, C = xyz.shape
        S = self.npoint
        new_xyz = index_points(xyz, farthest_point_sample(xyz, S))
        new_points_list = []
        for i, radius in enumerate(self.radius_list):
            K = self.nsample_list[i]
            group_idx = query_ball_point(radius, K, xyz, new_xyz)
            grouped_xyz = index_points(xyz, group_idx)
            grouped_xyz -= new_xyz.view(B, S, 1, C)
            print(grouped_xyz.shape)
            if points is not None:
                grouped_points = index_points(points, group_idx)
                grouped_points = torch.cat([grouped_points, grouped_xyz], dim=-1)
            else:
                grouped_points = grouped_xyz

            grouped_points = grouped_points.permute(0, 3, 2, 1)  # [B, D, K, S]
            for j in range(len(self.conv_blocks[i])):
                conv = self.conv_blocks[i][j]
                bn = self.bn_blocks[i][j]
                grouped_points =  F.relu(bn(conv(grouped_points)))
            new_points = torch.max(grouped_points, 2)[0]  # [B, D', S]
            new_points_list.append(new_points)

        new_xyz = new_xyz.permute(0, 2, 1)
        new_points_concat = torch.cat(new_points_list, dim=1)
        return new_xyz, new_points_concat


class PointNetFeaturePropagation(nn.Module):
    def __init__(self, in_channel, mlp):
        super(PointNetFeaturePropagation, self).__init__()
        self.mlp_convs = nn.ModuleList()
        self.mlp_bns = nn.ModuleList()
        last_channel = in_channel
        for out_channel in mlp:
            self.mlp_convs.append(nn.Conv1d(last_channel, out_channel, 1))
            self.mlp_bns.append(nn.BatchNorm1d(out_channel))
            last_channel = out_channel

    def forward(self, xyz1, xyz2, points1, points2):
        """
        Input:
            xyz1: input points position data, [B, C, N]
            xyz2: sampled input points position data, [B, C, S]
            points1: input points data, [B, D, N]
            points2: input points data, [B, D, S]
        Return:
            new_points: upsampled points data, [B, D', N]
        """
        xyz1 = xyz1.permute(0, 2, 1)
        xyz2 = xyz2.permute(0, 2, 1)

        points2 = points2.permute(0, 2, 1)
        B, N, C = xyz1.shape
        _, S, _ = xyz2.shape

        if S == 1:
            interpolated_points = points2.repeat(1, N, 1)
        else:
            dists = square_distance(xyz1, xyz2)
            dists, idx = dists.sort(dim=-1)
            dists, idx = dists[:, :, :3], idx[:, :, :3]  # [B, N, 3]

            dist_recip = 1.0 / (dists + 1e-8)
            norm = torch.sum(dist_recip, dim=2, keepdim=True)
            weight = dist_recip / norm
            interpolated_points = torch.sum(index_points(points2, idx) * weight.view(B, N, 3, 1), dim=2)

        if points1 is not None:
            points1 = points1.permute(0, 2, 1)
            new_points = torch.cat([points1, interpolated_points], dim=-1)
        else:
            new_points = interpolated_points

        new_points = new_points.permute(0, 2, 1)
        for i, conv in enumerate(self.mlp_convs):
            bn = self.mlp_bns[i]
            new_points = F.relu(bn(conv(new_points)))
        return new_points
python/difffacto/models/encoders/pointnet2_utils.py
diffFacto/diffFacto
import torch 
import torch.nn as nn
import torch.nn.functional as F
from difffacto.utils.registry import ENCODERS
from difffacto.models.modules import ConditionalBatchNorm1d
from difffacto.utils.misc import timestep_embedding

@ENCODERS.register_module()
class SupSegsEncoder(nn.Module):
    def __init__(self, sup_segs_dim, part_latent_dim, num_anchors):
        super().__init__()
        dim = sup_segs_dim
        self.num_anchors = num_anchors
        self.sup_seg_enc = PartglotSupSegsEncoder(dim)

        self.first_conv = nn.Sequential(
            nn.Conv1d(dim, 128, 1),
            nn.BatchNorm1d(128),
            nn.ReLU(inplace=True),
            nn.Conv1d(128, 256, 1)
        )

        self.second_conv = nn.Sequential(
            nn.Conv1d(512, 512, 1),
            nn.BatchNorm1d(512),
            nn.ReLU(inplace=True),
            nn.Conv1d(512, 1024, 1)
        )

        self.latent_mlp = nn.Sequential(
            nn.Linear(1024, 1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, 1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, part_latent_dim * num_anchors)
        )

    def forward(self, x, mask):
        neg_inf_mask = torch.logical_not(mask) * -10000
        B, n_seg, npoint, _ = x.shape
        x = x.reshape(B*n_seg, npoint, -1)
        out = self.sup_seg_enc(x) # (B * n_seg,  64, 1)
        feature_super_seg = torch.max(out, dim=1)[0].reshape(B, n_seg, -1).transpose(1,2)                       # (B, 64, n_seg)
        feature = self.first_conv(feature_super_seg)
        masked_feature = feature + neg_inf_mask.unsqueeze(1)
        feature_global = torch.max(masked_feature, dim=2, keepdim=True)[0]
        masked_feature = self.second_conv(torch.cat([feature, feature_global.expand(-1, -1, n_seg)], dim=1)) + neg_inf_mask.unsqueeze(1)
        feature_global = torch.max(masked_feature, dim=2, keepdim=False)[0]

        part_latent = self.latent_mlp(feature_global).reshape(B, self.num_anchors, -1)
        

        return part_latent

class PartglotSupSegsEncoder(nn.Module):
    def __init__(self, sup_segs_dim):
        super().__init__()
        dim = sup_segs_dim
        self.conv1 = nn.Sequential(nn.Conv1d(3, dim, 1), nn.BatchNorm1d(dim))
        self.conv2 = nn.Sequential(nn.Conv1d(dim, dim, 1), nn.BatchNorm1d(dim))
        self.conv3 = nn.Sequential(nn.Conv1d(dim, dim, 1), nn.BatchNorm1d(dim))
        self.conv4 = nn.Sequential(nn.Conv1d(dim, dim, 1), nn.BatchNorm1d(dim))
        self.fc = nn.Conv1d(dim, dim, 1)

    def forward(self, x):
        x = x.transpose(1, 2)
        out1 = F.relu(self.conv1(x))
        out2 = F.relu(self.conv2(out1))
        out3 = F.relu(self.conv3(out2))
        out4 = F.relu(self.conv4(out3))

        out = self.fc(out4)
        out = out.transpose(1, 2)

        return out

@ENCODERS.register_module()
class PartglotSupSegsEncoderWithCBN(nn.Module):
    def __init__(self, sup_segs_dim, embed_t_size=128, emb_size=256):
        super().__init__()
        dim = sup_segs_dim
        self.conv1 = nn.Conv1d(3, dim, 1)
        self.bn1=ConditionalBatchNorm1d(dim, embed_t_size, emb_size)
        self.conv2 = nn.Conv1d(dim, dim, 1)
        self.bn2 = ConditionalBatchNorm1d(dim, embed_t_size, emb_size)
        self.conv3 = nn.Conv1d(dim, dim, 1)
        self.bn3 = ConditionalBatchNorm1d(dim, embed_t_size, emb_size)
        self.conv4 = nn.Conv1d(dim, dim, 1)
        self.bn4 = ConditionalBatchNorm1d(dim, embed_t_size, emb_size)
        self.fc = nn.Conv1d(dim, dim, 1)

    def forward(self, x, t):
        x = x.transpose(1, 2)
        out1 = F.relu(self.bn1(self.conv1(x)), t)
        out2 = F.relu(self.bn2(self.conv2(out1), t))
        out3 = F.relu(self.bn3(self.conv3(out2), t))
        out4 = F.relu(self.bn4(self.conv4(out3), t))

        out = self.fc(out4)
        out = out.transpose(1, 2)

        return out
python/difffacto/models/encoders/super_seg_encoder.py
diffFacto/diffFacto
import torch
import torch.nn as nn
import torch.nn.functional as F
from difffacto.utils.registry import GENERATORS


# Taken from https://discuss.pytorch.org/t/implementing-truncated-normal-initializer/4778/15
def truncated_normal(tensor, mean=0, std=1, trunc_std=2):
    size = tensor.shape
    tmp = tensor.new_empty(size + (4,)).normal_()
    valid = (tmp < trunc_std) & (tmp > -trunc_std)
    ind = valid.max(-1, keepdim=True)[1]
    tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1))
    tensor.data.mul_(std).add_(mean)
    return tensor

@GENERATORS.register_module()
class Generator(nn.Module):

    def __init__(self, inp_dim, out_dim, use_bn, output_bn, dims, prior='gaussian', gaussian_scale=1.0):
        super().__init__()


        self.inp_dim = inp_dim
        self.out_dim = out_dim
        self.use_bn = use_bn
        self.output_bn = output_bn
        self.dims = dims
        self.gaussian_scale=gaussian_scale
        curr_dim = self.inp_dim
        self.layers = []
        self.bns = []
        for hid in self.dims:
            self.layers.append(nn.Linear(curr_dim, hid))
            self.bns.append(nn.BatchNorm1d(hid))
            curr_dim = hid
        self.layers = nn.ModuleList(self.layers)
        self.bns = nn.ModuleList(self.bns)
        self.out = nn.Linear(curr_dim, self.out_dim)
        self.out_bn = nn.BatchNorm1d(self.out_dim)
        self.prior_type = prior

    def get_prior(self, bs):
        if self.prior_type == "truncate_gaussian":
            noise = (torch.randn(bs, self.inp_dim) * self.gaussian_scale).cuda()
            noise = truncated_normal(
                noise, mean=0, std=self.gaussian_scale, trunc_std=2.)
            return noise
        elif self.prior_type == "gaussian":
            return torch.randn(bs, self.inp_dim).cuda() * self.gaussian_scale
        else:
            raise NotImplementedError(
                "Invalid prior type:%s" % self.prior_type)

    def forward(self, z=None, bs=None):
        if z is None:
            assert bs is not None
            z = self.get_prior(bs).cuda()

        y = z
        for layer, bn in zip(self.layers, self.bns):
            y = layer(y)
            if self.use_bn:
                y = bn(y)
            y = F.relu(y)
        y = self.out(y)

        if self.output_bn:
            y = self.out_bn(y)
        return y


python/difffacto/models/generators/generator.py
diffFacto/diffFacto
import numpy as np

import torch


def normal_kl(mean1, logvar1, mean2, logvar2):
    """
    Compute the KL divergence between two gaussians.
    Shapes are automatically broadcasted, so batches can be compared to
    scalars, among other use cases.
    """
    tensor = None
    for obj in (mean1, logvar1, mean2, logvar2):
        if isinstance(obj, torch.Tensor):
            tensor = obj
            break
    assert tensor is not None, "at least one argument must be a Tensor"

    # Force variances to be Tensors. Broadcasting helps convert scalars to
    # Tensors, but it does not work for th.exp().
    logvar1, logvar2 = [
        x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor)
        for x in (logvar1, logvar2)
    ]

    return 0.5 * (
        -1.0
        + logvar2
        - logvar1
        + torch.exp(logvar1 - logvar2)
        + ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
    )
python/difffacto/models/losses/losses.py
diffFacto/diffFacto
from .cbn import ConditionalBatchNorm1d
python/difffacto/models/modules/__init__.py
diffFacto/diffFacto
import torch
import torch.nn as nn
from difffacto.utils.misc import timestep_embedding


'''
CBN (Conditional Batch Normalization layer)
    uses an MLP to predict the beta and gamma parameters in the batch norm equation
    Reference : https://papers.nips.cc/paper/7237-modulating-early-visual-processing-by-language.pdf
'''
class ConditionalBatchNorm1d(nn.Module):

    def __init__(self, channels, embed_t_size=128, emb_size=256, use_betas=True, use_gammas=True, eps=1.0e-5):
        super(ConditionalBatchNorm1d, self).__init__()

        self.embed_t_size = embed_t_size # size of the lstm emb which is input to MLP
        self.emb_size = emb_size # size of hidden layer of MLP
        self.use_betas = use_betas
        self.use_gammas = use_gammas

        self.channels = channels

        # beta and gamma parameters for each channel - defined as trainable parameters
        self.betas = nn.Parameter(torch.zeros(self.channels))
        self.gammas = nn.Parameter(torch.ones(self.channels))
        self.betas.requires_grad = False
        self.gammas.requires_grad = False
        self.eps = eps

        # MLP used to predict betas and gammas
        self.fc_gamma = nn.Sequential(
          nn.Linear(self.embed_t_size, self.emb_size),
          nn.ReLU(inplace=True),
          nn.Linear(self.emb_size, self.channels),
        )

        self.fc_beta = nn.Sequential(
          nn.Linear(self.embed_t_size, self.emb_size),
          nn.ReLU(inplace=True),
          nn.Linear(self.emb_size, self.channels),
        )

        # initialize weights using Xavier initialization and biases with constant value
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform(m.weight)
                nn.init.constant(m.bias, 0.1)

    '''
    Predicts the value of delta beta and delta gamma for each channel
    Arguments:
        lstm_emb : lstm embedding of the question
    Returns:
        delta_betas, delta_gammas : for each layer
    '''
    def create_cbn_input(self, t):
        t_embed = timestep_embedding(t, self.embed_t_size)
        if self.use_betas:
            delta_betas = self.fc_beta(t_embed)
        else:
            delta_betas = torch.zeros(self.channels).to(t.device)

        if self.use_gammas:
            delta_gammas = self.fc_gamma(t_embed)
        else:
            delta_gammas = torch.zeros(self.channels).to(t.device)

        return delta_betas, delta_gammas

    '''
    Computer Normalized feature map with the updated beta and gamma values
    Arguments:
        feature : feature map from the previous layer
        lstm_emb : lstm embedding of the question
    Returns:
        out : beta and gamma normalized feature map
        lstm_emb : lstm embedding of the question (unchanged)
    Note : lstm_emb needs to be returned since CBN is defined within nn.Sequential
           and subsequent CBN layers will also require lstm question embeddings
    '''
    def forward(self, feature, t):
        # get delta values
        delta_betas, delta_gammas = self.create_cbn_input(t)

        betas_cloned = self.betas.clone()
        gammas_cloned = self.gammas.clone()

        # update the values of beta and gamma
        betas_cloned += delta_betas
        gammas_cloned += delta_gammas

        # get the mean and variance for the batch norm layer
        batch_mean = torch.mean(feature, dim=(0,2), keepdim=True)
        batch_var = torch.var(feature, dim=(0,2), keepdim=True)

        # extend the betas and gammas of each channel across the height and width of feature map

        # normalize the feature map
        feature_normalized = (feature-batch_mean)/torch.sqrt(batch_var+self.eps)

        # get the normalized feature map with the updated beta and gamma values
        out = feature_normalized, gammas_cloned.unsqueeze(2) + betas_cloned.unsqueeze(2)

        return out
python/difffacto/models/modules/cbn.py
diffFacto/diffFacto
import torch
import torch.nn as nn
from difffacto.utils.registry import MODELS, METRICS, build_from_cfg

@MODELS.register_module()
class PCN(nn.Module):
    def __init__(self, num_pred, encoder_channel, loss):
        super().__init__()
        self.number_fine = num_pred
        self.encoder_channel = encoder_channel
        grid_size = 4 # set default
        self.grid_size = grid_size
        assert self.number_fine % grid_size**2 == 0
        self.number_coarse = self.number_fine // (grid_size ** 2 )
        self.first_conv = nn.Sequential(
            nn.Conv1d(3,128,1),
            nn.BatchNorm1d(128),
            nn.ReLU(inplace=True),
            nn.Conv1d(128,256,1)
        )
        self.second_conv = nn.Sequential(
            nn.Conv1d(512,512,1),
            nn.BatchNorm1d(512),
            nn.ReLU(inplace=True),
            nn.Conv1d(512,self.encoder_channel,1)
        )
        self.mlp = nn.Sequential(
            nn.Linear(self.encoder_channel,1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024,1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024,3*self.number_coarse)
        )
        self.final_conv = nn.Sequential(
            nn.Conv1d(1024+3+2,512,1),
            nn.BatchNorm1d(512),
            nn.ReLU(inplace=True),
            nn.Conv1d(512,512,1),
            nn.BatchNorm1d(512),
            nn.ReLU(inplace=True),
            nn.Conv1d(512,3,1)
        )
        a = torch.linspace(-0.05, 0.05, steps=grid_size, dtype=torch.float).view(1, grid_size).expand(grid_size, grid_size).reshape(1, -1)
        b = torch.linspace(-0.05, 0.05, steps=grid_size, dtype=torch.float).view(grid_size, 1).expand(grid_size, grid_size).reshape(1, -1)
        self.folding_seed = torch.cat([a, b], dim=0).view(1, 2, grid_size ** 2).cuda() # 1 2 S
        self.loss_func = build_from_cfg(loss, METRICS)

    def get_loss(self, coarse, fine, gt):
        loss_coarse = self.loss_func(coarse, gt)
        loss_fine = self.loss_func(fine, gt)
        return dict(loss_coarse=loss_coarse, loss_fine=loss_fine)

    def forward(self, gt, xyz, device):
        bs , n , _ = xyz.shape
        # encoder
        feature = self.first_conv(xyz.transpose(2,1))  # B 256 n
        feature_global = torch.max(feature,dim=2,keepdim=True)[0]  # B 256 1
        feature = torch.cat([feature_global.expand(-1,-1,n), feature], dim=1)# B 512 n
        feature = self.second_conv(feature) # B 1024 n
        feature_global = torch.max(feature,dim=2,keepdim=False)[0] # B 1024
        # decoder
        coarse = self.mlp(feature_global).reshape(-1,self.number_coarse,3) # B M 3
        
        point_feat = coarse.unsqueeze(2).expand(-1,-1,self.grid_size**2,-1) # B M S 3
        point_feat = point_feat.reshape(-1,self.number_fine,3).transpose(2,1) # B 3 N

        seed = self.folding_seed.unsqueeze(2).expand(bs,-1,self.number_coarse, -1) # B 2 M S
        seed = seed.reshape(bs,-1,self.number_fine)  # B 2 N

        feature_global = feature_global.unsqueeze(2).expand(-1,-1,self.number_fine) # B 1024 N
        feat = torch.cat([feature_global, seed, point_feat], dim=1) # B C N
        
        fine = self.final_conv(feat) + point_feat   # B 3 N
        
        if self.training:
            return self.get_loss(coarse.contiguous(), fine.transpose(2,1).contiguous(), gt)
        return dict(coarse_pred=coarse.contiguous(), pred=fine.transpose(2,1).contiguous(), ref=gt, input=xyz)

python/difffacto/models/networks/PCN.py
diffFacto/diffFacto
from .completion import DiffuCompletion
from .anchor_gen import AnchorDiffAE
from .PCN import PCN
from .coarsePCN import CoarsePCN
from .anchor_gen_supseg import AnchorDiffGenSuperSegments
from .anchor_gen_supseg_partglot import AnchorDiffGenPartglot
from .listener import PartglotClassifier
from .part_editing import PartLatentEditor
from .anchor_gen_uncond import AnchoredGeneration
from .latent_diffusion import LatentDiffEdit
from .point_diff import PDM
from .flow import PointFlow
from .anchor_gen_parameter_trainer import ParameterTrainer
python/difffacto/models/networks/__init__.py
diffFacto/diffFacto
import torch
import torch.nn as nn 
import torch.nn.functional as F
from torch.nn import Module, functional
from difffacto.models.diffusions.nets.attention import MLP
from difffacto.utils.misc import *
from difffacto.utils.registry import build_from_cfg, MODELS, ENCODERS, DIFFUSIONS, SAMPLERS, METRICS, DECOMPOSERS, SEGMENTORS, GENERATORS, DISCRIMINATORS
from difffacto.metrics.common import gradient_penalty, gen_loss, dis_loss, triplet_loss
from pointnet2_ops.pointnet2_utils import gather_operation
from .language_utils.language_util import tokenizing
from einops import rearrange
VOCAB_SIZE=2787

def Normalize(in_channels):
    return nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)

def identity(*args, **kwargs):
    return nn.Identity()

def zero_module(module):
    """
    Zero out the parameters of a module and return it.
    """
    for p in module.parameters():
        p.detach().zero_()
    return module


@MODELS.register_module()
class AnchorDiffAE(Module):

    def __init__(
        self, 
        encoder, 
        diffusion, 
        sampler, 
        num_anchors,
        num_timesteps, 
        npoints=2048, 
        zero_anchors=False,
        gen=False,
        sample_noise_num=20,
        cimle=False,
        cimle_sample_num=10,
        diffusion_loss_weight=1.0,
        use_input=False,


        learn_var=False,
        detach_variance=True,
        detach_anchor=True,
        
        global_shift=False,
        global_scale=False,
        vertical_only=True,

        ret_traj=False, 
        ret_interval=20, 
        forward_sample=False,
        interpolate=False,
        interpolate_part_id=2,
        fix_part_ids = None,
        combine=False,
        drift_anchors=False,
        save_pred_xstart=False,
        save_dir=None,
        save_weights=False,
        
        
        noise_reg_loss=True,
        reg_loss_weight=1.0,
        
        pretrain_prior=False,
        train_language=False,
        language_encoder=None,
        clip_weight=1.0,
        triplet_weight=1.0,
        triplet_thresh=0.1,
        
        ):
        super().__init__()
        self.pretrain_prior=pretrain_prior
        self.encoder = build_from_cfg(encoder, ENCODERS)
        if pretrain_prior:
            self.diffusion = nn.ModuleList([build_from_cfg(diffusion, DIFFUSIONS, num_timesteps=num_timesteps) for _ in range(num_anchors)])
        else:
            self.diffusion = build_from_cfg(diffusion, DIFFUSIONS, num_timesteps=num_timesteps)
        self.sampler = build_from_cfg(sampler, SAMPLERS, num_timesteps=num_timesteps)
        self.diffusion_loss_weight=diffusion_loss_weight
        self.sample_noise_num=sample_noise_num
        self.interpolate_part_id=interpolate_part_id
        self.cimle=cimle
        self.cimle_sample_num=cimle_sample_num
        self.fix_part_ids=fix_part_ids
        self.learn_var = learn_var
        self.gen=gen
        self.use_input=use_input
        self.save_weights=save_weights
        self.save_dir=save_dir
        self.global_shift = global_shift
        self.global_scale=global_scale
        self.vertical_only=vertical_only
        self.zero_anchors=zero_anchors

        self.drift_anchors=drift_anchors
        self.num_timesteps = int(num_timesteps)
        self.num_anchors = num_anchors
        
        self.detach_anchor=detach_anchor
        self.detach_variance=detach_variance
        self.fixed_id= [0, 0, 0, 0]
        assert len(self.fixed_id) == self.num_anchors
        
        self.train_language=train_language
        if train_language:
            self.language_encoder = build_from_cfg(language_encoder, ENCODERS)
            self.update_mlp = zero_module(MLP([self.language_encoder.text_dim + self.encoder.zdim, self.encoder.zdim * 2, self.encoder.zdim*2, self.encoder.zdim], use_linear=True))
            self.temp = nn.Parameter(torch.zeros(1))
            self.clip_weight=clip_weight
            self.triplet_weight=triplet_weight
            self.triplet_thresh=triplet_thresh
        
        self.npoints = npoints
        self.points_per_anchor = self.npoints // self.num_anchors
        self.forward_sample = forward_sample
        self.interpolate=interpolate
        self.ret_traj=ret_traj
        self.ret_interval=ret_interval
        self.save_pred_xstart=save_pred_xstart
        self.combine=combine
        self.noise_reg_loss=noise_reg_loss
        self.reg_loss_weight=reg_loss_weight
        
        # lgan training
        print(self)

    def train_language_module(self):
        self.language_encoder.train()
        self.update_mlp.train()
        
    def eval_language_module(self):
        self.language_encoder.eval()
        self.update_mlp.eval()
        
    def decode(self, anchors, ctx=None, noise=None, variance=None, anchor_assignments=None, valid_id=None, device='cuda'):
        final = dict() 
        np = anchor_assignments.shape[1]
        bs = anchors.shape[0]
        for t, sample in self.diffusion.p_sample_loop_progressive(
            [bs, 3, np],
            anchors=anchors,
            variance=variance,
            ctx=ctx,
            noise=noise,
            anchor_assignment=anchor_assignments,
            valid_id=valid_id,
            device=device,
            progress=False,
        ):
            if t == 0:
                final['pred'] = sample['sample'].transpose(2, 1)
                if self.save_pred_xstart:
                    final['pred_xstart'] = sample['pred_xstart'].transpose(1,2) 
            elif (self.ret_traj and t % self.ret_interval == 0):
                final[t] = sample['sample'].transpose(2, 1) 
                if self.save_pred_xstart:
                    final[f'pred_xstart_{t}'] = sample['pred_xstart'].transpose(1,2)
                
        return final

    def q_sample(self, gt, anchors, noise=None, device='cuda',variance=None):
        final = {'pred':gt.transpose(2,1)}
        
        for t, sample in self.diffusion.q_sample_loop_progressive(
            gt_pcd=gt,
            anchors=anchors,
            noise=noise,
            device=device,
            variance=variance,
            progress=False,
        ):
            if t == self.num_timesteps - 1:
                final[t] = sample.transpose(2, 1)
            elif (self.ret_traj and t % self.ret_interval == 0):
                final[t] = sample.transpose(2, 1)
                
        return final
    @torch.no_grad()
    def interpolate_two_shapes(self, pcds1, pcds2, part_id, param_shift, param_scale, mid_num=10):
        B = pcds1.shape[0]
        xyz1, mask1 = torch.split(pcds1, [3,1], dim=-1)
        xyz2, mask2 = torch.split(pcds2, [3,1], dim=-1)
        # import pickle
        # pickle.dump(dict(pred=xyz1.cpu().numpy(), pred_seg_mask=mask1.cpu().numpy()), open("/orion/u/w4756677/diffusion/anchorDIff/weights/interpolation/out.pkl", "wb"))
        # exit()
        mask1 = mask1.reshape(B, -1).long()
        mask2 = mask2.reshape(B, -1).long()
        valid_id = torch.zeros(B, 4).cuda()
        for i in range(4):
            m = torch.any(mask1 == i, dim=1)
            valid_id[m] = 1
        anchor_per_point, ctx, variance_per_point, anchor_assignments, valid_id = self.encoder.interpolate_two_shape(xyz1, mask1, xyz2, mask2, self.npoints, part_id, param_shift, param_scale, valid_id, mid_num)
        pred = self.decode(anchor_per_point, ctx=ctx, device='cuda', variance=variance_per_point, anchor_assignments=anchor_assignments, valid_id=valid_id) 
        return pred['pred'].reshape(B, mid_num, self.npoints, 3), anchor_assignments.reshape(B, mid_num, self.npoints)
        
    def interpolate_latent(self, device, pcds):
        out_dict = {}
        ref = pcds['ref'].to(device)
        input = pcds['input'].to(device)
        ref_seg_mask = pcds['ref_seg_mask'].to(device)
        seg_flag = pcds['attn_map'].to(device)
        valid_id = pcds['present'].to(device)
        gt_part_shift = pcds.get('part_shift', torch.zeros_like(ref).transpose(1,2))
        gt_part_var = pcds.get('part_scale', torch.ones_like(ref).transpose(1,2))
        B = input.shape[0]
        if not self.encoder.origin_scale:
            gt_part_var = gt_part_var ** 2
        if self.cimle:
            noise, _ = self.encoder.sample_noise(pcds, device, 1)
            noise = noise.squeeze(1)
        else:
            noise = None
        if self.gen:
            part_code = torch.randn(B, self.encoder.zdim, self.num_anchors).to(device) * math.sqrt(self.encoder.prior_var)
            if self.encoder.use_flow:
                part_code_ = []
                for i in range(self.num_anchors):
                    _part_code = part_code[..., i]
                    part_code_.append(self.encoder.flow[i](_part_code, reverse=True).view(B, self.encoder.zdim))
                part_code = torch.stack(part_code_, dim=-1)
            valid_id[..., self.interpolate_part_id] = 1.
            pred_seg_mask = torch.arange(self.num_anchors).cuda()[None] * valid_id + torch.argmax(valid_id, dim=1, keepdims=True) * (1 - valid_id)
            pred_seg_mask = pred_seg_mask.repeat_interleave(self.npoints // self.num_anchors, dim=1)
        else:
            part_code_means, part_code_logvars = self.encoder.get_part_code(input, seg_flag)
            part_code = reparameterize_gaussian(mean=part_code_means, logvar=part_code_logvars).transpose(1,2)
            pred_seg_mask = ref_seg_mask.reshape(-1, ref_seg_mask.shape[1], 1).expand(-1, -1, self.npoints // ref_seg_mask.shape[1]).reshape(-1, self.npoints)
            
        K = 10
        dx = torch.linspace(0, 1, steps=K).reshape(1, -1, 1).to(device)
        part_code_perm = torch.randperm(B)
        permuted_ref = ref[part_code_perm]
        permuted_ref_seg_mask = ref_seg_mask[part_code_perm]
        interpolated_part_code = part_code[..., self.interpolate_part_id].unsqueeze(1) + (part_code[part_code_perm,:, self.interpolate_part_id].unsqueeze(1) - part_code[..., self.interpolate_part_id].unsqueeze(1)) * dx
        part_code = part_code.unsqueeze(1).repeat_interleave(K, dim=1)
        part_code[..., self.interpolate_part_id] = interpolated_part_code
        part_code = part_code.reshape(B*K, -1, self.num_anchors)
        gt_part_shift = gt_part_shift.repeat_interleave(K, dim=0)
        gt_part_var = gt_part_var.repeat_interleave(K, dim=0)
        noise = noise.repeat_interleave(K, dim=0)
        valid_id = valid_id.repeat_interleave(K, dim=0)
        mean, logvar = self.encoder.get_params_from_part_code(part_code, valid_id, gt_mean=gt_part_shift, gt_var=gt_part_var, noise=noise)
            
        anchor_assignments = pred_seg_mask.repeat_interleave(K, dim=0).int()
        anchor_per_point, logvar_per_point, _ = self.encoder.gather_all(anchor_assignments, anchors=mean, variances=logvar)
        ctx = self.encoder.prepare_ctx(part_code, mean, logvar, anchor_assignments=anchor_assignments)
        
        N = anchor_assignments.shape[1]
        if self.npoints < N:
            anchor_per_point, logvar_per_point  = map(lambda r: r[..., :self.npoints], [anchor_per_point, logvar_per_point])
            ctx = [r[..., self.npoints] if r.shape[-1] > self.npoints else r for r in ctx]
            ref = ref[:, :self.npoints]
            anchor_assignments = anchor_assignments[:, :self.npoints]
        variance_per_point = torch.exp(logvar_per_point)
        priors = torch.randn_like(variance_per_point).transpose(1,2).to(variance_per_point.device) * variance_per_point.transpose(1,2) + anchor_per_point.transpose(1,2)
        from tqdm import tqdm
        from collections import defaultdict
        pred = defaultdict(list)
        chunk = 50
        # B = 5
        for k in tqdm(range(0, B*K, chunk)):
            anchor_per_point_chunk = anchor_per_point[k: k+chunk]
            ctx_chunk = [c[k: k+chunk] for c in ctx]
            variance_per_point_chunk = variance_per_point[k: k+chunk]
            anchor_assignments_chunk = anchor_assignments[k: k+chunk]
            valid_id_chunk = valid_id[k: k+chunk]
            # import ipdb; ipdb.set_trace()
            pred_chunk = self.decode(anchor_per_point_chunk, ctx=ctx_chunk, device=device, variance=variance_per_point_chunk, anchor_assignments=anchor_assignments_chunk, valid_id=valid_id_chunk) 
            pred_chunk = {key: v.detach().cpu() for key, v in pred_chunk.items()}
            for key in pred_chunk.keys():
                pred[key].append(pred_chunk[key].detach().cpu())
        # ref = torch.zeros([self.npoints, 3]).cuda()
        # _ref = torch.cat([gt1[seg_mask[id1] != anchor_id, :3], gt2[seg_mask[id2] == anchor_id, :3]], dim=0)
        # ref[:_ref.shape[0]] = _ref
        # seg_mask_ref = torch.zeros_like(seg_mask[0])
        # _seg_mask_ref = torch.cat([seg_mask[id1][seg_mask[id1] != anchor_id], seg_mask[id2][seg_mask[id2] == anchor_id]], dim=0)
        # seg_mask_ref[:_seg_mask_ref.shape[0]] = _seg_mask_ref
        pred = {key: torch.cat(v) for key, v in pred.items()}
        _pred = pred['pred'].reshape(B, K, self.npoints, 3)
        # priors = priors.reshape(B, K, self.npoints, 3)
        
        for i in range(K):
            out_dict[f"interpolate sample {i}"] = _pred[:,i]
        # for i in range(K):
        #     out_dict[f"interpolate sample prior {i}"] = priors[:, i]
            
        out_dict['pred_seg_mask'] = pred_seg_mask[:B]
        out_dict['ref_seg_mask'] = ref_seg_mask
        out_dict['pred'] = out_dict["interpolate sample 0"]
        out_dict['input_ref'] = ref
        out_dict['permuted_ref'] = permuted_ref
        out_dict['permuted_ref_seg_mask'] = permuted_ref_seg_mask
        out_dict['shift'] = pcds['shift']
        out_dict['scale'] = pcds['scale']
        return out_dict
    
    def sample_one_part(self, code, valid_id, mean, logvar, seg_mask, part_id, sample_num_each, fix_size, param_sample_num, selective_param_sample): 
        bs = code.shape[0]
        np = seg_mask.shape[1]
        ctx, \
        mean_per_point, \
        logvar_per_point, \
        seg_mask, \
        valid_id, \
        latents = \
            self.encoder.sample_with_fixed_latents(
            code,
            valid_id,
            mean, 
            logvar,
            seg_mask, 
            part_id, 
            sample_num_each,
            fix_size, 
            param_sample_num,
            selective_param_sample
        )
        codes, noise_latents, means, logvars = latents
        variance_per_point = torch.exp(logvar_per_point)
        pred = self.decode(mean_per_point, ctx=ctx, device='cuda', variance=variance_per_point, anchor_assignments=seg_mask.to(torch.int32), valid_id=valid_id) 
        pred = pred['pred'].reshape(bs, sample_num_each, param_sample_num, np, 3)
        seg_mask = seg_mask.reshape(bs, sample_num_each, param_sample_num, np)
        valid_id = valid_id.reshape(bs, sample_num_each, param_sample_num, self.num_anchors)
        codes = codes.reshape(bs, sample_num_each, param_sample_num, self.encoder.zdim, self.num_anchors)
        means = means.reshape(bs, sample_num_each, param_sample_num, 3, self.num_anchors)
        logvars = logvars.reshape(bs, sample_num_each, param_sample_num, 3, self.num_anchors)
        return pred, seg_mask, valid_id, codes, noise_latents, means, logvars
    def interpolate_params(self, device, pcds):
        out_dict = {}
        ref = pcds['ref'].to(device)
        input = pcds['input'].to(device)
        ref_seg_mask = pcds['ref_seg_mask'].to(device)
        seg_flag = pcds['attn_map'].to(device)
        valid_id = pcds['present'].to(device)
        gt_part_shift = pcds.get('part_shift', torch.zeros_like(ref).transpose(1,2))
        gt_part_var = pcds.get('part_scale', torch.ones_like(ref).transpose(1,2))
        B = input.shape[0]
        if not self.encoder.origin_scale:
            gt_part_var = gt_part_var ** 2
        if self.cimle:
            noise, _ = self.encoder.sample_noise(pcds, device, 1)
            noise = noise.squeeze(1)
        else:
            noise = None
        part_code_means, part_code_logvars = self.encoder.get_part_code(input, seg_flag)
        if self.encoder.gen:
            part_code = reparameterize_gaussian(mean=part_code_means, logvar=part_code_logvars).transpose(1,2)  # (B, F)
        else:
            part_code = part_code_means.transpose(1,2)
            
        K = self.cimle_sample_num
        dx = torch.linspace(1, 5, steps=K).reshape(1, -1).to(device).expand(B, -1).reshape(B*K, 1)
        part_code = part_code.repeat_interleave(K, dim=0)
        gt_part_shift = gt_part_shift.repeat_interleave(K, dim=0)
        gt_part_var = gt_part_var.repeat_interleave(K, dim=0)
        noise = noise.repeat_interleave(K, dim=0)
        valid_id = valid_id.repeat_interleave(K, dim=0)
        mean, logvar = self.encoder.get_params_from_part_code(part_code, valid_id, gt_mean=gt_part_shift, gt_var=gt_part_var, noise=noise)
        mean[:, 1, [0, 2]] = (mean[:, 1, [0, 2]] * torch.sqrt(dx))
        logvar[:, 1, [0, 2]] = (logvar[:, 1, [0, 2]] + torch.log(dx))
            
        
        
        anchor_assignments = ref_seg_mask.repeat_interleave(K, dim=0).int()
        anchor_per_point, logvar_per_point, _ = self.encoder.gather_all(anchor_assignments, anchors=mean, variances=logvar + self.encoder.log_scale_var if logvar is not None else logvar)
        ctx = self.encoder.prepare_ctx(part_code, mean, logvar, anchor_assignments=anchor_assignments)
        
        N = anchor_assignments.shape[1]
        if self.npoints < N:
            anchor_per_point, logvar_per_point  = map(lambda r: r[..., :self.npoints], [anchor_per_point, logvar_per_point])
            ctx = [r[..., self.npoints] if r.shape[-1] > self.npoints else r for r in ctx]
            ref = ref[:, :self.npoints]
            anchor_assignments = anchor_assignments[:, :self.npoints]
        variance_per_point = torch.exp(logvar_per_point)
        priors = torch.randn(B*K, self.npoints, 3).to(variance_per_point.device) * variance_per_point.transpose(1,2) + anchor_per_point.transpose(1,2)
        

        pred = self.decode(anchor_per_point, ctx=ctx, device=device, variance=variance_per_point, anchor_assignments=anchor_assignments, valid_id=valid_id) 
        # ref = torch.zeros([self.npoints, 3]).cuda()
        # _ref = torch.cat([gt1[seg_mask[id1] != anchor_id, :3], gt2[seg_mask[id2] == anchor_id, :3]], dim=0)
        # ref[:_ref.shape[0]] = _ref
        # seg_mask_ref = torch.zeros_like(seg_mask[0])
        # _seg_mask_ref = torch.cat([seg_mask[id1][seg_mask[id1] != anchor_id], seg_mask[id2][seg_mask[id2] == anchor_id]], dim=0)
        # seg_mask_ref[:_seg_mask_ref.shape[0]] = _seg_mask_ref
        _pred = pred['pred'].reshape(B, K, self.npoints, 3)
        priors = priors.reshape(B, K, self.npoints, 3)
        
        for i in range(K):
            out_dict[f"interpolate sample {i}"] = _pred[:,i]
        for i in range(K):
            out_dict[f"interpolate sample prior {i}"] = priors[:, i]
            
        out_dict['pred_seg_mask'] = ref_seg_mask
        out_dict['ref_seg_mask'] = ref_seg_mask
        out_dict['seg_mask'] = ref_seg_mask
        out_dict['pred'] = out_dict["interpolate sample 0"]
        out_dict['input_ref'] = ref
        out_dict['shift'] = pcds['shift']
        out_dict['scale'] = pcds['scale']
        return out_dict

    def combine_latent_specific(self, inputs, device):
        # idss = [[61, 51, 49, 40, 38, 25, 23, 21, 13], [51, 34, 12], [62, 56, 54, 38, 29, 28, 13, 11, 7], [42, 36, 34, 32, 62]]
        # idss = [[21, 34, 13, 62], [61, 51, 62, 42], [61, 34, 56, 42], [61, 51, 56, 42], [51, 12, 38, 36], [51, 34, 29, 42], [51, 34, 11, 62]]
        assert len(inputs) == self.num_anchors
        valid_id = torch.tensor([ 1 if torch.any(inp !=0) else 0 for i, inp in enumerate(inputs) ]).cuda().unsqueeze(0)
        seg_flag = torch.cat([torch.eye(self.num_anchors).cuda()[i].reshape(1, self.num_anchors).repeat_interleave(inp.shape[0], dim=0) for i, inp in enumerate(inputs) if torch.any(inp !=0)], dim=0).unsqueeze(0)
        inputs = [inp for inp in inputs if torch.any(inp !=0)]
        input = torch.cat(inputs, dim=0).unsqueeze(0)
        print(seg_flag.shape, input.shape, valid_id)
        part_code_means, part_code_logvars = self.encoder.get_part_code(input, seg_flag)
        part_code = part_code_means.transpose(1,2)
        
        if self.cimle:
            K=self.cimle_sample_num
            noise = torch.randn(K, self.encoder.part_aligner.noise_dim).cuda()
            noise = noise.reshape(K, -1)
        else:
            K=1
            noise = None
        part_code, valid_id = map(lambda a: a.repeat_interleave(K, dim=0), [part_code, valid_id])
        out = self.encoder.get_params_from_part_code(part_code, valid_id, gt_mean=None, gt_var=None, noise=noise)
        if len(out) == 3:
            mean, logvar, _ = out
        else:
            mean, logvar = out
        anchor_assignments = torch.arange(self.num_anchors).unsqueeze(0).to(device) * valid_id 
        anchor_assignments = anchor_assignments.repeat_interleave(self.npoints // self.num_anchors, dim=1).to(torch.int32)
        mean_per_point, logvar_per_point, _ = self.encoder.gather_all(anchor_assignments, anchors=mean, variances=logvar + self.encoder.log_scale_var if logvar is not None else logvar)
        variance_per_point = torch.exp(logvar_per_point)
        ctx = self.encoder.prepare_ctx(part_code, mean, logvar, anchor_assignments=anchor_assignments)
        _pred = self.decode(mean_per_point, ctx=ctx, device=device, variance=variance_per_point, anchor_assignments=anchor_assignments, valid_id=valid_id) 
        pred = dict()
        for k, v in _pred.items():
            v = rearrange(v, "(b k) ... -> b k ...", k=K)
            for i in range(K):
                pred[f"{k}_sample_{i}"] = v[:, i]
        pred['pred'] = pred['pred_sample_0']
        pred['input'] = input
        pred['pred_seg_mask'] = anchor_assignments
        pred['seg_mask'] = torch.argmax(seg_flag, dim=2)
            
        pred['shift'] = torch.zeros(1, 1, 3).cuda()
        pred['scale'] = torch.ones(1, 1, 1).cuda()
        return pred

    def combine_latent(self, pcds, device):
        # idss = [[61, 51, 49, 40, 38, 25, 23, 21, 13], [51, 34, 12], [62, 56, 54, 38, 29, 28, 13, 11, 7], [42, 36, 34, 32, 62]]
        # idss = [[21, 34, 13, 62], [61, 51, 62, 42], [61, 34, 56, 42], [61, 51, 56, 42], [51, 12, 38, 36], [51, 34, 29, 42], [51, 34, 11, 62]]
        ref = pcds['ref'].to(device)
        refs = [ref for _ in range(self.num_anchors)]
        input = pcds['input'].to(device)
        seg_flag = pcds['attn_map'].to(device)
        refs_seg_mask = [pcds['ref_seg_mask'] for _ in range(self.num_anchors)]
        valid_id = pcds['present'].to(device)  
        gt_part_shift = pcds.get('part_shift', torch.zeros_like(ref).transpose(1,2))
        gt_part_var = pcds.get('part_scale', torch.ones_like(ref).transpose(1,2))
        B = input.shape[0]
        if not self.encoder.origin_scale:
            gt_part_var = gt_part_var ** 2
        
        
        part_code_means, part_code_logvars = self.encoder.get_part_code(input, seg_flag)
        if self.encoder.gen:
            part_code = reparameterize_gaussian(mean=part_code_means, logvar=part_code_logvars).transpose(1,2)  # (B, F)
        else:
            part_code = part_code_means.transpose(1,2)
        
        if self.cimle:
            if self.encoder.selective_noise_sampling:
                K = 100
            else:
                K=self.cimle_sample_num
            noise, _ = self.encoder.sample_noise(pcds, device, K)
            noise = noise.reshape(B*K, -1)
        else:
            K=1
            noise = None
        for i in range(self.num_anchors):
            perm = torch.randperm(B)
            refs[i] = refs[i][perm]
            refs_seg_mask[i] = pcds['ref_seg_mask'][perm]
            part_code[..., i] = part_code[perm][..., i]
            valid_id[:, i] = valid_id[perm][..., i] * valid_id[:, i]
            gt_part_shift[..., i] = gt_part_shift[perm][..., i]
            gt_part_var[..., i] = gt_part_var[perm][..., i]
        
        part_code, valid_id, gt_part_shift, gt_part_var = map(lambda a: a.repeat_interleave(K, dim=0), [part_code, valid_id, gt_part_shift, gt_part_var])
        out = self.encoder.get_params_from_part_code(part_code, valid_id, gt_mean=gt_part_shift, gt_var=gt_part_var, noise=noise)
        if len(out) == 3:
            mean, logvar, _ = out
        else:
            mean, logvar = out
        
        if self.encoder.selective_noise_sampling:
            mean, logvar = self.encoder.subsample_params(mean.reshape(B, K, 3, self.num_anchors), logvar.reshape(B, K, 3, self.num_anchors), valid_id.reshape(B, K, self.num_anchors)[:, 0], num=10)
            K = 10
            valid_id = valid_id.reshape(B, 100, self.num_anchors)[:, :K].reshape(B*K, self.num_anchors)
            part_code = part_code.reshape(B, 100, -1, self.num_anchors)[:, :K].reshape(B*K, -1, self.num_anchors)
        anchor_assignments = torch.arange(self.num_anchors).unsqueeze(0).to(device) * valid_id  + torch.argmax(valid_id, dim=1).unsqueeze(1) * (1 - valid_id)
        anchor_assignments = anchor_assignments.repeat_interleave(self.npoints // self.num_anchors, dim=1).to(torch.int32)
        mean_per_point, logvar_per_point, _ = self.encoder.gather_all(anchor_assignments, anchors=mean, variances=logvar + self.encoder.log_scale_var if logvar is not None else logvar)
        variance_per_point = torch.exp(logvar_per_point)
        ctx = self.encoder.prepare_ctx(part_code, mean, logvar, anchor_assignments=anchor_assignments)

        _pred = self.decode(mean_per_point, ctx=ctx, device=device, variance=variance_per_point, anchor_assignments=anchor_assignments, valid_id=valid_id) 
        pred = dict()
        for k, v in _pred.items():
            v = rearrange(v, "(b k) ... -> b k ...", k=K)
            for i in range(K):
                pred[f"{k}_sample_{i}"] = v[:, i]
        pred['pred'] = pred['pred_sample_0']
        pred['pred_seg_mask'] = anchor_assignments.reshape(B, K, -1)[:, 0]
        pred['input_ref'] = ref
        pred['ref_seg_mask'] = pcds['ref_seg_mask']
        for i in range(self.num_anchors):
            pred[f'input_ref{i}'] = refs[i]
            pred[f"ref_seg_mask{i}"] = refs_seg_mask[i]
            
        pred['shift'] = pcds['shift']
        pred['scale'] = pcds['scale']
        return pred

    def language_train_step(self, data, device):
        input_pcds = data['input'].to(device)
        ref_pcds = data['ref'].to(device)
        attn_maps = data['attn_map'].to(device)
        ref_attn_maps = data['ref_attn_map'].to(device)
        seg_mask = data['seg_mask'].to(device)
        ref_seg_mask = data['ref_seg_mask'].to(device)
        part_indicators = data['part_indicator'].to(device)
        texts = data['text'].to(device)
        B,_, npoint = input_pcds.shape[:-1]
        language_enc_f, language_enc_attn = self.language_encoder(texts)
    
        part_code_means, part_code_logvars = self.encoder.get_part_code(input_pcds.reshape(B*3, npoint, 3), attn_maps.reshape(B*3, npoint, self.num_anchors))
        if self.encoder.gen:
            part_code = reparameterize_gaussian(mean=part_code_means, logvar=part_code_logvars).transpose(1,2).reshape(B, 3, self.encoder.zdim, self.num_anchors)  # (B, F)
        else:
            part_code = part_code_means.transpose(1,2).reshape(B, 3, self.encoder.zdim, self.num_anchors)
            
        

        part_id = part_indicators.max(-1)[1].to(torch.int)
        referenced_part_latent = gather_operation(part_code.reshape(B*3, self.encoder.zdim, self.num_anchors).contiguous(), part_id.reshape(B*3, 1)).reshape(B, 3, self.encoder.zdim)
        tgt_latent, latent_for_edit, distractor_latent = referenced_part_latent.split(1, dim=1)
        tgt_latent, latent_for_edit, distractor_latent = tgt_latent.reshape(B, self.encoder.zdim), latent_for_edit.reshape(B, self.encoder.zdim), distractor_latent.reshape(B, self.encoder.zdim)
        modified_part_latent_res = self.update_mlp(torch.cat([language_enc_f, latent_for_edit], -1))
        modified_part_latent = latent_for_edit + modified_part_latent_res
        loss_dict = dict()
        # clip loss
        cos_sim = F.normalize(modified_part_latent, dim=1)[..., None] * F.normalize(torch.stack([tgt_latent, distractor_latent], dim=-1), dim=1)# * torch.exp(self.temp).reshape(1, 1, 1)
        cos_sim = cos_sim.sum(1)
        label = torch.tensor([1,0]).cuda().reshape(1, 2).expand(B, -1).float()
        clip_l = F.binary_cross_entropy(F.sigmoid(cos_sim), label)
        # triplet loss
        loss_dict['clip_loss'] = self.clip_weight * clip_l
        loss_dict['cos_sim_pos'] = cos_sim[:, 0]
        loss_dict['cos_sim_neg'] = cos_sim[:, 1]
        trip_l, pos_diff, neg_diff = triplet_loss(modified_part_latent, tgt_latent, distractor_latent, thresh=self.triplet_thresh) 
        loss_dict['triplet_loss'] = self.triplet_weight * trip_l
        loss_dict['pos_diff'] = pos_diff
        loss_dict['neg_diff'] = neg_diff
        return loss_dict
    
    @torch.no_grad()
    def language_edit_step(self, data, device):
        prompts = [
            "leg",
            "four leg", 
            "one leg",
            "thick leg", 
            "thin leg", 
            "wheel", 
            "long leg", 
            "short leg", 
            "seat", 
            "thin seat", 
            "thick seat",
            "back",
            "straight back", 
            "slanted back",
            "tall back", 
            "short back",
            "square back",
            "diamond back", 
            "arm rest", 
        ]
        # ids = [2]
        ids = [2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 0, 0, 0, 0, 0,0, 0, 3]
        prompts = [f"a chair with {p}" for p in prompts]
        tokenized_prompts = [tokenizing(f"a chair with {p}") for p in prompts]
        # if self.partglot_dataset:
        #     target_pcds = data['target'].to(device)
        #     target_attn_maps = data['target_attn_map'].to(device)
        #     distractor_pcds = data['distractor'].to(device)
        #     distractor_attn_maps = data['distractor_attn_map'].to(device)
        #     part_indicators = data['part_indicator'].to(device).repeat_interleave(3, dim=0)
        #     texts = data['text'].to(device)
        #     B, npoint = target_pcds.shape[:-1]
        #     pcds = torch.stack([target_pcds, distractor_pcds, distractor_pcds], dim=1).reshape(B*3, npoint, 3)
        #     attn_maps = torch.stack([target_attn_maps, distractor_attn_maps, distractor_attn_maps], dim=1).reshape(B*3, npoint, -1)
        #     input = torch.cat([pcds, attn_maps], dim=-1)
        #     language_enc_f, language_enc_attn = self.language_encoder(texts)
        
        
        #     part_latent = self.encoder(input)
        #     n_part, n_dim = part_latent.shape[1:]
            

        #     part_id = part_indicators.max(1)[1].to(torch.int)
        #     referenced_part_latent = gather_operation(part_latent.transpose(1,2).contiguous(), part_id.reshape(B*3, 1)).reshape(B, 3, 1, -1)
        #     tgt_latent, latent_for_edit, distractor_latent = referenced_part_latent.reshape(B, 3, 1, n_dim).split(1, dim=1)
        #     tgt_latent, latent_for_edit, distractor_latent = tgt_latent.reshape(B, n_dim), latent_for_edit.reshape(B, n_dim), distractor_latent.reshape(B, n_dim)
        #     modified_part_latent_res = self.update_mlp(torch.cat([language_enc_f, latent_for_edit], -1))
        #     modified_part_latent = latent_for_edit + modified_part_latent_res
        #     modified_part_latent = torch.stack([tgt_latent, modified_part_latent, distractor_latent], dim=1).reshape(B*3, 1, n_dim)
        #     total_modified_part_latent = (1 - part_indicators).reshape(B*3, n_part, 1) * part_latent + part_indicators.reshape(B*3, n_part, 1) * modified_part_latent
        #     global_feature, anchors, part_latent = self.decomposer(total_modified_part_latent)
            
            
        #     assert anchors.shape[1] == part_latent.shape[1] == self.num_anchors
        #     param_per_point = anchors.reshape(B*3, self.num_anchors, 1, 3).repeat([1, 1, self.points_per_anchor, 1]).reshape(B*3, self.npoints, 3).transpose(1,2)
        #     assert param_per_point.shape[1] == 3
        #     anchor_per_point = param_per_point

        #     latent_per_point = part_latent.reshape(B*3, self.num_anchors, 1, n_dim).repeat([1, 1, self.points_per_anchor, 1]).reshape(B*3, self.npoints, n_dim)
        #     pred = self.decode(anchor_per_point, code=global_feature, pointwise_latent=latent_per_point, device=device, variance=None) 
        #     pred_out = dict()
        #     for k, v in pred.items():
        #         target, edit, distractor = v.reshape(B, 3, self.npoints, 3).split(1, dim=1)
        #         pred_out[f"distractor_{k}"] = distractor.squeeze(1)
        #         pred_out[f"distractor_edit_{k}"] = edit.squeeze(1)
        #         pred_out[f"target_{k}"] = target.squeeze(1)
            

        #     pred_out['text'] = texts
        #     pred_out['input_distractor'] = distractor_pcds
        #     pred_out['input_target'] = target_pcds
        #     pred_out['anchors'] = anchors
        #     pred_out['seg_mask'] = distractor_attn_maps.max(-1)[1]
        #     pred_out['seg_mask_ref'] = target_attn_maps.max(-1)[1]
        #     pred_out['assigned_anchor'] = anchor_per_point.transpose(1, 2)
        #     pred_out['target_shift'] = data['target_shift']
        #     pred_out['target_scale'] = data['target_scale']
        #     pred_out['distractor_shift'] = data['distractor_shift']
        #     pred_out['distractor_scale'] = data['distractor_scale']
        #     return pred_out
        # else:
        input = data['input'].to(device)
        ref = data['ref'].to(device)
        attn_map = data['attn_map'].to(device) # ->[B, npoints, n_class]
        seg_mask = data['seg_mask'].to(device)
        valid_id = data['present'].to(device).float()
        shift = data['shift'].to(device)
        scale = data['scale'].to(device)
        out_dict = dict()
        B = input.shape[0]
        if len(input.shape) == 4:
            input = input[:, 1]
            ref = ref[:, 1]
            attn_map = attn_map[:, 1]
            valid_id=valid_id[:,1]
            seg_mask=seg_mask[:, 1]
            shift = shift[:, 1]
            scale = scale[:, 1]
        part_code_means, part_code_logvars = self.encoder.get_part_code(input, attn_map)
        if self.encoder.gen:
            part_code = reparameterize_gaussian(mean=part_code_means, logvar=part_code_logvars).transpose(1,2).reshape(B, self.encoder.zdim, self.num_anchors)  # (B, F)
        else:
            part_code = part_code_means.transpose(1,2).reshape(B, self.encoder.zdim, self.num_anchors)
        part_latent_for_edit = part_code
        for i, prompt in enumerate(tokenized_prompts):
            part_latent_edit = part_latent_for_edit.clone()
            print(prompts[i])
            prompt = prompt.to(device).unsqueeze(0)
            edit_id = ids[i]
            language_enc_f, language_enc_attn = self.language_encoder(prompt)
            language_enc_f = language_enc_f.repeat_interleave(B, dim=0)
            referenced_part_latent = part_latent_for_edit[..., edit_id]
            if self.cimle:
                K=self.cimle_sample_num
                noise = torch.randn([B*K, self.encoder.part_aligner.noise_dim]).cuda()
            else:
                noise = None
            
            modified_part_latent = self.update_mlp(torch.cat([language_enc_f, referenced_part_latent], -1)) + referenced_part_latent
            part_latent_edit[..., edit_id] = modified_part_latent
            part_latent = torch.cat([part_latent_for_edit, part_latent_edit], dim=0)
            _valid_id = valid_id[None].repeat_interleave(2, dim=0).reshape(2, B, self.num_anchors)
            noise = noise[None].repeat_interleave(2, dim=0).reshape(2*B*K, self.encoder.part_aligner.noise_dim)
            _valid_id[1, :, edit_id] = 1
            _valid_id = _valid_id.reshape(2*B, self.num_anchors).repeat_interleave(K, dim=0)
            part_latent = part_latent.repeat_interleave(K, dim=0)
            mean, logvar = self.encoder.get_params_from_part_code(part_latent, _valid_id, noise=noise)
            anchor_assignments = torch.arange(self.num_anchors).unsqueeze(0).cuda() * _valid_id + torch.argmax(_valid_id, dim=1, keepdim=True) * (1 - _valid_id)
            anchor_assignments = anchor_assignments.reshape(2*B, 1, self.num_anchors, 1).expand(-1, K, -1, self.npoints // self.num_anchors).reshape(2*B*K, self.npoints).int()
            anchor_per_point, logvar_per_point, _ = self.encoder.gather_all(anchor_assignments, anchors=mean, variances=logvar)
            ctx = self.encoder.prepare_ctx(part_latent, mean, logvar, anchor_assignments=anchor_assignments)
            variance_per_point = torch.exp(logvar_per_point)
            pred = self.decode(anchor_per_point, ctx=ctx, device=device, variance=variance_per_point, anchor_assignments=anchor_assignments, valid_id=_valid_id) 
            pred_pcd = pred['pred'].reshape(2, B, K, self.npoints, 3)
            for i in range(K):
                out_dict[f"{prompts[i]}_edit_sample_{i}"] = pred_pcd[1, :, i]
                out_dict[f"{prompts[i]}_original_sample_{i}"] = pred_pcd[0, :, i]
        out_dict['pred_seg_mask'] = anchor_assignments
        out_dict['ref_seg_mask'] = seg_mask
        out_dict['seg_mask'] = seg_mask
        out_dict['pred'] = ref
        out_dict['input_ref'] = ref
        out_dict['shift'] = shift
        out_dict['scale'] = scale
        return [(out_dict, "language")]

        # if self.sample_by_seg_mask:
        #     anchor_assignments = seg_mask.to(torch.int32)
        #     scale = int(self.npoints / anchor_assignments.shape[1])
        #     anchor_assignments = anchor_assignments.repeat_interleave(scale, dim=1)
        #     input_latent_per_point = gather_operation(input_part_latent.transpose(2,1).contiguous(), anchor_assignments).reshape(B, -1, self.npoints).transpose(1,2)
        #     input_anchor_per_point = gather_operation(input_anchor.transpose(2, 1).contiguous(), anchor_assignments).reshape(B, 3, self.npoints)
        #     assert input_anchor_per_point.shape[1] == 3
        # else:
        #     input_anchor_per_point = input_anchor.reshape(B, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B, self.npoints, -1).transpose(1,2)
        #     assert input_anchor_per_point.shape[1] == 3
        #     input_latent_per_point = input_part_latent.reshape(B, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B, self.npoints, -1)
        
        # pred = self.decode(input_anchor_per_point, code=input_global_feature, pointwise_latent=input_latent_per_point, device=device, variance=None) 
        # pred['input'] = input
        # pred['input_ref'] = ref
        # pred['anchors'] = input_anchor
        # pred['seg_mask'] = seg_mask 
        # pred['seg_mask_ref'] = seg_mask 
        # pred['assigned_anchor'] = input_anchor_per_point.transpose(1, 2)
        # pred['shift'] = data['shift']
        # pred['scale'] = data['scale']

        # if self.latent_diffusion is not None:
        #     part_latent_for_edit = part_latent_for_edit.transpose(1,2)
        #     input_anchor = input_anchor.transpose(1,2)
        # for i, prompt in enumerate(tokenized_prompts):
        #     print(prompts[i])
        #     prompt = prompt.to(device).unsqueeze(0)
        #     one_hot_id = torch.eye(self.num_anchors).to(device)[ids[i]].unsqueeze(0).repeat_interleave(B, dim=0)
        #     language_enc_f, language_enc_attn = self.language_encoder(prompt)
        #     language_enc_f = language_enc_f.repeat_interleave(B, dim=0)
        #     if self.latent_diffusion is not None:
                
        #         language_enc_f = language_enc_f.transpose(1,2)
        #         # print(part_latent_for_edit.shape, language_enc_f.shape)
        #         total_modified_part_latent = self.latent_diffusion.p_sample_loop(part_latent_for_edit.shape, part_latent_for_edit, language_enc_f, anchors=input_anchor)
        #         # print(total_modified_part_latent.shape)
        #         total_modified_part_latent = total_modified_part_latent.transpose(1, 2).unsqueeze(1)
        #         K = total_modified_part_latent.shape[1]
        #         # total_modified_part_latent = part_latent_for_edit + torch.rand_like(part_latent_for_edit)
        #     else:
        #         referenced_part_latent = part_latent_for_edit[:, ids[i]]
        #         if self.icmle:
        #             K=10
        #             conditional = torch.randn([B, K, self.conditional_dim]).cuda()
        #         else:
        #             K=1
        #             conditional = None
        #         total_modified_part_latent = self.latent_language_encoder(part_latent_for_edit, one_hot_id, language_enc_f, conditional=conditional)
        #     edited_global_feature, edited_anchors, edited_part_latent = self.decomposer(total_modified_part_latent.reshape(B*K, self.num_anchors, -1))
            
            
        #     if self.sample_by_seg_mask and not (ids[i] == 3):
        #         anchor_assignments = seg_mask.to(torch.int32)
        #         scale = int(self.npoints / anchor_assignments.shape[1])
        #         anchor_assignments = anchor_assignments.repeat_interleave(scale, dim=1).repeat_interleave(K, dim=0)
        #         edited_latent_per_point = gather_operation(edited_part_latent.reshape(B*K, self.num_anchors, -1).transpose(2,1).contiguous(), anchor_assignments).reshape(B*K, -1, self.npoints).transpose(1,2)
        #         edited_anchor_per_point = gather_operation(edited_anchors.reshape(B*K, self.num_anchors, -1).transpose(2, 1).contiguous(), anchor_assignments).reshape(B*K, 3, self.npoints)
        #         assert edited_anchor_per_point.shape[1] == 3
        #     else:
        #         edited_anchor_per_point = edited_anchors.reshape(B*K, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B*K, self.npoints, -1).transpose(1,2)
        #         assert edited_anchor_per_point.shape[1] == 3
        #         edited_latent_per_point = edited_part_latent.reshape(B*K, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B*K, self.npoints, -1)
            

        #     pred_edit = self.decode(edited_anchor_per_point, code=edited_global_feature, pointwise_latent=edited_latent_per_point, device=device, variance=None)
        #     for k in range(K):
        #         # for t in range(20, 200, 20):
        #         #     pred[f'edit_pred: "{prompts[i]}" sample {k} T={t}'] = pred_edit[t].reshape(B, K, self.npoints, 3)[:, k]
        #         pred[f'edit_pred: "{prompts[i]}" sample {k}'] = pred_edit['pred'].reshape(B, K, self.npoints, 3)[:, k]
        


    def sample(self, sample_num, fixed_id, valid_id, device, epoch, K=10):
        part_code=None
        fixed_id = torch.tensor(fixed_id).to(device)
        return self.encoder.sample_latents(sample_num, self.npoints, device, fixed_id=fixed_id, valid_id=valid_id, epoch=epoch, K=self.cimle_sample_num, part_code=part_code)
        # return self.encoder.sample_with_fixed_latents()
        # return self.encoder.sample_variation_with_fixed_latents()
        # return self.encoder.sample_with_one_fixed_latents()
        # return self.encoder.sample_with_given_latents()

    def cache_noise(self, pcds, device, eval_whole=False):
        # cache noise for CIMLE 
        from difffacto.datasets.evaluation_utils import distChamfer
        noise, id = self.encoder.sample_noise(pcds, device, self.sample_noise_num)
        if not eval_whole:
            selected_noise = gather_operation(noise.transpose(1,2).contiguous(), id.unsqueeze(1).to(torch.int32))
            return selected_noise.squeeze(-1)
        
        ref = pcds['ref'].to(device)
        seg_mask = pcds['ref_seg_mask'].to(device)
        valid_id = pcds['present'].to(device)
        B, N, C = ref.shape
        ctx, mean_per_point, logvar_per_point, flag_per_point, _ = self.encoder(pcds, device, noise=noise)
        variance_per_point = torch.exp(logvar_per_point)
        seg_mask, valid_id = map(lambda t: t.repeat_interleave(self.sample_noise_num, dim=0), [seg_mask, valid_id])
        variance_per_point = torch.exp(logvar_per_point)
        pred = self.decode(mean_per_point, ctx=ctx, device=device, variance=variance_per_point, anchor_assignments=seg_mask.to(torch.int32), valid_id=valid_id) 
        shift = pcds['shift'].to(device)
        scale = pcds['scale'].to(device)
        preds = rearrange(pred['pred'] , '(b n) ... -> b n ...', n=self.sample_noise_num) * scale.unsqueeze(1) + shift.unsqueeze(1)
        ref = ref * scale + shift
        selected_noise = []
        for i in range(ref.shape[0]):
            pred_pcd = preds[i]
            ref_pcd = ref[i].unsqueeze(0).repeat_interleave(self.sample_noise_num, dim=0)
            dl, dr = distChamfer(pred_pcd, ref_pcd)
            dist = dl.mean(1) + dr.mean(1)
            mid = dist.argmin()
            selected_noise.append(noise[i, mid])
        return torch.stack(selected_noise, dim=0)
    def cimle_forward(self, pcds,device='cuda', noise=None):
        input = pcds['input'].to(device)
        ref = pcds['ref'].to(device)
        input_seg_mask = pcds['seg_mask'].to(device)
        seg_mask = pcds['ref_seg_mask'].to(device)
        valid_id = pcds['present'].to(device)
        B, N, C = ref.shape
        if noise is None:
            noise, _ = self.encoder.sample_noise(pcds, device, 10)
            
        K = noise.shape[1]
        ctx, mean_per_point, logvar_per_point, flag_per_point, _, _ = self.encoder(pcds, device, noise=noise)
        variance_per_point = torch.exp(logvar_per_point)
        priors = torch.randn_like(variance_per_point).transpose(1,2) * torch.sqrt(variance_per_point.transpose(1,2)) + mean_per_point.transpose(1,2)
        seg_mask, valid_id = map(lambda t: t.repeat_interleave(K, dim=0), [seg_mask, valid_id])
        variance_per_point = torch.exp(logvar_per_point)
        _pred = self.decode(mean_per_point, ctx=ctx, device=device, variance=variance_per_point, anchor_assignments=seg_mask.to(torch.int32), valid_id=valid_id) 
        pred = dict()
        for i in range(K):
            for k, v in _pred.items():
                v = rearrange(v, "(b h) ... -> b h ...", h=K)
                pred[f"{k}_sample {i}"] = v[:, i]
        for i in range(K):    
            pred[f"sample prior {i}"] = priors.reshape(B, K, N, C)[:, i]
        pred['pred'] = rearrange(_pred['pred'], "(b h) ... -> b h ...", h=K)[:, 0]
        pred['input'] = input
        pred['input_ref'] = ref
        pred['seg_mask'] = input_seg_mask 
        pred['pred_seg_mask'] = seg_mask 
        pred['ref_seg_mask'] = pcds['ref_seg_mask'] 
        pred['shift'] = pcds['shift']
        pred['scale'] = pcds['scale']
        pred = {k:v.detach().cpu() for k, v in pred.items()}
        return pred
    
    def edit_latent(self, z, input, seg_flag, valid_id,  ref_means, ref_vars, fix_ids, edit_part_id, edit_part_mean, edit_part_var, fit_weight=1):
        fit_loss_dict=dict()
        part_code_means, _ = self.encoder.get_part_code(input, seg_flag)
        part_code = part_code_means.transpose(1,2)
        mean, logvar = self.encoder.get_params_from_part_code(part_code, valid_id, noise=z)
        fit_loss = F.mse_loss(torch.cat([mean, logvar], dim=1), torch.cat([ref_means, torch.log(ref_vars)], dim=1), reduction='none')
        fit_loss = fit_loss * (valid_id * fix_ids).unsqueeze(1)
        fit_loss = fit_loss.sum(dim=(-1, -2)) / (valid_id * fix_ids).sum(dim=-1)   
        if edit_part_mean is not None:
            edit_l_mean = F.mse_loss(mean[..., edit_part_id], edit_part_mean)
        else:
            edit_l_mean = torch.zeros(1).cuda()
        if edit_part_var is not None:
            edit_l_var = F.mse_loss(logvar[..., edit_part_id], torch.log(edit_part_var))
        else:
            edit_l_var = torch.zeros(1).cuda()
        edit_loss = edit_l_var + edit_l_mean
        fit_loss_dict['fit_loss'] = fit_weight * fit_loss
        fit_loss_dict['edit_loss'] = edit_loss
        if self.noise_reg_loss:
            fit_loss_dict['reg_loss'] = self.reg_loss_weight * (z ** 2).sum(1)
        return fit_loss_dict
    
    def optimize_latent(self, pcds, z, device='cuda'):
        # print(pcds.keys())
        input = pcds['input'].to(device)
        ref = pcds['ref'].to(device)
        input_seg_mask = pcds['seg_mask'].to(device)
        seg_mask = pcds['ref_seg_mask'].to(device)
        valid_id = pcds.get('present', None)
        dp_valid_id = pcds.get('dp_present', None)
        if valid_id is not None:
            valid_id = valid_id.to(device)
        if dp_valid_id is not None:
            dp_valid_id = dp_valid_id.to(device)
        loss_dict = dict()
        B, N, C = ref.shape

        ctx, mean_per_point, logvar_per_point, flag_per_point, fit_los_dict, _ = self.encoder(pcds, device, noise=z)
        if self.noise_reg_loss:
            fit_los_dict['reg_loss'] = self.reg_loss_weight * (z ** 2).sum(1)
        return fit_los_dict
    
    def pretrain(self, inputs, seg_flags, device='cuda'):
        loss_dict=dict(mse_loss=0)
        B = inputs[0].shape[0]
        if self.encoder.per_part_encoder:
            part_code_means, part_code_logvars = [], []
            for i in range(self.num_anchors):
                _m, _v = self.encoder.encoder[i](inputs[i])
                part_code_means.append(_m)
                part_code_logvars.append(_v)
            part_code_means = torch.stack(part_code_means, dim=1)
            part_code_logvars = torch.stack(part_code_logvars, dim=1)
        else:
            part_code_means, part_code_logvars = self.encoder.get_part_code(torch.cat(inputs, dim=1), torch.cat(seg_flags, dim=1))
        part_code = reparameterize_gaussian(mean=part_code_means, logvar=part_code_logvars).transpose(1,2)  # (B, F, M)
        prior_dict = self.encoder.get_prior_loss(part_code, part_code_means, part_code_logvars, torch.ones([B, self.num_anchors]).cuda())
        loss_dict.update(prior_dict)
        for i in range(self.num_anchors):
            ctx = [part_code[..., i].unsqueeze(1)]
            input = inputs[i]
            t, _ = self.sampler.sample(B, device)
            diffusion_loss = self.diffusion[i].training_losses(input.transpose(1,2), t, ctx=ctx)
            loss_dict['mse_loss'] = diffusion_loss['mse_loss'] + loss_dict['mse_loss']
        loss_dict['mse_loss'] = loss_dict['mse_loss'] / self.num_anchors
        return loss_dict
    def pretrain_part(self, input, seg_flag, part_id, device='cuda'):
        loss_dict=dict()
        B = input.shape[0]
        part_code_means, part_code_logvars = self.encoder.encoder[part_id](input)
        part_code = reparameterize_gaussian(mean=part_code_means, logvar=part_code_logvars)  # (B, F)
        prior_dict = self.encoder.get_prior_loss_part( part_code, part_code_means, part_code_logvars, part_id)
        loss_dict.update(prior_dict)
        t, _ = self.sampler.sample(B, device)
        diffusion_loss = self.diffusion[part_id].training_losses(input.transpose(1,2), t, ctx=[part_code.unsqueeze(1)])
        loss_dict.update(diffusion_loss)
        return loss_dict
    
    def pretrain_validate(self, sample_num, num_points, device='cuda'):
        part_codes = torch.randn(sample_num, self.encoder.zdim, self.num_anchors).cuda()
        finals = []
        for i in range(self.num_anchors):
            final = dict()
            ctx = [part_codes[..., i]]
            for t, sample in self.diffusion[i].p_sample_loop_progressive(
                [sample_num, 3, num_points],
                anchors=torch.zeros([sample_num, 3, num_points]).cuda(),
                variance=torch.ones([sample_num, 3, num_points]).cuda(),
                ctx=[part_codes[..., i].unsqueeze(1)],
                noise=None,
                device=device,
                progress=False,
            ):
                if t == 0:
                    finals.append(sample['sample'].transpose(2, 1))
        return finals
        
    def forward(self, pcds, device='cuda', epoch=0, **kwargs):
        """
            input: [B, npoints, 3]
            attn_map: [B, npoints, n_class]
            seg_mask: [B, npoints]
        """
        if self.save_weights:
            self.eval()
            torch.save(self.encoder.state_dict(), os.path.join(self.save_dir, "encoder_ckpt.pth"))
            torch.save(self.decomposer.state_dict(), os.path.join(self.save_dir, "decomposer_ckpt.pth"))
            exit()
        
        input = pcds['input'].to(device)
        ref = pcds['ref'].to(device)
        input_seg_mask = pcds['seg_mask'].to(device)
        seg_mask = pcds['ref_seg_mask'].to(device)
        valid_id = pcds.get('present', None)
        dp_valid_id = pcds.get('dp_present', None)
        if valid_id is not None:
            valid_id = valid_id.to(device)
        if dp_valid_id is not None:
            dp_valid_id = dp_valid_id.to(device)
        loss_dict = dict()
        B, N, C = ref.shape

        ctx, mean_per_point, logvar_per_point, flag_per_point, fit_los_dict, part_code = self.encoder(pcds, device, epoch=epoch)
        if self.zero_anchors:
            mean_per_point = torch.zeros_like(mean_per_point)
        if self.npoints < N:
            mean_per_point, logvar_per_point  = map(lambda r: r[..., :self.npoints], [mean_per_point, logvar_per_point])
            ctx = [r[..., self.npoints] if r.shape[-1] > self.npoints else r for r in ctx]
            ref = ref[:, :self.npoints]
            seg_mask = seg_mask[:, :self.npoints]
        variance_per_point = torch.exp(logvar_per_point)
        if self.training:
            
            t, _ = self.sampler.sample(B, device)
            

            loss_dict.update(fit_los_dict)

            if self.detach_anchor:
                mean_per_point = mean_per_point.detach()
            if self.detach_variance:
                logvar_per_point = logvar_per_point.detach()
            
            if dp_valid_id is not None:
                dp_flags = gather_operation(dp_valid_id.unsqueeze(1).contiguous(), seg_mask.to(torch.int32)).reshape(B, 1, self.npoints) 
            else:
                dp_flags = None
            diffusion_loss = self.diffusion.training_losses(ref.transpose(1,2) if not self.use_input else input.transpose(1,2), t, anchors=mean_per_point, variance=variance_per_point, ctx=ctx, anchor_assignment=seg_mask.to(torch.int32), valid_id=dp_valid_id, flags=dp_flags)
            diffusion_loss['mse_loss'] = self.diffusion_loss_weight * diffusion_loss['mse_loss']
            loss_dict.update(diffusion_loss)
            return loss_dict

        else:
            
            if self.interpolate and not self.training:
                return [(self.interpolate_latent(device, pcds), "interpolate")]
            if self.combine and not self.training:
                return [(self.combine_latent(pcds, device), "mixing")]
            if self.drift_anchors:
                return [(self.interpolate_params(device, pcds), "interpolate_params")] 
            
            if not self.training and self.gen:
                # fixed_ids = [[0,0,0,0], [1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 1, 1, 1]]
                fixed_ids = [0, 0, 0, 0]
                out_list = []
                if self.fix_part_ids is not None:
                    for i in self.fix_part_ids:
                        fixed_ids[i] = 1
                for fixed_id in [fixed_ids]:
                    ctx, mean_per_point, logvar_per_point, _seg_mask, _valid_id, latents = self.sample(B, fixed_id, valid_id, device, epoch, K=10)
                    fixed_codes, min_noises, min_means, min_logvars = latents
                    variance_per_point = torch.exp(logvar_per_point)
                    _pred = self.decode(mean_per_point, ctx=ctx, device=device, variance=variance_per_point, anchor_assignments=_seg_mask.to(torch.int32), valid_id=_valid_id) 
                    # pred[f"sample prior"] = priors
                    # _pred['pred_seg_mask'] = _seg_mask 
                    # _pred['means'] = min_means
                    # _pred['logvars'] = min_logvars
                    # _pred['noises'] = min_noises
                    # _pred['part_codes'] = fixed_codes
                    # return [(_pred,  "car_chosen_variation")]
                    pred = dict()
                    if self.cimle:
                        priors = torch.randn_like(variance_per_point.transpose(1,2)) * torch.sqrt(variance_per_point.transpose(1,2)) + mean_per_point.transpose(1,2)
                        if self.use_input:
                            _pred['pred'] = _pred['pred'] * torch.sqrt(variance_per_point.transpose(1,2)) + mean_per_point.transpose(1,2)
                        for i in range(self.cimle_sample_num):
                            for k, v in _pred.items():
                                v = rearrange(v, "(b h) ... -> b h ...", h=self.cimle_sample_num)
                                pred[f"{k}_sample {i}"] = v[:, i]
                        for i in range(self.cimle_sample_num):    
                            pred[f"sample prior {i}"] = priors.reshape(B, self.cimle_sample_num, self.npoints, C)[:, i]
                        pred['pred'] = rearrange(_pred['pred'], "(b h) ... -> b h ...", h=self.cimle_sample_num)[:, 0]
                        pred['pred_seg_mask'] = rearrange(_seg_mask, "(b h) ... -> b h ...", h=self.cimle_sample_num)[:, 0]
                        pred['anchors'] = rearrange(mean_per_point, "(b h) ... -> b h ...", h=self.cimle_sample_num)[:, 0].transpose(1, 2)
                    else:
                        pred = _pred
                        priors = torch.randn_like(variance_per_point.transpose(1,2)) *  torch.sqrt(variance_per_point.transpose(1,2)) + mean_per_point.transpose(1,2)
                        pred[f"sample prior"] = priors
                        pred['pred_seg_mask'] = _seg_mask 
                        pred['anchors'] = mean_per_point.transpose(1, 2)
                        
                    pred['input'] = input
                    pred['input_ref'] = ref
                    pred['ref_seg_mask'] = pcds['ref_seg_mask']
                    pred['seg_mask'] = input_seg_mask 
                    
                    pred['present'] = valid_id
                    pred['shift'] = pcds['shift']
                    pred['scale'] = pcds['scale']
                    pred = {k:v.detach().cpu() for k, v in pred.items()}
                    out_list.append((pred,  f"gen_fixed" + "".join(map(lambda i: str(i), fixed_id))))
                return out_list
            elif self.cimle:
                noise, _ = self.encoder.sample_noise(pcds, device, self.cimle_sample_num)
                ctx, mean_per_point, logvar_per_point, _, _, latents = self.encoder(pcds, device, noise=noise)
                part_code, mean, logvar, noise = latents
                seg_mask, valid_id = map(lambda t: t.repeat_interleave(self.cimle_sample_num, dim=0), [seg_mask, valid_id])
                variance_per_point = torch.exp(logvar_per_point)
            N = mean_per_point.shape[-1]
            if self.npoints > mean_per_point.shape[-1]:
                mean_per_point, variance_per_point = mean_per_point.repeat_interleave(self.npoints // N, dim=-1), variance_per_point.repeat_interleave(self.npoints // N, dim=-1)
                seg_mask = seg_mask.repeat_interleave(self.npoints // N, dim=-1)
            _pred = self.decode(mean_per_point, ctx=ctx, device=device, variance=variance_per_point, anchor_assignments=seg_mask.to(torch.int32), valid_id=valid_id) 
            pred = dict()
            if self.cimle:
                for i in range(self.cimle_sample_num):
                    for k, v in _pred.items():
                        v = rearrange(v, "(b h) ... -> b h ...", h=self.cimle_sample_num)
                        pred[f"{k}_sample {i}"] = v[:, i]
                for i in range(self.cimle_sample_num):    
                    priors = torch.randn_like(variance_per_point).transpose(1,2) *  torch.sqrt(variance_per_point.transpose(1,2)) + mean_per_point.transpose(1,2)
                    pred[f"sample prior {i}"] = priors.reshape(B, self.cimle_sample_num, self.npoints, C)[:, i]
                    pred[f"noise latent {i}"] = noise.reshape(B, self.cimle_sample_num, 32)[:, i]
                    pred[f"sample {i} mean"] = mean.reshape(B, self.cimle_sample_num, 3, self.num_anchors)[:, i]
                    pred[f"sample {i} logvar"] = logvar.reshape(B, self.cimle_sample_num, 3, self.num_anchors)[:, i]
                pred['pred'] = rearrange(_pred['pred'], "(b h) ... -> b h ...", h=self.cimle_sample_num)[:, 0]
                pred['pred_seg_mask'] = rearrange(seg_mask, "(b h) ... -> b h ...", h=self.cimle_sample_num)[:, 0]
                pred['anchors'] = rearrange(mean_per_point, "(b h) ... -> b h ...", h=self.cimle_sample_num)[:, 0].transpose(1, 2)
                pred['part_latents'] = rearrange(part_code, "(b h) ... -> b h ...", h=self.cimle_sample_num)[:, 0]
                pred['valid_id'] = rearrange(valid_id, "(b h) ... -> b h ...", h=self.cimle_sample_num)[:, 0]
            else:
                pred = _pred
                pred['pred_seg_mask'] = seg_mask 
                pred['anchors'] = mean_per_point.transpose(1, 2)
                priors = torch.randn_like(variance_per_point.transpose(1,2)) * torch.sqrt(variance_per_point.transpose(1,2)) + mean_per_point.transpose(1,2)
                pred[f"sample prior"] = priors
            # import ipdb; ipdb.set_trace()
            pred['input'] = input
            pred['input_ref'] = ref
            pred['ref_seg_mask'] = pcds['ref_seg_mask']
            pred['seg_mask'] = input_seg_mask 
            pred['token'] = pcds['token']
            pred['present'] = valid_id
            pred['shift'] = pcds['shift']
            pred['scale'] = pcds['scale']
            pred = {k:v.detach().cpu()  if isinstance(v, torch.Tensor) else v for k, v in pred.items()}
            
            out_list = [(pred, "sample")]
            # out_list.append((self.interpolate_latent(device, pcds), "interpolate"))
            # out_list.append((self.combine_latent(pcds, device), "mixing"))
            # out_list.append((self.interpolate_params(device, pcds), "interpolate_params"))
            return out_list

python/difffacto/models/networks/anchor_gen.py
diffFacto/diffFacto
from torch.nn import Module
from difffacto.utils.misc import *
from difffacto.utils.registry import build_from_cfg, MODELS, ENCODERS

@MODELS.register_module()
class ParameterTrainer(Module):

    def __init__(
        self, 
        encoder, 
        num_anchors,
        npoints=2048, 


        ):
        super().__init__()
        self.encoder = build_from_cfg(encoder, ENCODERS)
        self.num_anchors = num_anchors
        
        self.npoints = npoints
        self.points_per_anchor = self.npoints // self.num_anchors

        print(self)


    def forward(self, pcds, device='cuda', **kwargs):
        """
            input: [B, npoints, 3]
            attn_map: [B, npoints, n_class]
            seg_mask: [B, npoints]
        """


        ctx, mean_per_point, logvar_per_point, flag_per_point, fit_los_dict = self.encoder(pcds, device)
        
        return fit_los_dict

python/difffacto/models/networks/anchor_gen_parameter_trainer.py
diffFacto/diffFacto
import torch
import torch.nn as nn 
import torch.nn.functional as F
from torch.nn import Module, functional
from difffacto.utils.misc import *
from difffacto.utils.registry import build_from_cfg, MODELS, ENCODERS, DIFFUSIONS, SAMPLERS, METRICS, DECOMPOSERS, SEGMENTORS
from pointnet2_ops.pointnet2_utils import gather_operation
from .language_utils.language_util import tokenizing
VOCAB_SIZE=2787

def Normalize(in_channels):
    return nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)

def identity(*args, **kwargs):
    return nn.Identity()

def zero_module(module):
    """
    Zero out the parameters of a module and return it.
    """
    for p in module.parameters():
        p.detach().zero_()
    return module


@MODELS.register_module()
class AnchorDiffGenSuperSegments(Module):

    def __init__(
        self, 
        encoder, 
        decomposer,
        diffusion, 
        sampler, 
        num_anchors,
        num_timesteps, 
        npoints=2048, 
        anchor_loss_weight=1.,
        anchor_weight_annealing=False,
        loss=None,
        completion=False,
        include_attn_weight_in_encoder=True,
        contrastive_loss=None,
        contrastive_weight=1.,
        detach_anchor=False,
        guidance=False,
        part_latent_dropout_prob=0.2,
        global_latent_dropout_prob=0.2,
        language_edit=False,
        language_encoder=None,
        latent_language_fuser=None,
        language_encoder_ckpt=None,
        update_mlp_ckpt=None,
        latent_diffusion=None,
        ldm_ckpt=None,
        use_primary=False,
        use_zero_anchor=False,
        use_global_anchor=False,
        use_gt_anchors=False,
        project_latent=False,
        project_pe_type = 0,
        post_ff=False,
        post_norm='gn',
        post_dp=0.2,
        use_log_for_scale=True,
        share_projection=True,
        learn_var=False,
        detach_variance=True,
        part_dim=256,
        global_shift=False,
        global_scale=False,
        vertical_only=True,

        sample_by_seg_mask=False,
        ret_traj=False, 
        ret_interval=20, 
        forward_sample=False,
        interpolate=False,
        combine=False,
        drift_anchors=False,
        save_pred_xstart=False,
        partglot_dataset=False,
        save_dir=None,
        save_weights=False,
        normal_diffusion=False,
        annealing_epoch=500,
        intervaled_training=False,
        cache_interval=100,
        freeze_interval=50,
        freeze_diffusion=False,
        freeze_encoder=False
        ):
        super().__init__()
        self.encoder = build_from_cfg(encoder, ENCODERS, num_anchors=num_anchors)
        self.diffusion = build_from_cfg(diffusion, DIFFUSIONS, num_timesteps=num_timesteps)
        self.decomposer = build_from_cfg(decomposer, DECOMPOSERS, num_anchors=num_anchors, point_dim=3)
        self.sampler = build_from_cfg(sampler, SAMPLERS, num_timesteps=num_timesteps)
        self.loss_func=build_from_cfg(loss, METRICS)
        self.use_zero_anchor=use_zero_anchor
        self.use_global_anchor=use_global_anchor
        self.use_gt_anchors=use_gt_anchors
        self.intervaled_training=intervaled_training
        self.cache_interval=cache_interval
        self.freeze_interval=freeze_interval
        self.freeze_diffusion=freeze_diffusion
        self.freeze_encoder=freeze_encoder
        if cache_interval is not None and freeze_interval is not None:
            assert cache_interval > freeze_interval
        self.normal_diffusion=normal_diffusion
        self.learn_var = learn_var
        param_dim = 6 if learn_var else 3
        self.param_dim=param_dim
        if latent_diffusion is not None:
            self.latent_diffusion = build_from_cfg(latent_diffusion, DIFFUSIONS, num_part=num_anchors)
            self.latent_diffusion.load_state_dict(torch.load(ldm_ckpt))
        self.save_weights=save_weights
        self.save_dir=save_dir
        self.completion = completion
        self.global_shift = global_shift
        self.global_scale=global_scale
        self.vertical_only=vertical_only
        self.anchor_weight_annealing=anchor_weight_annealing
        if anchor_weight_annealing:
            self.start_weight=0.
            self.end_weight=1.
            self.ratio = (self.end_weight - self.start_weight) / annealing_epoch
            self.annealing_epoch=annealing_epoch
        self.drift_anchors=drift_anchors
        self.sample_by_seg_mask=sample_by_seg_mask
        self.use_log_for_scale=use_log_for_scale
        self.include_attn_weight_in_encoder=include_attn_weight_in_encoder
        self.num_timesteps = int(num_timesteps)
        self.num_anchors = num_anchors
        self.detach_anchor=detach_anchor
        self.detach_variance=detach_variance
        self.use_primary=use_primary
        self.language_edit=language_edit
        self.project_latent=project_latent 
        self.part_dim=part_dim
        if project_latent:
            self.project_pe_type = project_pe_type
            self.post_ff = post_ff 
            if post_norm == 'bn':
                post_norm = nn.BatchNorm1d
            elif post_norm == 'gn':
                post_norm = Normalize
            elif post_norm is None:
                post_norm = identity
            self.share_projection=share_projection
            if share_projection:
                if self.project_pe_type == 0:
                    self.down_proj = nn.Linear(part_dim, part_dim + param_dim)
                    self.up_proj = nn.Linear(param_dim, self.part_dim)
                elif self.project_pe_type == 1:
                    self.down_proj = nn.Sequential(
                        nn.ReLU(inplace=True),
                        nn.Linear(part_dim, 128 ),
                        nn.ReLU(inplace=True),
                        nn.Linear(128 , 64),
                        nn.ReLU(inplace=True),
                        nn.Linear(64 , param_dim),
                    )
                    self.up_proj = nn.Linear(param_dim, self.part_dim)
                elif self.project_pe_type == 2:
                    self.down_proj = nn.Sequential(
                        nn.ReLU(inplace=True),
                        nn.Linear(part_dim * self.num_anchors, 256 ),
                        nn.ReLU(inplace=True),
                        nn.Linear(256 , 64),
                        nn.ReLU(inplace=True),
                        nn.Linear(64 , param_dim * self.num_anchors),
                    )
                    self.up_proj = nn.Linear(param_dim, self.part_dim)
            else:
                self.down_proj = nn.Conv1d(part_dim * self.num_anchors, (part_dim + param_dim) * self.num_anchors, 1, groups=self.num_anchors)
                self.up_proj = nn.Conv1d(param_dim * self.num_anchors, part_dim * self.num_anchors,  1, groups=self.num_anchors)
            if self.post_ff:
                self.out_layers = nn.Sequential(
                    post_norm(part_dim),
                    nn.SiLU(),
                    nn.Dropout(p=post_dp),
                    zero_module(
                        nn.Conv1d(part_dim, part_dim, 1)
                    ),
                )
        self.partglot_dataset=partglot_dataset
        if self.language_edit:
            self.language_encoder = build_from_cfg(language_encoder, ENCODERS, vocab_size=VOCAB_SIZE)
            self.latent_language_encoder = build_from_cfg(latent_language_fuser, ENCODERS, num_part=num_anchors)
            self.language_encoder.load_state_dict(torch.load(language_encoder_ckpt))
            if self.latent_language_encoder is not None:
                self.latent_language_encoder.load_state_dict(torch.load(update_mlp_ckpt))
                self.icmle = self.latent_language_encoder.conditional
                self.conditional_dim=self.latent_language_encoder.conditional_dim
        self.npoints = npoints
        self.points_per_anchor = self.npoints // self.num_anchors
        self.forward_sample = forward_sample
        self.anchor_loss_weight=anchor_loss_weight
        self.interpolate=interpolate
        self.ret_traj=ret_traj
        self.ret_interval=ret_interval
        self.save_pred_xstart=save_pred_xstart
        self.combine=combine
        self.contrastive_loss=build_from_cfg(contrastive_loss, METRICS)
        self.contrastive_weight=contrastive_weight
        self.guidance=guidance
        self.part_latent_dropout_prob = part_latent_dropout_prob
        self.global_latent_dropout_prob=global_latent_dropout_prob
        print(self)

    
    def decode(self, anchors, code=None, pointwise_latent=None, noise=None, variance=None, anchor_assignments=None, device='cuda'):
        final = dict() 
        bs = anchors.shape[0]
        for t, sample in self.diffusion.p_sample_loop_progressive(
            [bs, 3, self.npoints],
            anchors=anchors,
            code=code,
            variance=variance,
            pointwise_latent=pointwise_latent,
            noise=noise,
            anchor_assignment=anchor_assignments,
            device=device,
            progress=False,
        ):
            if t == 0:
                final['pred'] = sample['sample'].transpose(2, 1)
                if self.save_pred_xstart:
                    final['pred_xstart'] = sample['pred_xstart'].transpose(1,2) 
            elif (self.ret_traj and t % self.ret_interval == 0):
                final[t] = sample['sample'].transpose(2, 1) 
                if self.save_pred_xstart:
                    final[f'pred_xstart_{t}'] = sample['pred_xstart'].transpose(1,2)
                
        return final

    def q_sample(self, gt, anchors, noise=None, device='cuda',variance=None):
        final = {'pred':gt.transpose(2,1)}
        
        for t, sample in self.diffusion.q_sample_loop_progressive(
            gt_pcd=gt,
            anchors=anchors,
            noise=noise,
            device=device,
            variance=variance,
            progress=False,
        ):
            if t == self.num_timesteps - 1:
                final[t] = sample.transpose(2, 1)
            elif (self.ret_traj and t % self.ret_interval == 0):
                final[t] = sample.transpose(2, 1)
                
        return final

    def interpolate_latent(self, seg_mask, device, pcds):
        # ids = [(19,20, 2), (17, 15, 0), (37, 38, 1), (122, 120, 1), (19, 18, 2), (66,65, 2), (1, 0, 2)]
        ids = [(0,1, 1), (6, 8, 2), (12, 11, 1), (11, 12, 2), (40, 38, 2), (51,50, 1), (52, 53, 2)]
        out_dict = {f"interpolate sample {i}":[] for i in range(10)}
        # out_dict['combined'] = []
        # out_dict['seg_mask_ref'] = []
        out_dict['input1_ref'] = []
        out_dict['input2_ref'] = []
        out_dict['seg_mask1_ref'] = []
        out_dict['seg_mask2_ref'] = []
        out_dict['seg_mask'] = []
        out_dict.update({f"interpolate anchor sample {i}": [] for i in range(10)})
        latents = self.get_primary_latent(pcds['input'].to(device), attn=pcds['attn_map'].to(device))
        for id1, id2, anchor_id in ids:
            latent1, latent2 = latents[id1], latents[id2]

            dx = torch.zeros([10, self.num_anchors, 1]).to(device)
            dx[:, anchor_id] += torch.linspace(0, 1, steps=10).reshape(-1, 1).to(device)
            primary_latent = latent1 + (latent2 - latent1) * dx
            B = primary_latent.shape[0]
            if self.project_latent:
                primary_latent, shift, scale, _ = self.project_primary_latent(primary_latent, interpolate=True)
            global_feat, new_anchors, secondary_latent = self.decomposer(primary_latent)
            if self.project_latent:
                new_anchors = shift 
                new_vars = scale
            else:
                new_vars = None
            if self.use_primary:
                new_latents = primary_latent 
            else:
                new_latents = secondary_latent
            if self.use_zero_anchor:
                new_anchors = torch.zeros_like(new_anchors)
            if self.use_global_anchor:
                new_anchors = pcds['global_anchor_mean'].cuda()
            if self.use_gt_anchors:
                new_anchors = pcds['part_means'].cuda()[..., 1:]
            if self.sample_by_seg_mask:
                anchor_assignments = seg_mask[id1].to(torch.int32).unsqueeze(0).repeat_interleave(B, dim=0)
                multiple = int(self.npoints / anchor_assignments.shape[1])
                anchor_assignments = anchor_assignments.repeat_interleave(multiple, dim=1)
                latent_per_point, anchor_per_point, variance_per_point = self.gather_all(anchor_assignments, new_latents, new_anchors, new_vars)
            else:
                latent_per_point = new_latents.reshape(10, self.num_anchors, 1, -1).expand(-1, -1, self.points_per_anchor, -1).reshape(10, self.npoints, -1)
                anchor_per_point = new_anchors.reshape(10, self.num_anchors, 1, 3).expand(-1, -1, self.points_per_anchor, -1).reshape(10, self.npoints, 3).transpose(1,2)
                anchor_assignments = torch.arange(self.num_anchors).to(new_anchors.device).reshape(1, -1, 1).expand(10, -1, self.points_per_anchor).reshape(10, self.npoints)
                if self.learn_var:
                    variance_per_point = new_vars.reshape(10, self.num_anchors, 1, 3).expand(-1, -1, self.points_per_anchor, -1).reshape(10, self.npoints, 3)
                else:
                    variance_per_point=None
            if self.normal_diffusion:
                anchor_per_point = anchor_per_point * 0
                variance_per_point = torch.ones_like(variance_per_point)

            pred = self.decode(anchor_per_point, code=global_feat, pointwise_latent=latent_per_point, device=device, variance=variance_per_point, anchor_assignments=anchor_assignments)
            # ref = torch.zeros([self.npoints, 3]).cuda()
            # _ref = torch.cat([gt1[seg_mask[id1] != anchor_id, :3], gt2[seg_mask[id2] == anchor_id, :3]], dim=0)
            # ref[:_ref.shape[0]] = _ref
            # seg_mask_ref = torch.zeros_like(seg_mask[0])
            # _seg_mask_ref = torch.cat([seg_mask[id1][seg_mask[id1] != anchor_id], seg_mask[id2][seg_mask[id2] == anchor_id]], dim=0)
            # seg_mask_ref[:_seg_mask_ref.shape[0]] = _seg_mask_ref
            for i in range(pred['pred'].shape[0]):
                out_dict[f"interpolate sample {i}"].append(pred['pred'][i])
                if pred.get(self.num_timesteps, None) is not None:
                    if f"interpolate sample {i} priors" not in out_dict.keys():
                        out_dict[f"interpolate sample {i} priors"] = []
                    out_dict[f"interpolate sample {i} priors"].append(pred[self.num_timesteps][i])
                out_dict[f"interpolate anchor sample {i}"].append(new_anchors[i])
            # out_dict['combined'].append(ref)
            # out_dict['seg_mask_ref'].append(seg_mask_ref)
            out_dict['input1_ref'].append(pcds['input'][id1, :, 0:3])
            out_dict['input2_ref'].append(pcds['input'][id2, :, 0:3])
            out_dict['seg_mask1_ref'].append(seg_mask[id1])
            out_dict['seg_mask'].append(seg_mask[id1])
            out_dict['seg_mask2_ref'].append(seg_mask[id2])
        out_dict = {k:torch.stack(v, dim=0) for k, v in out_dict.items()}
        out_dict['pred'] = pcds['input'][0].unsqueeze(0).repeat_interleave(len(ids), dim=0)
        out_dict['input_ref'] = pcds['input'][0].unsqueeze(0).repeat_interleave(len(ids), dim=0)
        out_dict['shift'] = pcds['shift'][0].unsqueeze(0)
        out_dict['scale'] = pcds['scale'][0].unsqueeze(0)
        return out_dict

    def part_completion(self, x, seg_mask, id1, anchor_id, device):
        gt = x[id1]
        idx = gt[..., 3 + anchor_id] == 0.
        gt_part = gt[idx]
        latent = self.encoder(xyz=gt_part.unsqueeze(0))
        new_params, new_latents = self.decomposer(latent)
        latent_per_point = new_latents.reshape(1, self.num_anchors, 1, -1).expand(-1, -1, self.points_per_anchor, -1).reshape(1, self.npoints, -1)
        param_per_point = new_params.reshape(1, self.num_anchors, 1, 3).expand(-1, -1, self.points_per_anchor, -1).reshape(1, self.npoints, 3).transpose(1,2)
        assert param_per_point.shape[1] == 3
        anchor_per_point = param_per_point
        variance_per_point = None

        pred = self.decode(anchor_per_point, code=None, pointwise_latent=latent_per_point, device=device, variance=variance_per_point) 
        pred['input'] = gt_part[:, 0:3].unsqueeze(0)
        pred['ref'] = gt[:, 0:3].unsqueeze(0)
        pred['anchors'] = new_params[..., :3]
        
        return pred

    def combine_latent(self, x, pcds, seg_mask, ids, device):
        idss = [(18, 21, 12, 19)]
        primary_latents = self.get_primary_latent(pcds['input'].to(device), attn=pcds['attn_map'].to(device))
        for ids in idss:
            primary_latents = torch.stack([primary_latents[i] for i in ids], dim=0)
            primary_latents = torch.stack([primary_latents[i, i] for i in range(self.num_anchors)], dim=0).reshape(1, self.num_anchors, -1)
            if self.project_latent:
                primary_latents, shift, scale, _ = self.project_primary_latent(primary_latents, mixing=True)
            global_feature, new_anchors, secondary_latent = self.decomposer(primary_latents)
            if self.project_latent:
                new_anchors = shift 
                new_vars = scale
            else:
                new_vars = None
            if self.use_primary:
                new_latents = primary_latents
            else:
                new_latents = secondary_latent
            if self.use_zero_anchor:
                new_anchors = torch.zeros_like(new_anchors)
            if self.use_global_anchor:
                new_anchors = pcds['global_anchor_mean'].cuda()
            if self.use_gt_anchors:
                new_anchors = torch.stack([pcds['part_means'].cuda()[ids[i], i + 1] for i in range(self.num_anchors)], dim=0)
            latent_per_point = new_latents.reshape(1, self.num_anchors, 1, -1).expand(-1, -1, self.points_per_anchor, -1).reshape(1, self.npoints, -1)
            param_per_point = new_params.reshape(1, self.num_anchors, 1, 3).expand(-1, -1, self.points_per_anchor, -1).reshape(1, self.npoints, 3).transpose(1,2)
            assert param_per_point.shape[1] == 3
            anchor_per_point = param_per_point
            variance_per_point = None
        
        pred = self.decode(anchor_per_point, code=global_feature, pointwise_latent=latent_per_point, device=device, variance=variance_per_point) 
        ref = torch.cat([gts[i, seg_mask[ids[i]] == i, :3] for i in range(self.num_anchors)], dim=0)[:self.npoints]
        seg_mask_ref = torch.cat([seg_mask[ids[i]][seg_mask[ids[i]] == i] for i in range(self.num_anchors)], dim=0)
        print(ref.shape)
        pred['input_ref'] = ref.unsqueeze(0)
        pred['seg_mask_ref'] = seg_mask_ref.unsqueeze(0)
        for i in range(self.num_anchors):
            pred[f'input_{i}'] = gts[i, :, 0:3].unsqueeze(0)
            pred[f'seg_mask_{i}'] = seg_mask[ids[i]].unsqueeze(0)
        pred['anchors'] = new_params[..., :3]
        pred['shift'] = pcds['shift']
        pred['scale'] = pcds['scale']
        return pred

    def language_edit_step(self, data, device):
        prompts = [
            "leg",
            "four leg", 
            "one leg",
            "thick leg", 
            "thin leg", 
            "wheel", 
            "long leg", 
            "short leg", 
            "seat", 
            "thin seat", 
            "thick seat",
            "back",
            "straight back", 
            "slanted back",
            "tall back", 
            "short back",
            "square back",
            "diamond back", 
            "arm rest", 
        ]
        # ids = [2]
        ids = [2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 0, 0, 0, 0, 0,0, 0, 3]
        prompts = [f"a chair with {p}" for p in prompts]
        tokenized_prompts = [tokenizing(f"a chair with {p}") for p in prompts]
        if self.partglot_dataset:
            target_pcds = data['target'].to(device)
            target_attn_maps = data['target_attn_map'].to(device)
            distractor_pcds = data['distractor'].to(device)
            distractor_attn_maps = data['distractor_attn_map'].to(device)
            part_indicators = data['part_indicator'].to(device).repeat_interleave(3, dim=0)
            texts = data['text'].to(device)
            B, npoint = target_pcds.shape[:-1]
            pcds = torch.stack([target_pcds, distractor_pcds, distractor_pcds], dim=1).reshape(B*3, npoint, 3)
            attn_maps = torch.stack([target_attn_maps, distractor_attn_maps, distractor_attn_maps], dim=1).reshape(B*3, npoint, -1)
            input = torch.cat([pcds, attn_maps], dim=-1)
            language_enc_f, language_enc_attn = self.language_encoder(texts)
        
        
            part_latent = self.encoder(input)
            n_part, n_dim = part_latent.shape[1:]
            

            part_id = part_indicators.max(1)[1].to(torch.int)
            referenced_part_latent = gather_operation(part_latent.transpose(1,2).contiguous(), part_id.reshape(B*3, 1)).reshape(B, 3, 1, -1)
            tgt_latent, latent_for_edit, distractor_latent = referenced_part_latent.reshape(B, 3, 1, n_dim).split(1, dim=1)
            tgt_latent, latent_for_edit, distractor_latent = tgt_latent.reshape(B, n_dim), latent_for_edit.reshape(B, n_dim), distractor_latent.reshape(B, n_dim)
            modified_part_latent_res = self.update_mlp(torch.cat([language_enc_f, latent_for_edit], -1))
            modified_part_latent = latent_for_edit + modified_part_latent_res
            modified_part_latent = torch.stack([tgt_latent, modified_part_latent, distractor_latent], dim=1).reshape(B*3, 1, n_dim)
            total_modified_part_latent = (1 - part_indicators).reshape(B*3, n_part, 1) * part_latent + part_indicators.reshape(B*3, n_part, 1) * modified_part_latent
            global_feature, anchors, part_latent = self.decomposer(total_modified_part_latent)
            
            
            assert anchors.shape[1] == part_latent.shape[1] == self.num_anchors
            param_per_point = anchors.reshape(B*3, self.num_anchors, 1, 3).repeat([1, 1, self.points_per_anchor, 1]).reshape(B*3, self.npoints, 3).transpose(1,2)
            assert param_per_point.shape[1] == 3
            anchor_per_point = param_per_point

            latent_per_point = part_latent.reshape(B*3, self.num_anchors, 1, n_dim).repeat([1, 1, self.points_per_anchor, 1]).reshape(B*3, self.npoints, n_dim)
            pred = self.decode(anchor_per_point, code=global_feature, pointwise_latent=latent_per_point, device=device, variance=None) 
            pred_out = dict()
            for k, v in pred.items():
                target, edit, distractor = v.reshape(B, 3, self.npoints, 3).split(1, dim=1)
                pred_out[f"distractor_{k}"] = distractor.squeeze(1)
                pred_out[f"distractor_edit_{k}"] = edit.squeeze(1)
                pred_out[f"target_{k}"] = target.squeeze(1)
            

            pred_out['text'] = texts
            pred_out['input_distractor'] = distractor_pcds
            pred_out['input_target'] = target_pcds
            pred_out['anchors'] = anchors
            pred_out['seg_mask'] = distractor_attn_maps.max(-1)[1]
            pred_out['seg_mask_ref'] = target_attn_maps.max(-1)[1]
            pred_out['assigned_anchor'] = anchor_per_point.transpose(1, 2)
            pred_out['target_shift'] = data['target_shift']
            pred_out['target_scale'] = data['target_scale']
            pred_out['distractor_shift'] = data['distractor_shift']
            pred_out['distractor_scale'] = data['distractor_scale']
            return pred_out
        else:
            input = data['input'].to(device)
            ref = data['ref'].to(device)
            attn_map = data['attn_map'].to(device) # ->[B, npoints, n_class]
            seg_mask = data['seg_mask'].to(device)
            B = input.shape[0]
            x = torch.cat([input, attn_map], dim=-1).repeat_interleave(2, dim=0) # -> [B, npoints, n_class + 3]
            part_latent = self.encoder(x)
            input_part_latent, part_latent_for_edit = part_latent.reshape(B, 2, self.num_anchors, -1).split(1, dim=1)
            input_part_latent, part_latent_for_edit = input_part_latent.reshape(B, self.num_anchors, -1), part_latent_for_edit.reshape(B, self.num_anchors, -1)
            input_global_feature, input_anchor, input_part_latent = self.decomposer(input_part_latent)

            if self.sample_by_seg_mask:
                anchor_assignments = seg_mask.to(torch.int32)
                scale = int(self.npoints / anchor_assignments.shape[1])
                anchor_assignments = anchor_assignments.repeat_interleave(scale, dim=1)
                input_latent_per_point = gather_operation(input_part_latent.transpose(2,1).contiguous(), anchor_assignments).reshape(B, -1, self.npoints).transpose(1,2)
                input_anchor_per_point = gather_operation(input_anchor.transpose(2, 1).contiguous(), anchor_assignments).reshape(B, 3, self.npoints)
                assert input_anchor_per_point.shape[1] == 3
            else:
                input_anchor_per_point = input_anchor.reshape(B, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B, self.npoints, -1).transpose(1,2)
                assert input_anchor_per_point.shape[1] == 3
                input_latent_per_point = input_part_latent.reshape(B, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B, self.npoints, -1)
            
            pred = self.decode(input_anchor_per_point, code=input_global_feature, pointwise_latent=input_latent_per_point, device=device, variance=None) 
            pred['input'] = input
            pred['input_ref'] = ref
            pred['anchors'] = input_anchor
            pred['seg_mask'] = seg_mask 
            pred['seg_mask_ref'] = seg_mask 
            pred['assigned_anchor'] = input_anchor_per_point.transpose(1, 2)
            pred['shift'] = data['shift']
            pred['scale'] = data['scale']

            if self.latent_diffusion is not None:
                part_latent_for_edit = part_latent_for_edit.transpose(1,2)
                input_anchor = input_anchor.transpose(1,2)
            for i, prompt in enumerate(tokenized_prompts):
                print(prompts[i])
                prompt = prompt.to(device).unsqueeze(0)
                one_hot_id = torch.eye(self.num_anchors).to(device)[ids[i]].unsqueeze(0).repeat_interleave(B, dim=0)
                language_enc_f, language_enc_attn = self.language_encoder(prompt)
                language_enc_f = language_enc_f.repeat_interleave(B, dim=0)
                if self.latent_diffusion is not None:
                    
                    language_enc_f = language_enc_f.transpose(1,2)
                    # print(part_latent_for_edit.shape, language_enc_f.shape)
                    total_modified_part_latent = self.latent_diffusion.p_sample_loop(part_latent_for_edit.shape, part_latent_for_edit, language_enc_f, anchors=input_anchor)
                    # print(total_modified_part_latent.shape)
                    total_modified_part_latent = total_modified_part_latent.transpose(1, 2).unsqueeze(1)
                    K = total_modified_part_latent.shape[1]
                    # total_modified_part_latent = part_latent_for_edit + torch.rand_like(part_latent_for_edit)
                else:
                    referenced_part_latent = part_latent_for_edit[:, ids[i]]
                    if self.icmle:
                        K=10
                        conditional = torch.randn([B, K, self.conditional_dim]).cuda()
                    else:
                        K=1
                        conditional = None
                    total_modified_part_latent = self.latent_language_encoder(part_latent_for_edit, one_hot_id, language_enc_f, conditional=conditional)
                edited_global_feature, edited_anchors, edited_part_latent = self.decomposer(total_modified_part_latent.reshape(B*K, self.num_anchors, -1))
                
                
                if self.sample_by_seg_mask and not (ids[i] == 3):
                    anchor_assignments = seg_mask.to(torch.int32)
                    scale = int(self.npoints / anchor_assignments.shape[1])
                    anchor_assignments = anchor_assignments.repeat_interleave(scale, dim=1).repeat_interleave(K, dim=0)
                    edited_latent_per_point = gather_operation(edited_part_latent.reshape(B*K, self.num_anchors, -1).transpose(2,1).contiguous(), anchor_assignments).reshape(B*K, -1, self.npoints).transpose(1,2)
                    edited_anchor_per_point = gather_operation(edited_anchors.reshape(B*K, self.num_anchors, -1).transpose(2, 1).contiguous(), anchor_assignments).reshape(B*K, 3, self.npoints)
                    assert edited_anchor_per_point.shape[1] == 3
                else:
                    edited_anchor_per_point = edited_anchors.reshape(B*K, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B*K, self.npoints, -1).transpose(1,2)
                    assert edited_anchor_per_point.shape[1] == 3
                    edited_latent_per_point = edited_part_latent.reshape(B*K, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B*K, self.npoints, -1)
                

                pred_edit = self.decode(edited_anchor_per_point, code=edited_global_feature, pointwise_latent=edited_latent_per_point, device=device, variance=None)
                for k in range(K):
                    # for t in range(20, 200, 20):
                    #     pred[f'edit_pred: "{prompts[i]}" sample {k} T={t}'] = pred_edit[t].reshape(B, K, self.npoints, 3)[:, k]
                    pred[f'edit_pred: "{prompts[i]}" sample {k}'] = pred_edit['pred'].reshape(B, K, self.npoints, 3)[:, k]
        return pred
        
    def anchor_drift(self, pcds, x, global_feature, anchors, part_latent, seg_mask, device):
        out_dict = dict()
        B = global_feature.shape[0]
        direction = F.normalize(torch.randn(B, 10, self.num_anchors, 1).cuda())
        part_id = torch.tensor([1.,1., 1., 1.]).to(float).cuda().reshape(1, 1, -1, 1)
        scales = torch.linspace(0., 2, 10).to(float).cuda().reshape(1, -1, 1, 1)
        part_latent = part_latent.repeat_interleave(10, dim=0)
        global_feature = global_feature.repeat_interleave(10, dim=0)
        if self.project_latent:
            primary_part_latent = self.encoder(x)
            intrinsic, anchors = torch.split(self.down_proj(primary_part_latent), [self.part_dim, 3], dim=-1)
            intrinsic = intrinsic.repeat_interleave(10, dim=0)
        shifted_anchors = direction * scales  * part_id + anchors.unsqueeze(1) * part_id #+ anchors.unsqueeze(1)
        shifted_anchors = shifted_anchors.reshape(B*10, self.num_anchors, 3).to(torch.float32)
        if self.project_latent:
            primary_part_latent = intrinsic + self.up_proj(shifted_anchors)
            global_feature, _, part_latent = self.decomposer(primary_part_latent)
        if self.use_primary:
            part_latent = primary_part_latent 
        
        anchor_assignments = seg_mask.to(torch.int32).repeat_interleave(10, dim=0)
        scale = int(self.npoints / anchor_assignments.shape[1])
        anchor_assignments = anchor_assignments.repeat_interleave(scale, dim=1)
        latent_per_point = gather_operation(part_latent.transpose(2,1).contiguous(), anchor_assignments).reshape(B*10, -1, self.npoints).transpose(1,2)
        anchor_per_point = gather_operation(shifted_anchors.transpose(2, 1).contiguous(), anchor_assignments).reshape(B*10, 3, self.npoints)
        print(latent_per_point.shape, anchor_per_point.shape)
        assert anchor_per_point.shape[1] == 3
        pred = self.decode(anchor_per_point, code=global_feature, pointwise_latent=latent_per_point, device=device, variance=None) 
        for k in range(10):
            # for t in range(20, 200, 20):
            #     pred[f'edit_pred: "{prompts[i]}" sample {k} T={t}'] = pred_edit[t].reshape(B, K, self.npoints, 3)[:, k]
            out_dict[f'anchor sample {k}'] = shifted_anchors.reshape(B, 10, self.num_anchors, 3)[:, k]
            out_dict[f'pred anchor drift sample {k}'] = pred['pred'].reshape(B, 10, self.npoints, 3)[:, k]
        out_dict['input'] = pcds['input']
        out_dict['pred'] = pcds['input']
        out_dict['input_ref'] = pcds['input']
        
        out_dict['seg_mask'] = seg_mask 
        out_dict['seg_mask_ref'] = seg_mask 
        # out_dict['assigned_anchor'] = anchor_per_point.transpose(1, 2)
        out_dict['shift'] = pcds['shift']
        out_dict['scale'] = pcds['scale']
        return out_dict

    def get_params(self, pcds, device='cuda'):
        assert self.project_latent
        input = pcds['input'].to(device)
        attn_map = pcds['attn_map'].to(device) # ->[B, npoints, n_class]
        primary_part_latent = self.get_primary_latent(input, attn_map)
        _, shift, scale, _ = self.project_primary_latent(primary_part_latent)
        return shift, torch.log(scale)
    
    def get_primary_latent(self, input, attn=None):
        if self.include_attn_weight_in_encoder:
            assert attn is not None
            x = torch.cat([input, attn], dim=-1) # -> [B, npoints, n_class + 3]
        else:
            x = input


        primary_part_latent = self.encoder(x)
        return primary_part_latent

    def project_primary_latent(self, latent, interpolate=False, fix_shift=None, fix_scale=None, ref=None, mixing=False):
        B = latent.shape[0]
        if self.share_projection:
            if self.project_pe_type == 0:
                intrinsic, extrinsic = torch.split(self.down_proj(latent), [self.part_dim, self.param_dim], dim=-1)
            elif self.project_pe_type == 1:
                extrinsic = self.down_proj(latent)
                intrinsic = latent
            elif self.project_pe_type == 2:
                extrinsic = self.down_proj(latent.reshape(B, -1)).reshape(B, self.num_anchors, self.param_dim)
                intrinsic = latent
        else:
            intrinsic, extrinsic = torch.split(self.down_proj(latent.reshape(B, -1, 1)).reshape(B, self.num_anchors, self.part_dim + self.param_dim), [self.part_dim, self.param_dim], dim=-1)

        if self.learn_var:
            shift, scale = extrinsic.split(3, dim=-1)
        else:
            shift = extrinsic
            scale = torch.ones_like(shift)
        part_shifts = torch.zeros_like(shift)
        part_scales = torch.ones_like(scale)
        if self.training and self.global_shift:
            assert ref is not None
            rand_shift = torch.rand(B, 1, 3).cuda() - 0.5
            if self.vertical_only:
                rand_shift[..., [0,2]] = 0.
            part_shifts = rand_shift
            ref = ref + rand_shift
        if self.training and self.global_scale:
            assert ref is not None
            part_scales = torch.rand(B,1,3).cuda() / 2 + 0.7
            ref = ref * part_scales


        corrected_anchor = (shift + part_shifts) * part_scales 
        corrected_scale = scale + torch.log(part_scales)
        if interpolate:
            fix_shift = shift[0]
            fix_scale = scale[0]
        if mixing:
            fix_shift = torch.stack([shift[i, i] for i in range(self.num_anchors)], dim=0)
            fix_scale = torch.stack([scale[i, i] for i in range(self.num_anchors)], dim=0)
        if fix_shift is not None:
            while len(fix_shift.shape) < len(shift.shape):
                fix_shift = fix_shift.unsqueeze(0)
            fix_shift = fix_shift.expand([i if fix_shift.shape[id] == 1 else -1 for id, i in enumerate(shift.shape)])
            for i, j in zip(fix_shift.shape, shift.shape):
                assert i == j
            corrected_anchor = fix_shift
        if fix_scale is not None:
            while len(fix_scale.shape) < len(scale.shape):
                fix_scale = fix_scale.unsqueeze(0)
            fix_scale = fix_scale.expand([i if fix_scale.shape[id] == 1 else -1 for id, i in enumerate(scale.shape)])
            for i, j in zip(fix_scale.shape, scale.shape):
                assert i == j
            corrected_scale = fix_scale
        params = torch.cat([corrected_anchor, corrected_scale], dim=-1) if self.learn_var else corrected_anchor

        if self.share_projection:
            _primary_part_latent = intrinsic + self.up_proj(params)
        else:
            _primary_part_latent = intrinsic + self.up_proj(params.reshape(B, -1, 1)).reshape(B, self.num_anchors, self.part_dim)

        if self.post_ff:
            latent = latent + self.out_layers(_primary_part_latent.transpose(1,2)).transpose(1,2)
        else:
            latent = _primary_part_latent
        if self.use_log_for_scale:
            corrected_scale = torch.exp(corrected_scale)
        else:
            corrected_scale = F.relu(corrected_scale)
        return latent, corrected_anchor, torch.exp(corrected_scale), ref
    
    def gather_all(self, anchor_assignments, part_latent, anchors=None, variances=None):
        B = anchor_assignments.shape[0]
        latent_per_point = gather_operation(part_latent.transpose(2,1).contiguous(), anchor_assignments).reshape(B, -1, self.npoints).transpose(1,2) 
        if anchors is not None:
            anchor_per_point = gather_operation(anchors.transpose(2, 1).contiguous(), anchor_assignments).reshape(B, 3, self.npoints)  
        if variances is not None and  not isinstance(variances, float):
            variance_per_point = gather_operation(variances.transpose(2, 1).contiguous(), anchor_assignments).reshape(B, 3, self.npoints).transpose(1,2) 
        else:
            variance_per_point = 1.
        assert anchor_per_point.shape[1] == 3
        
        return latent_per_point, anchor_per_point, variance_per_point
    def modify_grad(self, epoch):
        if self.freeze_diffusion and (epoch % self.cache_interval >= self.freeze_interval):
            for param in self.diffusion.parameters():
                param.grad = None
            for param in self.decomposer.parameters():
                param.grad = None
        if self.freeze_encoder and (epoch % self.cache_interval < self.freeze_interval):
            for param in self.encoder.parameters():
                param.grad = None
            for param in self.down_proj.parameters():
                param.grad = None
            for param in self.up_proj.parameters():
                param.grad = None
    def forward(self, pcds, device='cuda', iter=-1, epoch=-1, fixing_shift=None, fixing_scale=None):
        """
            input: [B, npoints, 3]
            attn_map: [B, npoints, n_class]
            seg_mask: [B, npoints]
        """
        if self.save_weights:
            self.eval()
            print(1)
            torch.save(self.encoder.state_dict(), os.path.join(self.save_dir, "encoder_ckpt.pth"))
            torch.save(self.decomposer.state_dict(), os.path.join(self.save_dir, "decomposer_ckpt.pth"))
            exit()
        if self.language_edit:
            return self.language_edit_step(pcds, device)
        input = pcds['input'].to(device)
        ref = pcds['ref'].to(device)
        attn_map = pcds['attn_map'].to(device) # ->[B, npoints, n_class]
        seg_mask = pcds['seg_mask'].to(device)
        loss_dict = dict()


        primary_part_latent = self.get_primary_latent(input, attn=attn_map)
        if self.project_latent:
            if self.intervaled_training and (epoch % self.cache_interval < self.freeze_interval):
                primary_part_latent, corrected_anchor, corrected_scale, _ = self.project_primary_latent(primary_part_latent, fix_shift=fixing_shift, fix_scale=fixing_scale)
            else:
                primary_part_latent, corrected_anchor, corrected_scale, ref = self.project_primary_latent(primary_part_latent, ref=ref)
                
        global_feature, anchors, secondary_part_latent = self.decomposer(primary_part_latent)
        assert anchors.shape[1] == secondary_part_latent.shape[1] == self.num_anchors
        variances = 1.
        if self.use_primary:
            part_latent = primary_part_latent 
        else:
            part_latent = secondary_part_latent
        if self.project_latent:
            anchors = corrected_anchor
            if self.learn_var:
                variances = corrected_scale
        if self.use_zero_anchor:
            anchors = torch.zeros_like(anchors)
        if self.use_global_anchor:
            anchors = pcds['global_anchor_mean'].cuda()
        if self.use_gt_anchors:
            anchors = pcds['part_means'].cuda()[..., 1:]
        
        if self.training:
            if self.guidance:
                idxs = torch.rand(part_latent.shape[:-1]).to(part_latent.device) >= self.part_latent_dropout_prob
                part_latent = part_latent * idxs.unsqueeze(-1)

                global_idxs = torch.rand(global_feature.shape[0]).to(part_latent.device) >= self.global_latent_dropout_prob
                global_feature = global_feature * global_idxs.unsqueeze(-1)

            t, _ = self.sampler.sample(input.shape[0], device)
            latent_per_point, anchor_per_point, variance_per_point = self.gather_all(seg_mask.to(torch.int32), part_latent, anchors, variances)
            
            if self.anchor_weight_annealing:
                anchor_loss_weight = self.ratio * min(epoch, self.annealing_epoch) + self.start_weight
            elif self.use_zero_anchor or self.use_global_anchor or self.use_gt_anchors:
                anchor_loss_weight = 0
            else:
                anchor_loss_weight = self.anchor_loss_weight
            anchor_loss = self.loss_func(ref, anchor_per_point.transpose(1,2), var = variance_per_point)
            loss_dict['anchor_loss'] = anchor_loss_weight * anchor_loss
            loss_dict['anchor_weight'] = torch.ones(1).cuda() * anchor_loss_weight
            
            if self.normal_diffusion:
                anchor_per_point = anchor_per_point * 0
                variance_per_point = torch.ones_like(variance_per_point)

            if self.detach_anchor:
                anchor_per_point = anchor_per_point.detach()
            if self.detach_variance and not isinstance(variance_per_point, float):
                variance_per_point = variance_per_point.detach()
            if self.freeze_diffusion and (epoch % self.cache_interval > self.freeze_interval):
                with torch.no_grad():
                    diffusion_loss = self.diffusion.training_losses(ref.transpose(1,2), t, anchors=anchor_per_point, code=global_feature, pointwise_latent=latent_per_point, variance=variance_per_point, reduce=True, anchor_assignment=seg_mask.to(torch.int32))
            else:
                diffusion_loss = self.diffusion.training_losses(ref.transpose(1,2), t, anchors=anchor_per_point, code=global_feature, pointwise_latent=latent_per_point, variance=variance_per_point, reduce=True, anchor_assignment=seg_mask.to(torch.int32))
                
            loss_dict.update(diffusion_loss)
            return loss_dict

        else:
            
            if self.interpolate and not self.training:
                return self.interpolate_latent(seg_mask, device, pcds)
            if self.combine and not self.training:
                ids = [15, 50, 14,19]
                return self.combine_latent(x, pcds, seg_mask, ids, device)
            if self.drift_anchors:
                return self.anchor_drift(pcds, x, global_feature, anchors, part_latent, seg_mask, device)

            if self.sample_by_seg_mask:
                anchor_assignments = seg_mask.to(torch.int32)
                multiple = int(self.npoints / anchor_assignments.shape[1])
                anchor_assignments = anchor_assignments.repeat_interleave(multiple, dim=1)
                latent_per_point, anchor_per_point, variance_per_point = self.gather_all(anchor_assignments, part_latent, anchors, variances)

            else:
                B = input.shape[0]
                anchor_per_point = anchors.reshape(B, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B, self.npoints, -1).transpose(1,2)
                assert anchor_per_point.shape[1] == 3
                latent_per_point = part_latent.reshape(B, self.num_anchors, 1, 3).repeat([1, 1, self.points_per_anchor, 1]).reshape(B, self.npoints, 3)
                if not isinstance(variances, float):
                    variance_per_point = variances.reshape(B, self.num_anchors, 1, 3).repeat([1, 1, self.points_per_anchor, 1]).reshape(B, self.npoints, 3).transpose(1,2)
                else:
                    variance_per_point = 1.
                anchor_assignments=torch.arange(self.num_anchors).reshape(1, -1, 1).expand(B, -1, self.points_per_anchor).reshape(B, self.npoints)
            if self.normal_diffusion:
                anchor_per_point = anchor_per_point * 0
                variance_per_point = torch.ones_like(variance_per_point)

           
            pred = self.q_sample(ref.transpose(2,1), anchors=anchor_per_point, device=device, variance=None) if self.forward_sample else self.decode(anchor_per_point, code=global_feature, pointwise_latent=latent_per_point, device=device, variance=variance_per_point, anchor_assignments=anchor_assignments) 
            pred['input'] = input
            pred['input_ref'] = ref
            pred['anchors'] = anchors
            pred['seg_mask'] = seg_mask 
            pred['seg_mask_ref'] = seg_mask 
            pred['assigned_anchor'] = anchor_per_point.transpose(1, 2)
            pred['shift'] = pcds['shift']
            pred['scale'] = pcds['scale']
            pred = {k:v.detach().cpu() for k, v in pred.items()}
            return pred

python/difffacto/models/networks/anchor_gen_supseg.py
diffFacto/diffFacto
import torch
from torch.nn import Module, functional
from difffacto.utils.misc import *
from difffacto.utils.registry import build_from_cfg, MODELS, ENCODERS, DIFFUSIONS, SAMPLERS, METRICS, DECOMPOSERS, SEGMENTORS
from pointnet2_ops.pointnet2_utils import gather_operation

@MODELS.register_module()
class AnchorDiffGenPartglot(Module):

    def __init__(
        self, 
        encoder, 
        decomposer,
        diffusion, 
        sampler, 
        num_anchors,
        num_timesteps, 
        npoints=2048, 
        anchor_loss_weight=1.,
        loss=None,


        ret_traj=False, 
        ret_interval=20, 
        forward_sample=False,
        interpolate=False,
        combine=False,
        save_pred_xstart=False
        ):
        super().__init__()
        self.encoder = build_from_cfg(encoder, ENCODERS, num_anchors=num_anchors)
        self.diffusion = build_from_cfg(diffusion, DIFFUSIONS, num_timesteps=num_timesteps, learn_variance=False)
        self.decomposer = build_from_cfg(decomposer, DECOMPOSERS, num_anchors=num_anchors, point_dim=3)
        self.sampler = build_from_cfg(sampler, SAMPLERS, num_timesteps=num_timesteps)
        self.loss_func=build_from_cfg(loss, METRICS, reduction=None)
        self.num_timesteps = int(num_timesteps)
        self.num_anchors = num_anchors
        self.npoints = npoints
        self.points_per_anchor = self.npoints // self.num_anchors
        self.forward_sample = forward_sample
        self.anchor_loss_weight=anchor_loss_weight
        self.interpolate=interpolate
        self.ret_traj=ret_traj
        self.ret_interval=ret_interval
        self.save_pred_xstart=save_pred_xstart
        self.combine=combine
    
    def decode(self, anchors, code=None, pointwise_latent=None, noise=None, variance=None, device='cuda'):
        final = dict() 
        bs = anchors.shape[0]
        for t, sample in self.diffusion.p_sample_loop_progressive(
            [bs, 3, self.npoints],
            anchors=anchors,
            code=code,
            variance=variance,
            pointwise_latent=pointwise_latent,
            noise=noise,
            device=device,
            progress=False,
        ):
            if t == 0:
                final['pred'] = sample['sample'].transpose(2, 1)
                final['pred_xstart'] = sample['pred_xstart'].transpose(1,2) 
            elif (self.ret_traj and t % self.ret_interval == 0):
                final[t] = sample['sample'].transpose(2, 1) 
                final[f'pred_xstart_{t}'] = sample['pred_xstart'].transpose(1,2) 
                
        return final

    def q_sample(self, gt, anchors, noise=None, device='cuda',variance=None):
        final = {'pred':gt.transpose(2,1)}
        
        for t, sample in self.diffusion.q_sample_loop_progressive(
            gt_pcd=gt,
            anchors=anchors,
            noise=noise,
            device=device,
            variance=variance,
            progress=False,
        ):
            if t == self.num_timesteps - 1:
                final[t] = sample.transpose(2, 1)
            elif (self.ret_traj and t % self.ret_interval == 0):
                final[t] = sample.transpose(2, 1)
                
        return final

    def interpolate_latent(self, x, seg_mask, id1, id2, anchor_id, device):
        gt1, gt2 = x[id1], x[id2]
        latent1 = self.encoder(xyz=gt1.unsqueeze(0))
        latent2 = self.encoder(xyz=gt2.unsqueeze(0))
        latent1, latent2 = latent1.squeeze(0), latent2.squeeze(0) # [1, n_class, C]

        dx = torch.zeros([10, self.num_anchors, 1]).to(device)
        dx[:, anchor_id] += torch.linspace(0, 1, steps=10).reshape(-1, 1).to(device)
        latents = latent1.unsqueeze(0) + (latent2 - latent1).unsqueeze(0) * dx
        global_feat, new_params, new_latents = self.decomposer(latents)
        latent_per_point = new_latents.reshape(10, self.num_anchors, 1, -1).expand(-1, -1, self.points_per_anchor, -1).reshape(10, self.npoints, -1)
        param_per_point = new_params.reshape(10, self.num_anchors, 1, 3).expand(-1, -1, self.points_per_anchor, -1).reshape(10, self.npoints, 3).transpose(1,2)
        assert param_per_point.shape[1] == 3
        anchor_per_point = param_per_point
        variance_per_point = None


        pred = self.decode(anchor_per_point, code=global_feat, pointwise_latent=latent_per_point, device=device, variance=variance_per_point) 
        ref = torch.cat([gt1[seg_mask[id1] != anchor_id, :3], gt2[seg_mask[id2] == anchor_id, :3]], dim=0)
        seg_mask_ref = torch.cat([seg_mask[id1][seg_mask[id1] != anchor_id], seg_mask[id2][seg_mask[id2] == anchor_id]], dim=0)
        pred['input_ref'] = ref.unsqueeze(0).expand(10, -1,-1)
        pred['seg_mask_ref'] = seg_mask_ref.unsqueeze(0).expand(10,-1)
        pred['input1'] = gt1[:, 0:3].unsqueeze(0).expand(10, -1, -1)
        pred['input2'] = gt2[:, 0:3].unsqueeze(0).expand(10, -1, -1)
        pred['seg_mask1'] = seg_mask[id1].unsqueeze(0).expand(10, -1)
        pred['seg_mask2'] = seg_mask[id2].unsqueeze(0).expand(10, -1)
        pred['anchors'] = new_params[..., :3]
        return pred


    def combine_latent(self, x, seg_mask, ids, device):
        gts = []
        for i in ids:
            gts.append(x[i])
        gts = torch.stack(gts, dim=0)
        latents = self.encoder(xyz=gts)
        latents = torch.stack([latents[i, i] for i in range(self.num_anchors)], dim=0).reshape(1, self.num_anchors, -1)
        global_feature, new_params, new_latents = self.decomposer(latents)
        latent_per_point = new_latents.reshape(1, self.num_anchors, 1, -1).expand(-1, -1, self.points_per_anchor, -1).reshape(1, self.npoints, -1)
        param_per_point = new_params.reshape(1, self.num_anchors, 1, 3).expand(-1, -1, self.points_per_anchor, -1).reshape(1, self.npoints, 3).transpose(1,2)
        assert param_per_point.shape[1] == 3
        anchor_per_point = param_per_point
        variance_per_point = None
        
        pred = self.decode(anchor_per_point, code=global_feature, pointwise_latent=latent_per_point, device=device, variance=variance_per_point) 
        ref = torch.cat([gts[i, seg_mask[ids[i]] == i, :3] for i in range(self.num_anchors)], dim=0)
        seg_mask_ref = torch.cat([seg_mask[ids[i]][seg_mask[ids[i]] == i] for i in range(self.num_anchors)], dim=0)
        pred['input_ref'] = ref.unsqueeze(0)
        pred['seg_mask_ref'] = seg_mask_ref.unsqueeze(0)
        for i in range(self.num_anchors):
            pred[f'input_{i}'] = gts[i, :, 0:3].unsqueeze(0)
            pred[f'seg_mask_{i}'] = seg_mask[ids[i]].unsqueeze(0)
        pred['anchors'] = new_params[..., :3]
        return pred

    def forward(self, pcds, device='cuda'):
        """
            input: [B, n seg, npoints, 3]
            attn_map: [B, n_class, n_seg]
            geo_mask: [B, n_seg]
        """

        input = pcds['input'].to(device)
        ref = pcds['ref'].to(device)
        attn_map = pcds['attn_map'].to(device) # ->[B, npoints, n_class]
        seg_mask = pcds['seg_mask'].to(device)
        geo_mask = pcds['geo_mask'].to(device)
        loss_dict = dict()

        B, n_seg, npoint = input.shape[:-1]
        num_active = geo_mask.sum(1)
        

        part_latent = self.encoder(input, geo_mask)

        global_feature, anchors, part_latent = self.decomposer(part_latent)
        
        assert anchors.shape[1] == part_latent.shape[1] == self.num_anchors
        

        if self.training:

            t, _ = self.sampler.sample(B*n_seg, device)
            anchor_assignments = seg_mask.to(torch.int32)

            global_feature = global_feature.unsqueeze(1).expand(-1, n_seg, -1).reshape(B*n_seg, -1)
            latent_per_point = gather_operation(part_latent.transpose(2,1).contiguous(), anchor_assignments).reshape(B, -1, n_seg)
            anchor_per_point = gather_operation(anchors.transpose(2, 1).contiguous(), anchor_assignments).reshape(B, 3, n_seg)  
            assert anchor_per_point.shape[1] == 3
            
            anchor_per_point = anchor_per_point.transpose(1,2).unsqueeze(2).expand(-1, -1, npoint, -1)
            latent_per_point = latent_per_point.transpose(1,2).unsqueeze(2).expand(-1, -1, npoint, -1)
            anchor_loss = self.loss_func(anchor_per_point, ref)
            loss_dict['anchor_loss'] = self.anchor_loss_weight * ((anchor_loss.mean(-1) * geo_mask).sum(-1) / num_active).mean()

            diffusion_loss = self.diffusion.training_losses(ref.reshape(B*n_seg, npoint, 3).transpose(1,2), t, anchors=anchor_per_point.reshape(B*n_seg, npoint, 3).transpose(1,2), code=global_feature, pointwise_latent=latent_per_point.reshape(B*n_seg, npoint, -1), variance=None, reduce=False)
            diffusion_loss = (diffusion_loss.reshape(B, n_seg) * geo_mask).sum(1) / num_active 
            loss_dict['diffusion_loss'] = diffusion_loss.mean()
            return loss_dict

        else:
            
            if self.interpolate and not self.training:
                id1, id2 = 0,1
                anchor_id= 1
                return self.interpolate_latent(x, seg_mask, id1, id2, anchor_id, device)
            if self.combine and not self.training:
                ids = [86, 78, 100,98]
                return self.combine_latent(x, seg_mask, ids, device)

            param_per_point = anchors.reshape(B, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B, self.npoints, -1).transpose(1,2)
            assert param_per_point.shape[1] == 3
            anchor_per_point = param_per_point

            latent_per_point = part_latent.reshape(B, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B, self.npoints, -1)
            
            pred = self.q_sample(ref.transpose(2,1), anchors=anchor_per_point, device=device, variance=None) if self.forward_sample else self.decode(anchor_per_point, code=global_feature, pointwise_latent=latent_per_point, device=device, variance=None) 
            pred['input'] = input
            pred['input_ref'] = ref
            pred['anchors'] = anchors
            pred['seg_mask'] = seg_mask 
            pred['seg_mask_ref'] = seg_mask 
            pred['assigned_anchor'] = anchor_per_point.transpose(1, 2)
            return pred

python/difffacto/models/networks/anchor_gen_supseg_partglot.py
diffFacto/diffFacto
from distutils.command.build import build
import torch
import torch.nn as nn 
from torch.nn import Module, functional
from difffacto.utils.misc import *
from difffacto.utils.registry import build_from_cfg, MODELS, ENCODERS, DIFFUSIONS, SAMPLERS, METRICS, DECOMPOSERS, SEGMENTORS
from pointnet2_ops.pointnet2_utils import gather_operation
from .language_utils.language_util import tokenizing
VOCAB_SIZE=2787

@MODELS.register_module()
class AnchoredGeneration(Module):

    def __init__(
        self, 
        encoder, 
        decomposer,
        parameter_regressor,
        diffusion, 
        sampler, 
        num_anchors,
        num_timesteps, 
        npoints=2048, 
        anchor_loss_weight=1.,
        loss=None,
        completion=False,
        include_attn_weight_in_encoder=True,
        guidance=False,
        part_latent_dropout_prob=0.2,
        global_latent_dropout_prob=0.2,
        kl_weight=1e-3,
        ret_traj=False, 
        ret_interval=20, 
        forward_sample=False,
        interpolate=False,
        combine=False,
        save_pred_xstart=False,
        fix_part=False
        ):
        super().__init__()
        self.encoder = build_from_cfg(encoder, ENCODERS, num_anchors=num_anchors)
        self.diffusion = build_from_cfg(diffusion, DIFFUSIONS, num_timesteps=num_timesteps, learn_variance=False)
        self.decomposer = build_from_cfg(decomposer, DECOMPOSERS, num_anchors=num_anchors, point_dim=3)
        self.parameter_regressor = build_from_cfg(parameter_regressor, ENCODERS, num_class=num_anchors)
        self.sampler = build_from_cfg(sampler, SAMPLERS, num_timesteps=num_timesteps)
        self.loss_func=build_from_cfg(loss, METRICS)
        self.latent_dim = self.encoder.part_latent_dim
        self.fix_part=fix_part
        self.kl_weight=kl_weight
        self.completion = completion
        self.include_attn_weight_in_encoder=include_attn_weight_in_encoder
        self.num_timesteps = int(num_timesteps)
        self.num_anchors = num_anchors
        self.npoints = npoints
        self.points_per_anchor = self.npoints // self.num_anchors
        self.forward_sample = forward_sample
        self.anchor_loss_weight=anchor_loss_weight
        self.interpolate=interpolate
        self.ret_traj=ret_traj
        self.ret_interval=ret_interval
        self.save_pred_xstart=save_pred_xstart
        self.combine=combine
        self.guidance=guidance
        self.part_latent_dropout_prob = part_latent_dropout_prob
        self.global_latent_dropout_prob=global_latent_dropout_prob
    
    def decode(self, anchors, code=None, pointwise_latent=None, noise=None, variance=None, device='cuda'):
        final = dict() 
        bs = anchors.shape[0]
        for t, sample in self.diffusion.p_sample_loop_progressive(
            [bs, 3, self.npoints],
            anchors=anchors,
            code=code,
            variance=variance,
            pointwise_latent=pointwise_latent,
            noise=noise,
            device=device,
            progress=False,
        ):
            if t == 0:
                final['pred'] = sample['sample'].transpose(2, 1)
                final['pred_xstart'] = sample['pred_xstart'].transpose(1,2) 
            elif (self.ret_traj and t % self.ret_interval == 0):
                final[t] = sample['sample'].transpose(2, 1) 
                final[f'pred_xstart_{t}'] = sample['pred_xstart'].transpose(1,2) 
                
        return final

    def q_sample(self, gt, anchors, noise=None, device='cuda',variance=None):
        final = {'pred':gt.transpose(2,1)}
        
        for t, sample in self.diffusion.q_sample_loop_progressive(
            gt_pcd=gt,
            anchors=anchors,
            noise=noise,
            device=device,
            variance=variance,
            progress=False,
        ):
            if t == self.num_timesteps - 1:
                final[t] = sample.transpose(2, 1)
            elif (self.ret_traj and t % self.ret_interval == 0):
                final[t] = sample.transpose(2, 1)
                
        return final

    def interpolate_latent(self, x, seg_mask, id1, id2, anchor_id, device, pcds):
        gt1, gt2 = x[id1], x[id2]
        latent1 = self.encoder(xyz=gt1.unsqueeze(0))
        latent2 = self.encoder(xyz=gt2.unsqueeze(0))
        latent1, latent2 = latent1.squeeze(0), latent2.squeeze(0) # [1, n_class, C]

        dx = torch.zeros([10, self.num_anchors, 1]).to(device)
        dx[:, anchor_id] += torch.linspace(0, 1, steps=10).reshape(-1, 1).to(device)
        latents = latent1.unsqueeze(0) + (latent2 - latent1).unsqueeze(0) * dx
        global_feat, new_params, new_latents = self.decomposer(latents)
        latent_per_point = new_latents.reshape(10, self.num_anchors, 1, -1).expand(-1, -1, self.points_per_anchor, -1).reshape(10, self.npoints, -1)
        param_per_point = new_params.reshape(10, self.num_anchors, 1, 3).expand(-1, -1, self.points_per_anchor, -1).reshape(10, self.npoints, 3).transpose(1,2)
        assert param_per_point.shape[1] == 3
        anchor_per_point = param_per_point
        variance_per_point = None


        pred = self.decode(anchor_per_point, code=global_feat, pointwise_latent=latent_per_point, device=device, variance=variance_per_point) 
        ref = torch.cat([gt1[seg_mask[id1] != anchor_id, :3], gt2[seg_mask[id2] == anchor_id, :3]], dim=0)
        seg_mask_ref = torch.cat([seg_mask[id1][seg_mask[id1] != anchor_id], seg_mask[id2][seg_mask[id2] == anchor_id]], dim=0)
        pred['input_ref'] = ref.unsqueeze(0).expand(10, -1,-1)
        pred['seg_mask_ref'] = seg_mask_ref.unsqueeze(0).expand(10,-1)
        pred['input1'] = gt1[:, 0:3].unsqueeze(0).expand(10, -1, -1)
        pred['input2'] = gt2[:, 0:3].unsqueeze(0).expand(10, -1, -1)
        pred['seg_mask1'] = seg_mask[id1].unsqueeze(0).expand(10, -1)
        pred['seg_mask2'] = seg_mask[id2].unsqueeze(0).expand(10, -1)
        pred['anchors'] = new_params[..., :3]
        pred['shift'] = pcds['shift']
        pred['scale'] = pcds['scale']
        return pred

    def part_completion(self, x, seg_mask, id1, anchor_id, device):
        gt = x[id1]
        idx = gt[..., 3 + anchor_id] == 0.
        gt_part = gt[idx]
        latent = self.encoder(xyz=gt_part.unsqueeze(0))
        new_params, new_latents = self.decomposer(latent)
        latent_per_point = new_latents.reshape(1, self.num_anchors, 1, -1).expand(-1, -1, self.points_per_anchor, -1).reshape(1, self.npoints, -1)
        param_per_point = new_params.reshape(1, self.num_anchors, 1, 3).expand(-1, -1, self.points_per_anchor, -1).reshape(1, self.npoints, 3).transpose(1,2)
        assert param_per_point.shape[1] == 3
        anchor_per_point = param_per_point
        variance_per_point = None

        pred = self.decode(anchor_per_point, code=None, pointwise_latent=latent_per_point, device=device, variance=variance_per_point) 
        pred['input'] = gt_part[:, 0:3].unsqueeze(0)
        pred['ref'] = gt[:, 0:3].unsqueeze(0)
        pred['anchors'] = new_params[..., :3]
        
        return pred

    def combine_latent(self, x, seg_mask, ids, device):
        gts = []
        for i in ids:
            gts.append(x[i])
        gts = torch.stack(gts, dim=0)
        latents = self.encoder(xyz=gts)
        latents = torch.stack([latents[i, i] for i in range(self.num_anchors)], dim=0).reshape(1, self.num_anchors, -1)
        global_feature, new_params, new_latents = self.decomposer(latents)
        latent_per_point = new_latents.reshape(1, self.num_anchors, 1, -1).expand(-1, -1, self.points_per_anchor, -1).reshape(1, self.npoints, -1)
        param_per_point = new_params.reshape(1, self.num_anchors, 1, 3).expand(-1, -1, self.points_per_anchor, -1).reshape(1, self.npoints, 3).transpose(1,2)
        assert param_per_point.shape[1] == 3
        anchor_per_point = param_per_point
        variance_per_point = None
        
        pred = self.decode(anchor_per_point, code=global_feature, pointwise_latent=latent_per_point, device=device, variance=variance_per_point) 
        ref = torch.cat([gts[i, seg_mask[ids[i]] == i, :3] for i in range(self.num_anchors)], dim=0)
        seg_mask_ref = torch.cat([seg_mask[i][seg_mask[ids[i]] == i] for i in range(self.num_anchors)], dim=0)
        pred['input_ref'] = ref.unsqueeze(0)
        pred['seg_mask_ref'] = seg_mask_ref.unsqueeze(0)
        for i in range(self.num_anchors):
            pred[f'ref_{i}'] = gts[i, :, 0:3].unsqueeze(0)
            pred[f'seg_mask_{i}'] = seg_mask[i].unsqueeze(0)
        pred['anchors'] = new_params[..., :3]
        return pred

    def forward(self, pcds, device='cuda'):
        """
            input: [B, npoints, 3]
            attn_map: [B, npoints, n_class]
            seg_mask: [B, npoints]
        """
        input = pcds['input'].to(device)
        ref = pcds['ref'].to(device)
        attn_map = pcds['attn_map'].to(device) # ->[B, npoints, n_class]
        seg_mask = pcds['seg_mask'].to(device)
        loss_dict = dict()

        npoint = input.shape[-2]
       

        if self.include_attn_weight_in_encoder:
            x = torch.cat([input, attn_map], dim=-1) # -> [B, npoints, n_class + 3]
        else:
            x = input

        B = input.shape[0]

        
        


        if self.training:
            part_latent = self.encoder(x)
            means, variances = self.parameter_regressor(part_latent)
            assert means.shape[1] == variances.shape[1] == self.num_anchors
            z = reparameterize_gaussian(mean=means, logvar=variances)  # (B, F)
            log_pz = standard_normal_logprob(z).sum(dim=2)  # (B, ), Independence assumption
            entropy = gaussian_entropy(logvar=variances)      # (B, )
            loss_prior = (- log_pz - entropy).mean()
            global_feature, anchors, part_latent = self.decomposer(z)
            loss_dict['prior_loss'] = self.kl_weight * loss_prior
            if self.guidance:
                idxs = torch.rand(part_latent.shape[:-1]).to(part_latent.device) >= self.part_latent_dropout_prob
                part_latent = part_latent * idxs.unsqueeze(-1)

                global_idxs = torch.rand(global_feature.shape[0]).to(part_latent.device) >= self.global_latent_dropout_prob
                global_feature = global_feature * global_idxs.unsqueeze(-1)

            t, _ = self.sampler.sample(B, device)
            anchor_assignments = seg_mask.to(torch.int32)

            latent_per_point = gather_operation(part_latent.transpose(2,1).contiguous(), anchor_assignments).reshape(B, -1, npoint)
            anchor_per_point = gather_operation(anchors.transpose(2, 1).contiguous(), anchor_assignments).reshape(B, 3, npoint)  
            assert anchor_per_point.shape[1] == 3

            anchor_loss = self.loss_func(anchor_per_point.transpose(1,2), ref)
            loss_dict['anchor_loss'] = self.anchor_loss_weight * anchor_loss

            diffusion_loss = self.diffusion.training_losses(ref.transpose(1,2), t, anchors=anchor_per_point, code=global_feature, pointwise_latent=latent_per_point.transpose(1,2), variance=None, reduce=True)
            loss_dict['diffusion_loss'] = diffusion_loss
            return loss_dict

        else:
            z = torch.randn([B, self.num_anchors, self.latent_dim]).cuda()
            
            
            if self.fix_part:
                anchor_id = 0
                z = torch.randn([1, self.num_anchors, self.latent_dim]).cuda().repeat_interleave(B, dim=0)
                z[:, anchor_id] = torch.randn([B, self.latent_dim]).cuda()
            if self.interpolate and not self.training:
                id1, id2 = 22,107
                anchor_id= 2
                return self.interpolate_latent(x, seg_mask, id1, id2, anchor_id, device, pcds)
            if self.combine and not self.training:
                ids = [23, 114, 19,89]
                return self.combine_latent(x, seg_mask, ids, device)
            global_feature, anchors, part_latent = self.decomposer(z)
            anchor_per_point = anchors.reshape(B, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B, self.npoints, -1).transpose(1,2)
            assert anchor_per_point.shape[1] == 3
            latent_per_point = part_latent.reshape(B, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B, self.npoints, -1)
            
            pred = self.q_sample(ref.transpose(2,1), anchors=anchor_per_point, device=device, variance=None) if self.forward_sample else self.decode(anchor_per_point, code=global_feature, pointwise_latent=latent_per_point, device=device, variance=None) 
            pred['input'] = input
            pred['input_ref'] = ref
            pred['anchors'] = anchors
            pred['seg_mask'] = seg_mask 
            pred['seg_mask_ref'] = seg_mask 
            pred['assigned_anchor'] = anchor_per_point.transpose(1, 2)
            pred['shift'] = pcds['shift']
            pred['scale'] = pcds['scale']
            pred = {k:v.detach().cpu() for k, v in pred.items()}
            return pred

python/difffacto/models/networks/anchor_gen_uncond.py
diffFacto/diffFacto
import torch
import torch.nn as nn
from difffacto.utils.registry import MODELS, METRICS, build_from_cfg

@MODELS.register_module()
class CoarsePCN(nn.Module):
    def __init__(self, num_anchors, encoder_channel, loss):
        super().__init__()
        self.encoder_channel = encoder_channel
        self.number_coarse = num_anchors
        self.first_conv = nn.Sequential(
            nn.Conv1d(3,128,1),
            nn.BatchNorm1d(128),
            nn.ReLU(inplace=True),
            nn.Conv1d(128,256,1)
        )
        self.second_conv = nn.Sequential(
            nn.Conv1d(512,512,1),
            nn.BatchNorm1d(512),
            nn.ReLU(inplace=True),
            nn.Conv1d(512,self.encoder_channel,1)
        )
        self.mlp = nn.Sequential(
            nn.Linear(self.encoder_channel,1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024,1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024,3*self.number_coarse)
        )
        self.loss_func = build_from_cfg(loss, METRICS)

    def forward(self, pcds, device):
        gt = pcds['pointcloud'].to(device)
        xyz = pcds['pointcloud'].to(device)
        bs , n , _ = xyz.shape
        # encoder
        feature = self.first_conv(xyz.transpose(2,1))  # B 256 n
        feature_global = torch.max(feature,dim=2,keepdim=True)[0]  # B 256 1
        feature = torch.cat([feature_global.expand(-1,-1,n), feature], dim=1)# B 512 n
        feature = self.second_conv(feature) # B 1024 n
        feature_global = torch.max(feature,dim=2,keepdim=False)[0] # B 1024
        # decoder
        coarse = self.mlp(feature_global).reshape(-1,self.number_coarse,3) # B M 3
        
        
        if self.training:
            loss_coarse = self.loss_func(coarse, gt)
            return dict(loss_coarse=loss_coarse)
        return dict(pred=coarse.contiguous(), ref=gt, input=xyz)

python/difffacto/models/networks/coarsePCN.py
diffFacto/diffFacto
import torch
from torch.nn import Module
from difffacto.utils.misc import *
from difffacto.utils.registry import build_from_cfg, MODELS, ENCODERS, DIFFUSIONS, SAMPLERS

@MODELS.register_module()
class DiffuCompletion(Module):

    def __init__(self, encoder, diffusion, sampler, npoints=2048, ret_traj=False, ret_interval=20):
        super().__init__()
        self.encoder = build_from_cfg(encoder, ENCODERS)
        self.npoints = npoints
        self.diffusion = build_from_cfg(diffusion, DIFFUSIONS)
        self.num_timesteps = self.diffusion.num_timesteps
        self.sampler = build_from_cfg(sampler, SAMPLERS, num_timesteps=self.num_timesteps)
        self.ret_traj=ret_traj
        self.ret_interval=ret_interval

    
    def decode(self, code, noise=None, device='cuda'):
        final = dict() 
        bs = code.shape[0]
        for t, sample in self.diffusion.p_sample_loop_progressive(
            [bs, 3, self.npoints],
            code=code,
            noise=noise,
            device=device,
            progress=False,
        ):
            if t == 0:
                final['pred'] = sample['sample'].transpose(2, 1)
            elif (self.ret_traj and t % self.ret_interval == 0):
                final[t] = sample['sample'].transpose(2, 1)
            
                
        return final
                


    def q_sample(self, input, anchors=None, anchor_assignments=None, ret_traj=False, ret_interval=1, device='cuda'):
        if anchors is None:
            anchors = fps(input, self.num_anchors)
            anchor_assignments = assign_anchor(input, anchors).to(device)
        return self.diffusion.forward_sample(input, anchors, anchor_assignments, ret_traj=ret_traj, ret_interval=ret_interval, device=device)
    
    def forward(self, pcds, device='cuda'):
        fine_gt = pcds['pointcloud'].to(device)
        input = pcds['partial'].to(device)
        code = self.encoder(input)
        if self.training:
            t, _ = self.sampler.sample(input.shape[0], device)
            recons_loss = self.diffusion.training_losses(fine_gt.transpose(2, 1), t, code)
            return recons_loss
        else:
            pred = self.decode(code, device=device)
            pred['ref'] = fine_gt 
            pred['input'] = input
            return pred
python/difffacto/models/networks/completion.py
diffFacto/diffFacto
import torch
import numpy as np
from torch import nn
from difffacto.utils.registry import build_from_cfg, MODELS, ENCODERS
from difffacto.utils.misc import standard_normal_logprob, truncated_normal, gaussian_log_likelihood
from pointnet2_ops.pointnet2_utils import gather_operation
from .flow_utils.flow import get_point_cnf, get_latent_cnf




@MODELS.register_module()
class PointFlow(nn.Module):
    def __init__(self, 
                 encoder, 
                 input_dim, 
                 zdim, 
                 use_latent_flow, 
                 use_deterministic_encoder, 
                 point_cnf,
                 latent_cnf,
                 part_aligner,
                 prior_weight=1., 
                 recon_weight=1.,
                 entropy_weight=1.,

                 n_class=4
                 ):
        super(PointFlow, self).__init__()
        self.input_dim = input_dim
        self.zdim = zdim
        self.use_latent_flow = use_latent_flow
        self.use_deterministic_encoder = use_deterministic_encoder
        self.prior_weight = prior_weight
        self.recon_weight = recon_weight
        self.entropy_weight = entropy_weight
        self.truncate_std = None
        self.encoder = build_from_cfg(encoder, ENCODERS)
        self.part_aligner=build_from_cfg(part_aligner, ENCODERS)
        self.point_cnf = get_point_cnf(point_cnf)
        self.latent_cnf = [get_latent_cnf(latent_cnf) for _ in range(n_class)] if use_latent_flow else nn.Sequential()
        self.n_class=n_class
        
    @staticmethod
    def sample_gaussian(size, truncate_std=None, gpu=None):
        y = torch.randn(*size).float().cuda()
        if truncate_std is not None:
            truncated_normal(y, mean=0, std=1, trunc_std=truncate_std)
        return y

    @staticmethod
    def reparameterize_gaussian(mean, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn(std.size()).to(mean)
        return mean + std * eps

    def encode(self, x):
        z_mu, z_sigma = self.encoder(x)
        if self.use_deterministic_encoder:
            return z_mu
        else:
            return self.reparameterize_gaussian(z_mu, z_sigma)
    
    @staticmethod
    def gaussian_entropy(logvar):
        const = 0.5 * float(logvar.size(1)) * (1. + np.log(np.pi * 2))
        ent = 0.5 * logvar.sum(dim=1, keepdim=False) + const
        return ent

    def gather_all(self, anchor_assignments, anchors=None, variances=None):
        B, N = anchor_assignments.shape
        anchor_per_point = torch.zeros(B, 3, N).to(anchor_assignments.device) 
        variance_per_point = torch.zeros(B, 3, N).to(anchor_assignments.device) 
        if anchors is not None:
            anchor_per_point = gather_operation(anchors.contiguous(), anchor_assignments).reshape(B, 3, N)  
        if variances is not None:
            variance_per_point = gather_operation(variances.contiguous(), anchor_assignments).reshape(B, 3, N) 
        
        return anchor_per_point, variance_per_point
    
    def forward(self, pcds, device='cuda', **kwargs):
        input = pcds['input'].to(device)
        ref = pcds['ref'].to(device)
        seg_mask = pcds['ref_seg_mask'].to(device).to(torch.int32)
        seg_flag = pcds['attn_map'].to(device)
        valid_id = pcds.get('present', None)
        if valid_id is not None:
            valid_id = valid_id.to(device)
        B, N, C = ref.shape
        if not self.training:
            _, pred = self.reconstruct(torch.cat([input, seg_flag], dim=-1), N, None, valid_id) if self.use_deterministic_encoder else self.sample(B, N, None, None, device, valid_id)
            return [(dict(pred=pred, input_ref=ref, shift=pcds['shift'], scale=pcds['scale']), "gen" if not self.use_deterministic_encoder else "sample")]
        z_mu, z_sigma = self.encoder(torch.cat([input, seg_flag], dim=-1))
        assert z_mu.shape[1] == z_sigma.shape[1] == self.n_class
        if self.use_deterministic_encoder:
            z = z_mu + 0 * z_sigma
        else:
            z = self.reparameterize_gaussian(z_mu, z_sigma)
        mean, logvar = self.part_aligner(z.transpose(1,2), mask=valid_id)
        mean_per_point, logvar_per_point = self.gather_all(seg_mask, anchors=mean, variances=logvar)
        assert mean_per_point.shape == logvar_per_point.shape == ref.transpose(1,2).shape
        if self.training and self.part_aligner is not None:
            fit_loss = -gaussian_log_likelihood(ref.transpose(1,2), mean_per_point, logvar_per_point).mean()
        # Compute H[Q(z|X)]
        if self.use_deterministic_encoder:
            entropy = torch.zeros(B).to(z)
        else:
            entropy = self.gaussian_entropy(z_sigma.reshape(B* self.n_class, -1))

        # Compute the prior probability P(z)
        if self.use_latent_flow:
            log_pz = torch.zeros(B, self.n_class, 1).to(input)
            for i in range(self.n_class):
                w, delta_log_pw = self.latent_cnf[i](z[:, i], None, torch.zeros(B, 1).to(z))
                log_pw = standard_normal_logprob(w).view(B, -1).sum(1, keepdim=True)
                delta_log_pw = delta_log_pw.view(B, 1)
                log_pz[:, i] = log_pw - delta_log_pw * valid_id[:, i].unsqueeze(1)
            log_pz = log_pz.sum(1, keepdim=True) / valid_id.sum(1, keepdim=True)
        else:
            log_pz = torch.zeros(B, 1).to(z)

        # Compute the reconstruction likelihood P(X|z)
        z_new = z.view(*z.size())
        z_new = z_new + (log_pz * 0.).mean()
        z_new = torch.cat([z_new, mean.transpose(1,2), torch.exp(logvar).transpose(1,2)], dim=-1)
        z_new = z_new.reshape(B, -1)
        y, delta_log_py = self.point_cnf(ref, z_new, torch.zeros(B, N, 1).to(ref))
        log_py = standard_normal_logprob(y).view(B, -1).sum(1, keepdim=True)
        delta_log_py = delta_log_py.view(B, N, 1).sum(1)
        log_px = log_py - delta_log_py

        # Loss
        entropy_loss = -entropy.mean() * self.entropy_weight
        recon_loss = -log_px.mean() * self.recon_weight
        prior_loss = -log_pz.mean() * self.prior_weight

        return {
            'entropy_loss':entropy_loss,
            'recon_loss':recon_loss,
            'prior_loss':prior_loss,
            'fit_loss': fit_loss
        }

    def decode(self, z, num_points, truncate_std=None, valid_id=None):
        # transform points from the prior to a point cloud, conditioned on a shape code
        y = self.sample_gaussian((z.size(0), num_points, self.input_dim), truncate_std)
        mean, logvar = self.part_aligner(z.transpose(1,2), mask=valid_id)
        z = torch.cat([z, mean.transpose(1,2), torch.exp(logvar).transpose(1,2)], dim=-1)
        x = self.point_cnf(y, z.reshape(z.size(0), -1), reverse=True).view(*y.size())
        return y, x

    def sample(self, batch_size, num_points, truncate_std=None, truncate_std_latent=None, gpu=None, valid_id=None):
        assert self.use_latent_flow, "Sampling requires `self.use_latent_flow` to be True."
        # Generate the shape code from the prior
        z = torch.zeros([batch_size, self.n_class, self.zdim]).to(gpu)
        for i in range(self.n_class):
            w = self.sample_gaussian((batch_size, self.zdim), truncate_std_latent, gpu=gpu)
            z[:, i] = self.latent_cnf[i](w, None, reverse=True).view(*w.size())
        # Sample points conditioned on the shape code
        y = self.sample_gaussian((batch_size, num_points, self.input_dim), truncate_std, gpu=gpu)
        mean, logvar = self.part_aligner(z.transpose(1,2), mask=valid_id)
        z_new = torch.cat([z, mean.transpose(1,2), torch.exp(logvar).transpose(1,2)], dim=-1)
        
        x = self.point_cnf(y, z_new.reshape(batch_size, -1), reverse=True).view(*y.size())
        return z, x

    def reconstruct(self, x, num_points=None, truncate_std=None, valid_id=None):
        num_points = x.size(1) if num_points is None else num_points
        z = self.encode(x)
        y, x = self.decode(z, num_points, truncate_std, valid_id=valid_id)
        return y, x

python/difffacto/models/networks/flow.py
diffFacto/diffFacto
from .odefunc import ODEfunc, ODEnet
from .normalization import MovingBatchNorm1d
from .cnf import CNF, SequentialFlow


def count_nfe(model):
    class AccNumEvals(object):

        def __init__(self):
            self.num_evals = 0

        def __call__(self, module):
            if isinstance(module, CNF):
                self.num_evals += module.num_evals()

    accumulator = AccNumEvals()
    model.apply(accumulator)
    return accumulator.num_evals


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def count_total_time(model):
    class Accumulator(object):

        def __init__(self):
            self.total_time = 0

        def __call__(self, module):
            if isinstance(module, CNF):
                self.total_time = self.total_time + module.sqrt_end_time * module.sqrt_end_time

    accumulator = Accumulator()
    model.apply(accumulator)
    return accumulator.total_time


def build_model(args, input_dim, hidden_dims, context_dim, num_blocks, conditional):
    def build_cnf():
        diffeq = ODEnet(
            hidden_dims=hidden_dims,
            input_shape=(input_dim,),
            context_dim=context_dim,
            layer_type=args['layer_type'],
            nonlinearity=args['nonlinearity'],
        )
        odefunc = ODEfunc(
            diffeq=diffeq,
        )
        cnf = CNF(
            odefunc=odefunc,
            T=args['time_length'],
            train_T=args['train_T'],
            conditional=conditional,
            solver=args['solver'],
            use_adjoint=args['use_adjoint'],
            atol=args['atol'],
            rtol=args['rtol'],
        )
        return cnf

    chain = [build_cnf() for _ in range(num_blocks)]
    if args['batch_norm']:
        bn_layers = [MovingBatchNorm1d(input_dim, bn_lag=args['bn_lag'], sync=args['sync_bn'])
                     for _ in range(num_blocks)]
        bn_chain = [MovingBatchNorm1d(input_dim, bn_lag=args['bn_lag'], sync=args['sync_bn'])]
        for a, b in zip(chain, bn_layers):
            bn_chain.append(a)
            bn_chain.append(b)
        chain = bn_chain
    model = SequentialFlow(chain)

    return model


def get_point_cnf(args):
    dims = tuple(map(int, args['dims'].split("-")))
    model = build_model(args, args['input_dim'], dims, args['cond_dim'] * args['n_class'], args['num_blocks'], True).cuda()
    print("Number of trainable parameters of Point CNF: {}".format(count_parameters(model)))
    return model


def get_latent_cnf(args):
    dims = tuple(map(int, args['latent_dims'].split("-")))
    model = build_model(args, args['zdim'], dims, 0, args['latent_num_blocks'], False).cuda()
    print("Number of trainable parameters of Latent CNF: {}".format(count_parameters(model)))
    return model

python/difffacto/models/networks/flow_utils/flow.py
diffFacto/diffFacto
import torch
import torch.nn as nn
from torch.nn import Parameter
import torch.distributed as dist

def reduce_tensor(tensor, world_size=None):
    rt = tensor.clone()
    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
    if world_size is None:
        world_size = dist.get_world_size()

    rt /= world_size
    return rt


__all__ = ['MovingBatchNorm1d']


class MovingBatchNormNd(nn.Module):
    def __init__(self, num_features, eps=1e-4, decay=0.1, bn_lag=0., affine=True, sync=False):
        super(MovingBatchNormNd, self).__init__()
        self.num_features = num_features
        self.sync = sync
        self.affine = affine
        self.eps = eps
        self.decay = decay
        self.bn_lag = bn_lag
        self.register_buffer('step', torch.zeros(1))
        if self.affine:
            self.weight = Parameter(torch.Tensor(num_features))
            self.bias = Parameter(torch.Tensor(num_features))
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)
        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))
        self.reset_parameters()

    @property
    def shape(self):
        raise NotImplementedError

    def reset_parameters(self):
        self.running_mean.zero_()
        self.running_var.fill_(1)
        if self.affine:
            self.weight.data.zero_()
            self.bias.data.zero_()

    def forward(self, x, c=None, logpx=None, reverse=False):
        if reverse:
            return self._reverse(x, logpx)
        else:
            return self._forward(x, logpx)

    def _forward(self, x, logpx=None):
        num_channels = x.size(-1)
        used_mean = self.running_mean.clone().detach()
        used_var = self.running_var.clone().detach()

        if self.training:
            # compute batch statistics
            x_t = x.transpose(0, 1).reshape(num_channels, -1)
            batch_mean = torch.mean(x_t, dim=1)

            if self.sync:
                batch_ex2 = torch.mean(x_t**2, dim=1)
                batch_mean = reduce_tensor(batch_mean)
                batch_ex2 = reduce_tensor(batch_ex2)
                batch_var = batch_ex2 - batch_mean**2
            else:
                batch_var = torch.var(x_t, dim=1)

            # moving average
            if self.bn_lag > 0:
                used_mean = batch_mean - (1 - self.bn_lag) * (batch_mean - used_mean.detach())
                used_mean /= (1. - self.bn_lag**(self.step[0] + 1))
                used_var = batch_var - (1 - self.bn_lag) * (batch_var - used_var.detach())
                used_var /= (1. - self.bn_lag**(self.step[0] + 1))

            # update running estimates
            self.running_mean -= self.decay * (self.running_mean - batch_mean.data)
            self.running_var -= self.decay * (self.running_var - batch_var.data)
            self.step += 1

        # perform normalization
        used_mean = used_mean.view(*self.shape).expand_as(x)
        used_var = used_var.view(*self.shape).expand_as(x)

        y = (x - used_mean) * torch.exp(-0.5 * torch.log(used_var + self.eps))

        if self.affine:
            weight = self.weight.view(*self.shape).expand_as(x)
            bias = self.bias.view(*self.shape).expand_as(x)
            y = y * torch.exp(weight) + bias

        if logpx is None:
            return y
        else:
            return y, logpx - self._logdetgrad(x, used_var).sum(-1, keepdim=True)

    def _reverse(self, y, logpy=None):
        used_mean = self.running_mean
        used_var = self.running_var

        if self.affine:
            weight = self.weight.view(*self.shape).expand_as(y)
            bias = self.bias.view(*self.shape).expand_as(y)
            y = (y - bias) * torch.exp(-weight)

        used_mean = used_mean.view(*self.shape).expand_as(y)
        used_var = used_var.view(*self.shape).expand_as(y)
        x = y * torch.exp(0.5 * torch.log(used_var + self.eps)) + used_mean

        if logpy is None:
            return x
        else:
            return x, logpy + self._logdetgrad(x, used_var).sum(-1, keepdim=True)

    def _logdetgrad(self, x, used_var):
        logdetgrad = -0.5 * torch.log(used_var + self.eps)
        if self.affine:
            weight = self.weight.view(*self.shape).expand(*x.size())
            logdetgrad += weight
        return logdetgrad

    def __repr__(self):
        return (
            '{name}({num_features}, eps={eps}, decay={decay}, bn_lag={bn_lag},'
            ' affine={affine})'.format(name=self.__class__.__name__, **self.__dict__)
        )


def stable_var(x, mean=None, dim=1):
    if mean is None:
        mean = x.mean(dim, keepdim=True)
    mean = mean.view(-1, 1)
    res = torch.pow(x - mean, 2)
    max_sqr = torch.max(res, dim, keepdim=True)[0]
    var = torch.mean(res / max_sqr, 1, keepdim=True) * max_sqr
    var = var.view(-1)
    # change nan to zero
    var[var != var] = 0
    return var


class MovingBatchNorm1d(MovingBatchNormNd):
    @property
    def shape(self):
        return [1, -1]

    def forward(self, x, context=None, logpx=None, integration_times=None, reverse=False):
        ret = super(MovingBatchNorm1d, self).forward(
                x, context, logpx=logpx, reverse=reverse)
        return ret

python/difffacto/models/networks/flow_utils/normalization.py
diffFacto/diffFacto
import torch 
import pickle
def tokenizing(text: str):
    """
    Input:
        word2int: dict
        text: str
    Output:
        tensor of token: [len_seq]
    """
    word2int = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/python/anchor_diff/models/networks/language_utils/word2int.pkl", "rb"))

    token = list(map(lambda x: word2int[x], text.split(" ")))
    token = torch.tensor(token)

    return token
python/difffacto/models/networks/language_utils/language_util.py
diffFacto/diffFacto
import pickle 
from six.moves import cPickle
import os.path as osp

def unpickle_data(file_name, python2_to_3=False):
    """Restore data previously saved with pickle_data().
    :param file_name: file holding the pickled data.
    :param python2_to_3: (boolean), if True, pickle happened under python2x, unpickling under python3x.
    :return: a generator over the un-pickled items.
    Note, about implementing the python2_to_3 see
        https://stackoverflow.com/questions/28218466/unpickling-a-python-2-object-with-python-3
    """

    in_file = open(file_name, "rb")
    if python2_to_3:
        size = cPickle.load(in_file, encoding="latin1")
    else:
        size = cPickle.load(in_file)

    for _ in range(size):
        if python2_to_3:
            yield cPickle.load(in_file, encoding="latin1")
        else:
            yield cPickle.load(in_file)
    in_file.close()


def process():
    (
        game_data,
        word2int,
        int2word,
        int2sn,
        sn2int,
        sorted_sn,
    ) = unpickle_data("/mnt/disk3/wang/diffusion/datasets/partglot_data/game_data.pkl")
    print(len(word2int))
    with open("/mnt/disk3/wang/diffusion/anchorDIff/python/anchor_diff/models/encoders/language_utils/word2int.pkl", "wb") as f:
        pickle.dump(word2int, f)
    with open("/mnt/disk3/wang/diffusion/anchorDIff/python/anchor_diff/models/encoders/language_utils/int2word.pkl", "wb") as ff:
        pickle.dump(int2word, ff)

process()
python/difffacto/models/networks/language_utils/process.py
diffFacto/diffFacto
from distutils.command.build import build
import torch
import torch.nn as nn 
from torch.nn import Module, functional
from difffacto.utils.misc import *
from difffacto.utils.registry import build_from_cfg, MODELS, ENCODERS, DIFFUSIONS, SAMPLERS, METRICS, DECOMPOSERS, SEGMENTORS
from pointnet2_ops.pointnet2_utils import gather_operation
from .language_utils.language_util import tokenizing
VOCAB_SIZE=2787

@MODELS.register_module()
class LatentDiffEdit(Module):

    def __init__(
        self, 
        encoder, 
        diffusion, 
        sampler, 
        num_anchors,
        num_timesteps, 
        guidance=False,
        language_encoder=None,
        pcd_encoder_ckpt=None,
        pcd_mixer=None,
        pcd_mixer_ckpt=None,
        part_latent_dropout_prob=0.2,
        language_latent_dropout_prob=0.2,
        zero_part_latent=False,

        save_dir=None,
        save_weights=False
        ):
        super().__init__()
        self.encoder = build_from_cfg(encoder, ENCODERS, num_anchors=num_anchors)
        self.diffusion = build_from_cfg(diffusion, DIFFUSIONS, num_timesteps=num_timesteps, num_part=num_anchors)
        self.encoder.load_state_dict(torch.load(pcd_encoder_ckpt))
        for m in self.encoder.parameters():
            m.requires_grad = False
        self.sampler = build_from_cfg(sampler, SAMPLERS, num_timesteps=num_timesteps)
        self.language_encoder = build_from_cfg(language_encoder, ENCODERS, vocab_size=VOCAB_SIZE)
        self.pcd_mixer = build_from_cfg(pcd_mixer,DECOMPOSERS, num_anchors=num_anchors, point_dim=3)
        if self.pcd_mixer is not None:
            self.pcd_mixer.load_state_dict(torch.load(pcd_mixer_ckpt))
        for m in self.pcd_mixer.parameters():
            m.requires_grad = False
        self.zero_part_latent=zero_part_latent
        self.guidance=guidance
        self.part_latent_dropout_prob=part_latent_dropout_prob
        self.language_latent_dropout_prob=language_latent_dropout_prob
        self.save_weights=save_weights
        self.save_dir=save_dir
        self.num_timesteps = int(num_timesteps)
        self.num_anchors = num_anchors
        
    def decode(self, prior, code=None, part_indicator=None, anchors=None, noise=None, device='cuda'):
        final = dict() 
        for t, sample in self.diffusion.p_sample_loop_progressive(
            prior.shape,
            prior=prior,
            code=code,
            part_indicator=part_indicator,
            anchors=anchors,
            noise=noise,
            device=device,
            progress=False,
        ):
            if t == 0:
                final['pred'] = sample['sample'].transpose(2, 1)
        return final['pred']
        
    def forward(self, pcds, device='cuda', iter=-1):
        """
            input: [B, npoints, 3]
            attn_map: [B, npoints, n_class]
            seg_mask: [B, npoints]
        """
        if not self.training and self.save_weights:
            self.eval()
            os.system(f"mkdir -p {self.save_dir}")
            torch.save(self.language_encoder.state_dict(), os.path.join(self.save_dir, "language_encoder_ckpt.pth"))
            torch.save(self.diffusion.state_dict(), os.path.join(self.save_dir, "ldm_ckpt.pth"))
            return dict(pred=pcds['target'],distractor=pcds['target'], target=pcds['target'], shift=pcds['target_shift'], scale=pcds['target_scale'])
        target = pcds['target'].to(device)
        distractors = pcds['distractor'].to(device)
        target_attn_map = pcds['target_attn_map'].to(device)
        distractor_attn_map = pcds['distractor_attn_map'].to(device)
        part_indicator = pcds['part_indicator'].to(device)
        text = pcds['text'].to(device)
        
        loss_dict = dict()
               
        input_pcds = torch.stack([target, distractors], dim=1) # [B, 2, npoint, 3]
        input_attn_map = torch.stack([target_attn_map, distractor_attn_map], dim=1)
        x = torch.cat([input_pcds, input_attn_map], dim=-1) # -> [B, 2, npoints, n_class + 3]

        B, K, npoint = input_pcds.shape[:-1]
        part_latent = self.encoder(x.reshape(B*2, npoint, -1))
        if self.pcd_mixer is not None:
            _, anchors, _ = self.pcd_mixer(part_latent)
        if iter == 0:
            self.diffusion.on_train_batch_start(part_latent, iter)
        target_part_latent, distractor_part_latent = torch.split(part_latent.reshape(B, 2, self.num_anchors, -1), 1, dim=1)
        target_anchors, distractor_anchors = torch.split(anchors.reshape(B, 2, self.num_anchors, 3), 1, dim=1)
        language_part_latent, language_enc_attn = self.language_encoder(text)
        target_part_latent, distractor_part_latent, target_anchors, distractor_anchors = map(lambda t: t.squeeze(1).transpose(1,2), [target_part_latent, distractor_part_latent, target_anchors, distractor_anchors])
        if len(language_part_latent.shape) == 3:
            language_part_latent = language_part_latent.transpose(1,2)
        if self.training:
            if self.guidance:
                idxs = torch.rand(distractor_part_latent.shape[0], distractor_part_latent.shape[2]).to(distractor_part_latent.device) >= self.part_latent_dropout_prob
                distractor_part_latent = distractor_part_latent * idxs.unsqueeze(1)

                idxxs = torch.rand(language_part_latent.shape[0], language_part_latent.shape[2]).to(language_part_latent.device) >= self.language_latent_dropout_prob
                language_part_latent = language_part_latent * idxxs.unsqueeze(1)
            if self.zero_part_latent:
                distractor_part_latent = torch.zeros_like(distractor_part_latent)
            t, _ = self.sampler.sample(B, device)
            diffusion_loss = self.diffusion.training_losses(target_part_latent, t, prior=distractor_part_latent, code=language_part_latent, part_indicator=part_indicator, anchors=distractor_anchors)
            loss_dict.update(diffusion_loss)
            return loss_dict
        else:
            pred = self.decode(prior=language_part_latent, code=language_part_latent, part_indicator=part_indicator, anchors=distractor_anchors, device=device) 
            return pred

python/difffacto/models/networks/latent_diffusion.py
diffFacto/diffFacto
from distutils.command.build import build
import torch 
import torch.nn as nn 
from difffacto.utils.registry import MODELS, ENCODERS, METRICS, DECOMPOSERS, build_from_cfg
import pickle
import os
VOCAB_SIZE=2787
part_names = ["back", "seat", "leg", "arm"]

@MODELS.register_module()
class PartglotClassifier(nn.Module):
    def __init__(self, text_dim, embedding_dim, classifier_encoder, pcd_encoder, pcd_mixer, cross_attention, loss, num_part_latent, pcd_encoder_ckpt_path,pcd_mixer_ckpt_path,  out_dim):
        super().__init__()
        word2int = pickle.load(open("/mnt/disk3/wang/diffusion/anchorDIff/python/difffacto/models/networks/language_utils/word2int.pkl", "rb"))
        pn_tensor = torch.tensor([word2int[x] for x in part_names])[None]
        self.register_buffer("pn_tensor", pn_tensor)
        self.clsf_encoder = build_from_cfg(classifier_encoder, ENCODERS, vocab_size=VOCAB_SIZE)
        self.attn_encoder = nn.Sequential(self.clsf_encoder.word_embedding, nn.Linear(embedding_dim, text_dim))
        self.pcd_encoder = build_from_cfg(pcd_encoder, ENCODERS, num_anchors=num_part_latent)
        self.pcd_mixer = build_from_cfg(pcd_mixer, DECOMPOSERS, num_anchors=num_part_latent, point_dim=3)
        self.cross_attention = build_from_cfg(cross_attention, ENCODERS)
        self.loss = build_from_cfg(loss, METRICS)
        self.pcd_encoder.load_state_dict(torch.load(pcd_encoder_ckpt_path))
        # freeze pcd encoder
        for p in self.pcd_encoder.parameters():
            p.requires_grad=False
        self.pcd_mixer.load_state_dict(torch.load(pcd_mixer_ckpt_path))
        # freeze pcd mixer
        for p in self.pcd_mixer.parameters():
            p.requires_grad=False

        self.mlp = nn.Sequential(
            nn.Linear(out_dim, 512, bias=True), nn.BatchNorm1d(512), nn.ReLU(inplace=True),
            nn.Linear(512, 128, bias=True), nn.BatchNorm1d(128), nn.ReLU(inplace=True),
            nn.Linear(128, 1, bias=True)
        )
        
    def forward(self, data, device='cuda'):
        # torch.save(self.clsf_encoder.state_dict(), "/mnt/disk3/wang/diffusion/anchorDIff/work_dirs/pn_aware_attn_anchor_feat_listener/checkpoints/classifier_encoder.pth")
        # torch.save(self.attn_encoder.state_dict(), "/mnt/disk3/wang/diffusion/anchorDIff/work_dirs/pn_aware_attn_anchor_feat_listener/checkpoints/attention_encoder.pth")
        # torch.save(self.mlp.state_dict(), "/mnt/disk3/wang/diffusion/anchorDIff/work_dirs/pn_aware_attn_anchor_feat_listener/checkpoints/listener_mlp.pth")
        # exit()
        pn_tensor = self.pn_tensor.to(device)
        target_pcds = data['target'].to(device)
        target_attn_maps = data['target_attn_map'].to(device)
        distractor_pcds = data['distractor'].to(device)
        distractor_attn_maps = data['distractor_attn_map'].to(device)
        part_indicators = data['part_indicator'].to(device).repeat_interleave(2, dim=0)
        texts = data['text'].to(device)
        B, npoint = target_pcds.shape[:-1]
        pcds = torch.stack([target_pcds, distractor_pcds], dim=1).reshape(B*2, npoint, 3)
        attn_maps = torch.stack([target_attn_maps, distractor_attn_maps], dim=1).reshape(B*2, npoint, -1)
        input = torch.cat([pcds, attn_maps], dim=-1)
        pcd_f = self.pcd_encoder(input)
        _, _, part_latent = self.pcd_mixer(pcd_f)

        attn_enc_f = self.attn_encoder(pn_tensor).repeat(B*2, 1, 1)

        clsf_enc_f, clsf_enc_attn = self.clsf_encoder(texts)
        clsf_enc_f = clsf_enc_f.repeat_interleave(2, dim=0).unsqueeze(1)

        cross_attn_f = self.cross_attention(attn_enc_f, part_latent, part_indicators)
        logits = self.mlp(torch.cat([clsf_enc_f, cross_attn_f], -1).squeeze(1)).reshape(B, 2)
        targets = torch.zeros(B, dtype=torch.int64).to(device)
        loss = self.loss(logits, targets)
        preds = torch.max(logits, 1)[1]
        num_correct = int(preds.shape[0] - preds.sum())
        if self.training:
            return dict(loss=loss) 
        
        return dict(
            logits=logits,
            num_correct=num_correct, 
            target=target_pcds, 
            distractor=distractor_pcds, 
            target_attn_map=target_attn_maps, 
            distractor_attn_map=distractor_attn_maps, 
            text=texts, 
            target_shift=data['target_shift'], 
            target_scale=data['target_scale'], 
            distractor_shift=data['distractor_shift'], 
            distractor_scale=data['distractor_scale'],
        )
python/difffacto/models/networks/listener.py
diffFacto/diffFacto
import torch
import numpy as np
from torch import nn
from difffacto.utils.registry import build_from_cfg, MODELS, ENCODERS
from difffacto.utils.misc import standard_normal_logprob, truncated_normal, gaussian_log_likelihood
from pointnet2_ops.pointnet2_utils import gather_operation
from .flow_utils.flow import get_point_cnf, get_latent_cnf




@MODELS.register_module()
class Aligner(nn.Module):
    def __init__(self, 
                 encoder,
                n_class
                 ):
        super(Aligner, self).__init__()
        self.encoder = build_from_cfg(encoder, ENCODERS)
        self.regressor = nn.Sequential(nn.Linear(self.encoder.zdim, 512), nn.ReLU(), 
                                       nn.Linear(512, 512), nn.ReLU(),
                                       nn.Linear(512, 512), nn.ReLU(),
                                       nn.Linear(512, 512), nn.ReLU(),
                                       nn.Linear(512, 64), nn.ReLU(),
                                       nn.Linear(64, 6))
        self.n_class=n_class
    def forward(self, pcds, device='cuda', **kwargs):
        ref = pcds['ref'].to(device)
        seg_mask = pcds['ref_seg_mask'].to(device)
        seg_flag = pcds['ref_attn_map'].to(device)
        valid_id = pcds.get('present', None)
        if valid_id is not None:
            valid_id = valid_id.to(device)

        loss_dict = dict()
        B, N, C = ref.shape

        codes = self.encoder(torch.cat([ref, seg_flag], dim=-1)).reshape(B, self.n_class, -1)
        params = self.regressor(codes)
        param_per_point = gather_operation(torch.cat([params.transpose(1,2), valid_id.unsqueeze(1)], dim=1).contiguous(), seg_mask.to(int)).reshape(B, 6+1, N)
        mean, logvar, flag = torch.split(param_per_point, [3, 3, 1], dim=1)
        fit_loss = - gaussian_log_likelihood(ref, mean, logvar).mean(1, keepdim=True) * flag
        fit_loss = fit_loss.sum(dim=(-1, -2)) / flag.sum(dim=(-1, -2))
        return dict(loss=fit_loss) if self.training else [(dict(index=pcds['id'] , mean=params[..., :3], logvar=params[..., 3:], loss=fit_loss), "params")]


python/difffacto/models/networks/parameter_regressor.py
diffFacto/diffFacto
from distutils.command.build import build
import torch 
import torch.nn as nn 
import torch.nn.functional as F
from difffacto.utils.registry import MODELS, ENCODERS, METRICS, DECOMPOSERS, build_from_cfg
import pickle
import os
from pointnet2_ops.pointnet2_utils import gather_operation
VOCAB_SIZE=2787
part_names = ["back", "seat", "leg", "arm"]

@MODELS.register_module()
class PartLatentEditor(nn.Module):
    def __init__(
        self, 
        text_dim, 
        embedding_dim, 
        part_latent_dim, 
        language_encoder, 
        latent_language_fuser,
        classifier_encoder, 
        pcd_encoder, 
        pcd_mixer, 
        cross_attention, 
        loss, 
        num_part_latent, 
        pcd_encoder_ckpt_path, 
        pcd_mixer_ckpt_path, 
        attn_encoder_ckpt_path, 
        classifier_encoder_ckpt_path, 
        listener_mlp_ckpt_path, 
        out_dim,
        direction_loss,
        magnitude_loss,
        logit_loss,
        logit_weight=1.0,
        direction_weight=1.0,
        magnitude_weight=1.0,
        supervise_before_mixer=False,
        icmle=False,
        conditional_dim=24,
        num_coditional_sample=10,
        save_weights=False,
        save_dir=None,
    ):
        super().__init__()
        self.part_latent_dim=part_latent_dim
        self.direction_loss=direction_loss
        self.magnitude_loss=magnitude_loss
        self.logit_loss=logit_loss
        self.num_coditional_sample=num_coditional_sample
        self.logit_weight=logit_weight
        self.direction_weight=direction_weight
        self.save_weights=save_weights
        self.icmle=icmle
        self.save_dir=save_dir
        self.conditional_dim=conditional_dim
        self.magnitude_weight=magnitude_weight
        self.supervise_before_mixer=supervise_before_mixer
        word2int = pickle.load(open("/mnt/disk3/wang/diffusion/anchorDIff/python/difffacto/models/networks/language_utils/word2int.pkl", "rb"))
        pn_tensor = torch.tensor([word2int[x] for x in part_names])[None]
        self.register_buffer("pn_tensor", pn_tensor)
        self.language_encoder = build_from_cfg(language_encoder, ENCODERS, vocab_size=VOCAB_SIZE)
        self.latent_language_fuser = build_from_cfg(latent_language_fuser, ENCODERS, num_part=num_part_latent, conditional=icmle, conditional_dim=conditional_dim)
        self.clsf_encoder = build_from_cfg(classifier_encoder, ENCODERS, vocab_size=VOCAB_SIZE)
        self.attn_encoder = nn.Sequential(self.clsf_encoder.word_embedding, nn.Linear(embedding_dim, text_dim))
        self.pcd_encoder = build_from_cfg(pcd_encoder, ENCODERS, num_anchors=num_part_latent)
        self.pcd_mixer = build_from_cfg(pcd_mixer, DECOMPOSERS, num_anchors=num_part_latent, point_dim=3)
        self.cross_attention = build_from_cfg(cross_attention, ENCODERS)
        self.loss = build_from_cfg(loss, METRICS)
        self.mlp = nn.Sequential(
            nn.Linear(out_dim, 512, bias=True), nn.BatchNorm1d(512), nn.ReLU(inplace=True),
            nn.Linear(512, 128, bias=True), nn.BatchNorm1d(128), nn.ReLU(inplace=True),
            nn.Linear(128, 1, bias=True)
        )

        # freeze mlp
        self.mlp.load_state_dict(torch.load(listener_mlp_ckpt_path))
        for p in self.mlp.parameters():
            p.requires_grad=False
        # freeze attn encoder
        self.attn_encoder.load_state_dict(torch.load(attn_encoder_ckpt_path))
        for p in self.attn_encoder.parameters():
            p.requires_grad=False
        # freeze classifier encoder
        self.clsf_encoder.load_state_dict(torch.load(classifier_encoder_ckpt_path))
        for p in self.clsf_encoder.parameters():
            p.requires_grad=False
        self.pcd_encoder.load_state_dict(torch.load(pcd_encoder_ckpt_path))
        # freeze pcd encoder
        for p in self.pcd_encoder.parameters():
            p.requires_grad=False
        self.pcd_mixer.load_state_dict(torch.load(pcd_mixer_ckpt_path))
        # freeze pcd mixer
        for p in self.pcd_mixer.parameters():
            p.requires_grad=False

        
        
    def forward(self, data, device='cuda'):
        if self.save_weights:
            torch.save(self.language_encoder.state_dict(), os.path.join(self.save_dir, "language_encoder.pth"))
            torch.save(self.latent_language_fuser.state_dict(), os.path.join(self.save_dir, "update_mlp.pth"))
            exit()
        loss_dict = dict()
        pn_tensor = self.pn_tensor.to(device)
        target_pcds = data['target'].to(device)
        target_attn_maps = data['target_attn_map'].to(device)
        distractor_pcds = data['distractor'].to(device)
        distractor_attn_maps = data['distractor_attn_map'].to(device)
        part_indicators = data['part_indicator'].to(device)
        texts = data['text'].to(device)
        B, npoint = target_pcds.shape[:-1]
        pcds = torch.stack([target_pcds, distractor_pcds], dim=1).reshape(B*2, npoint, 3)
        attn_maps = torch.stack([target_attn_maps, distractor_attn_maps], dim=1).reshape(B*2, npoint, -1)
        input = torch.cat([pcds, attn_maps], dim=-1)
        part_latent = self.pcd_encoder(input)
        

        n_part, n_dim = part_latent.shape[1:]
        tgt_latent, latent_for_edit = part_latent.reshape(B, 2, n_part, n_dim).split(1, dim=1)
        tgt_latent, latent_for_edit = tgt_latent.reshape(B, n_part, n_dim), latent_for_edit.reshape(B, n_part, n_dim)
        clsf_enc_f, clsf_enc_attn = self.clsf_encoder(texts)
        language_enc_f, language_enc_attn = self.language_encoder(texts)
        if self.icmle:
            K=self.num_coditional_sample
            conditional = torch.randn(B, K, self.conditional_dim).cuda()
        else:
            K=1
            conditional = None
        clsf_enc_f = clsf_enc_f.repeat_interleave(K, dim=0)
        tgt_latent = tgt_latent.unsqueeze(1).repeat_interleave(K, dim=1)
        edited_latent = self.latent_language_fuser(latent_for_edit, part_indicators, language_enc_f, conditional=conditional)
        
        part_id = part_indicators.repeat_interleave(2*K, dim=0).max(1)[1].to(torch.int)
        part_latent_pre_attn = torch.stack([tgt_latent, edited_latent], dim=1).reshape(B*2*K, n_part, n_dim)
        
        _, _, part_latent_attn_out = self.pcd_mixer(part_latent_pre_attn)
        if self.supervise_before_mixer:
            referenced_part_latent = gather_operation(part_latent_pre_attn.reshape(B*2*K, n_part, n_dim).transpose(1,2).contiguous(), part_id.reshape(B*2*K, 1)).reshape(B, 2, K, 1, -1)
        else:
            referenced_part_latent = gather_operation(part_latent_attn_out.reshape(B*2*K, n_part, n_dim).transpose(1,2).contiguous(), part_id.reshape(B*2*K, 1)).reshape(B, 2, K, 1, -1)
            
        tgt_latent, edited_latent = referenced_part_latent.reshape(B, 2, K, 1, n_dim).split(1, dim=1)
            
        tgt_latent, edited_latent = tgt_latent.reshape(B*K, n_dim), edited_latent.reshape(B*K, n_dim)
        loss = torch.zeros([B, K]).cuda()
        if self.direction_loss:
            cos_sim = F.cosine_similarity(tgt_latent, edited_latent, dim=1)
            dir_loss = 1 - cos_sim
            loss += dir_loss.reshape(B, K)
            loss_dict["direction_loss"] = self.direction_weight * dir_loss
        if self.magnitude_loss:
            tgt_norm, src_norm = torch.norm(tgt_latent, p=2, dim=1), torch.norm(edited_latent, p=2, dim=1)
            mag_loss = (tgt_norm - src_norm) ** 2
            loss += mag_loss.reshape(B, K)
            loss_dict["magnitude_loss"] = self.magnitude_weight * mag_loss
        if self.logit_loss:
            modified_logits = self.mlp(torch.cat([clsf_enc_f, edited_latent], -1))
            logits = self.mlp(torch.cat([clsf_enc_f, tgt_latent], -1))
            logit_loss = self.loss(modified_logits, logits)
            loss += logit_loss.reshape(B, K)
            loss_dict['logit_loss'] = self.logit_weight*logit_loss
            
        min_idx = loss.min(dim=1)[1]
        loss_dict = {k: torch.gather(v.reshape(B, K), 1, min_idx.unsqueeze(1)).mean() for k, v in loss_dict.items()}
        
        if self.training:
            return loss_dict

        return dict(
            num_correct=0,
            target=target_pcds, 
            distractor=distractor_pcds, 
            target_attn_map=target_attn_maps, 
            distractor_attn_map=distractor_attn_maps, 
            text=texts, 
            target_shift=data['target_shift'], 
            target_scale=data['target_scale'], 
            distractor_shift=data['distractor_shift'], 
            distractor_scale=data['distractor_scale'],
            modified_logits=modified_logits if self.logit_loss else torch.zeros(B, 2),
            logits=logits if self.logit_loss else torch.zeros(B, 2)
        )
python/difffacto/models/networks/part_editing.py
diffFacto/diffFacto
import torch
from torch.nn import Module, functional
from difffacto.utils.misc import *
from difffacto.utils.registry import build_from_cfg, MODELS, ENCODERS, DIFFUSIONS, SAMPLERS, METRICS, DECOMPOSERS, SEGMENTORS
from pointnet2_ops.pointnet2_utils import gather_operation

@MODELS.register_module()
class PartglotFinetune(Module):

    def __init__(
        self, 
        encoder, 
        diffusion, 
        sampler, 
        num_anchors,
        num_timesteps, 
        npoints=2048, 
        ):
        super().__init__()
        self.sup_segs_encoder = build_from_cfg(encoder, ENCODERS, num_anchors=num_anchors)
        self.diffusion = build_from_cfg(diffusion, DIFFUSIONS, num_timesteps=num_timesteps, learn_variance=False)
        self.sampler = build_from_cfg(sampler, SAMPLERS, num_timesteps=num_timesteps)
        self.num_timesteps = int(num_timesteps)
        self.num_anchors = num_anchors
        self.npoints = npoints
        self.points_per_anchor = self.npoints // self.num_anchors
    


    def forward(self, pcds, device='cuda'):
        pass
python/difffacto/models/networks/partglot_finetune.py
diffFacto/diffFacto
import torch
import torch.nn as nn 
import torch.nn.functional as F
from torch.nn import Module, functional
from difffacto.utils.misc import *
from difffacto.utils.registry import build_from_cfg, MODELS, ENCODERS, DIFFUSIONS, SAMPLERS, METRICS, DECOMPOSERS, SEGMENTORS
from pointnet2_ops.pointnet2_utils import gather_operation
from .language_utils.language_util import tokenizing
VOCAB_SIZE=2787

@MODELS.register_module()
class PDM(Module):

    def __init__(
        self, 
        encoder, 
        diffusion, 
        sampler, 
        num_timesteps, 
        npoints=2048, 
        include_attn_weight_in_encoder=True,

        ret_traj=False, 
        ret_interval=20, 
        interpolate=False,
        combine=False,
        drift_anchors=False,
        save_pred_xstart=False,
        partglot_dataset=False,
        save_dir=None,
        save_weights=False,
        ):
        super().__init__()
        self.encoder = build_from_cfg(encoder, ENCODERS)
        self.diffusion = build_from_cfg(diffusion, DIFFUSIONS, num_timesteps=num_timesteps)
        self.sampler = build_from_cfg(sampler, SAMPLERS, num_timesteps=num_timesteps)
        self.save_weights=save_weights
        self.save_dir=save_dir
        self.drift_anchors=drift_anchors
        self.include_attn_weight_in_encoder=include_attn_weight_in_encoder
        self.num_timesteps = int(num_timesteps)     
        self.partglot_dataset=partglot_dataset
        
        self.npoints = npoints
        self.interpolate=interpolate
        self.ret_traj=ret_traj
        self.ret_interval=ret_interval
        self.save_pred_xstart=save_pred_xstart
        self.combine=combine

        print(self)

    
    def decode(self, anchors, code=None, pointwise_latent=None, noise=None, variance=None, device='cuda'):
        final = dict() 
        bs = anchors.shape[0]
        for t, sample in self.diffusion.p_sample_loop_progressive(
            [bs, 3, self.npoints],
            anchors=anchors,
            code=code,
            variance=variance,
            pointwise_latent=pointwise_latent,
            noise=noise,
            device=device,
            progress=False,
        ):
            if t == 0:
                final['pred'] = sample['sample'].transpose(2, 1)
                final['pred_xstart'] = sample['pred_xstart'].transpose(1,2) 
            elif (self.ret_traj and t % self.ret_interval == 0):
                final[t] = sample['sample'].transpose(2, 1) 
                final[f'pred_xstart_{t}'] = sample['pred_xstart'].transpose(1,2) 
                
        return final

    def q_sample(self, gt, anchors, noise=None, device='cuda',variance=None):
        final = {'pred':gt.transpose(2,1)}
        
        for t, sample in self.diffusion.q_sample_loop_progressive(
            gt_pcd=gt,
            anchors=anchors,
            noise=noise,
            device=device,
            variance=variance,
            progress=False,
        ):
            if t == self.num_timesteps - 1:
                final[t] = sample.transpose(2, 1)
            elif (self.ret_traj and t % self.ret_interval == 0):
                final[t] = sample.transpose(2, 1)
                
        return final

    def interpolate_latent(self, x, seg_mask, device, pcds):
        ids = [(19,20, 2), (17, 15, 0), (37, 38, 1), (122, 120, 1), (19, 18, 2), (66,65, 2)]
        out_dict = {f"interpolate sample {i}":[] for i in range(10)}
        # out_dict['combined'] = []
        # out_dict['seg_mask_ref'] = []
        out_dict['input1_ref'] = []
        out_dict['input2_ref'] = []
        out_dict['seg_mask1'] = []
        out_dict['seg_mask2'] = []
        out_dict.update({f"interpolate anchor sample {i}": [] for i in range(10)})
        
        for id1, id2, anchor_id in ids:
            gt1, gt2 = x[id1], x[id2]
            latent1 = self.encoder(gt1.unsqueeze(0))
            latent2 = self.encoder(gt2.unsqueeze(0))
            latent1, latent2 = latent1.squeeze(0), latent2.squeeze(0) # [1, n_class, C]

            dx = torch.zeros([10, self.num_anchors, 1]).to(device)
            dx[:, anchor_id] += torch.linspace(0, 1, steps=10).reshape(-1, 1).to(device)
            primary_latent = latent1.unsqueeze(0) + (latent2 - latent1).unsqueeze(0) * dx
            B = primary_latent.shape[0]
            if self.project_latent:
                if self.share_projection:
                    intrinsic, extrinsic = torch.split(self.down_proj(primary_latent), [self.part_dim, self.param_dim], dim=-1)
                else:
                    intrinsic, extrinsic = torch.split(self.down_proj(primary_latent.reshape(B, -1, 1)).reshape(B, self.num_anchors, self.part_dim + self.param_dim), [self.part_dim, self.param_dim], dim=-1)
                
                if self.learn_var:
                    shift, scale = extrinsic.split(3, dim=-1)
                    scale = scale[0].unsqueeze(0).expand(B, -1, -1)
                else:
                    shift = extrinsic    
                
                shift = shift[0].unsqueeze(0).expand(B, -1, -1)
                
                if self.learn_var:
                    if self.share_projection:
                        primary_latent = intrinsic + self.up_proj(torch.cat([shift, scale], dim=-1))
                    else:
                        primary_latent = intrinsic + self.up_proj(torch.cat([shift, scale], dim=-1).reshape(B, -1, 1)).reshape(B, self.num_anchors, self.part_dim)
            
            global_feat, new_anchors, secondary_latent = self.decomposer(primary_latent)
            if self.project_latent:
                new_anchors = shift 
                if self.learn_var:
                    new_vars = scale
            if self.use_primary:
                new_latents = primary_latent 
            else:
                new_latents = secondary_latent
            latent_per_point = new_latents.reshape(10, self.num_anchors, 1, -1).expand(-1, -1, self.points_per_anchor, -1).reshape(10, self.npoints, -1)
            anchor_per_point = new_anchors.reshape(10, self.num_anchors, 1, 3).expand(-1, -1, self.points_per_anchor, -1).reshape(10, self.npoints, 3).transpose(1,2)
            if self.learn_var:
                variance_per_point = new_vars.reshape(10, self.num_anchors, 1, 3).expand(-1, -1, self.points_per_anchor, -1).reshape(10, self.npoints, 3)
                variance_per_point = torch.exp(variance_per_point)
            else:
                variance_per_point=None

            pred = self.decode(anchor_per_point, code=global_feat, pointwise_latent=latent_per_point, device=device, variance=variance_per_point) 
            # ref = torch.zeros([self.npoints, 3]).cuda()
            # _ref = torch.cat([gt1[seg_mask[id1] != anchor_id, :3], gt2[seg_mask[id2] == anchor_id, :3]], dim=0)
            # ref[:_ref.shape[0]] = _ref
            # seg_mask_ref = torch.zeros_like(seg_mask[0])
            # _seg_mask_ref = torch.cat([seg_mask[id1][seg_mask[id1] != anchor_id], seg_mask[id2][seg_mask[id2] == anchor_id]], dim=0)
            # seg_mask_ref[:_seg_mask_ref.shape[0]] = _seg_mask_ref
            for i in range(pred['pred'].shape[0]):
                out_dict[f"interpolate sample {i}"].append(pred['pred'][i])
                out_dict[f"interpolate anchor sample {i}"].append(new_anchors[i])
            # out_dict['combined'].append(ref)
            # out_dict['seg_mask_ref'].append(seg_mask_ref)
            out_dict['input1_ref'].append(gt1[:, 0:3])
            out_dict['input2_ref'].append(gt2[:, 0:3])
            out_dict['seg_mask1'].append(seg_mask[id1])
            out_dict['seg_mask2'].append(seg_mask[id2])
        out_dict = {k:torch.stack(v, dim=0) for k, v in out_dict.items()}
        out_dict['pred'] = pcds['input'][0].unsqueeze(0).repeat_interleave(len(ids), dim=0)
        out_dict['input_ref'] = pcds['input'][0].unsqueeze(0).repeat_interleave(len(ids), dim=0)
        out_dict['shift'] = pcds['shift'][0].unsqueeze(0)
        out_dict['scale'] = pcds['scale'][0].unsqueeze(0)
        return out_dict

    def part_completion(self, x, seg_mask, id1, anchor_id, device):
        gt = x[id1]
        idx = gt[..., 3 + anchor_id] == 0.
        gt_part = gt[idx]
        latent = self.encoder(xyz=gt_part.unsqueeze(0))
        new_params, new_latents = self.decomposer(latent)
        latent_per_point = new_latents.reshape(1, self.num_anchors, 1, -1).expand(-1, -1, self.points_per_anchor, -1).reshape(1, self.npoints, -1)
        param_per_point = new_params.reshape(1, self.num_anchors, 1, 3).expand(-1, -1, self.points_per_anchor, -1).reshape(1, self.npoints, 3).transpose(1,2)
        assert param_per_point.shape[1] == 3
        anchor_per_point = param_per_point
        variance_per_point = None

        pred = self.decode(anchor_per_point, code=None, pointwise_latent=latent_per_point, device=device, variance=variance_per_point) 
        pred['input'] = gt_part[:, 0:3].unsqueeze(0)
        pred['ref'] = gt[:, 0:3].unsqueeze(0)
        pred['anchors'] = new_params[..., :3]
        
        return pred

    def combine_latent(self, x, pcds, seg_mask, ids, device):
        gts = []
        for i in ids:
            gts.append(x[i])
        gts = torch.stack(gts, dim=0)
        latents = self.encoder(gts)
        latents = torch.stack([latents[i, i] for i in range(self.num_anchors)], dim=0).reshape(1, self.num_anchors, -1)
        global_feature, new_params, new_latents = self.decomposer(latents)
        latent_per_point = new_latents.reshape(1, self.num_anchors, 1, -1).expand(-1, -1, self.points_per_anchor, -1).reshape(1, self.npoints, -1)
        param_per_point = new_params.reshape(1, self.num_anchors, 1, 3).expand(-1, -1, self.points_per_anchor, -1).reshape(1, self.npoints, 3).transpose(1,2)
        assert param_per_point.shape[1] == 3
        anchor_per_point = param_per_point
        variance_per_point = None
        
        pred = self.decode(anchor_per_point, code=global_feature, pointwise_latent=latent_per_point, device=device, variance=variance_per_point) 
        ref = torch.cat([gts[i, seg_mask[ids[i]] == i, :3] for i in range(self.num_anchors)], dim=0)[:self.npoints]
        seg_mask_ref = torch.cat([seg_mask[ids[i]][seg_mask[ids[i]] == i] for i in range(self.num_anchors)], dim=0)
        print(ref.shape)
        pred['input_ref'] = ref.unsqueeze(0)
        pred['seg_mask_ref'] = seg_mask_ref.unsqueeze(0)
        for i in range(self.num_anchors):
            pred[f'input_{i}'] = gts[i, :, 0:3].unsqueeze(0)
            pred[f'seg_mask_{i}'] = seg_mask[ids[i]].unsqueeze(0)
        pred['anchors'] = new_params[..., :3]
        pred['shift'] = pcds['shift']
        pred['scale'] = pcds['scale']
        return pred

    def language_edit_step(self, data, device):
        prompts = [
            "leg",
            "four leg", 
            "one leg",
            "thick leg", 
            "thin leg", 
            "wheel", 
            "long leg", 
            "short leg", 
            "seat", 
            "thin seat", 
            "thick seat",
            "back",
            "straight back", 
            "slanted back",
            "tall back", 
            "short back",
            "square back",
            "diamond back", 
            "arm rest", 
        ]
        # ids = [2]
        ids = [2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 0, 0, 0, 0, 0,0, 0, 3]
        prompts = [f"a chair with {p}" for p in prompts]
        tokenized_prompts = [tokenizing(f"a chair with {p}") for p in prompts]
        if self.partglot_dataset:
            target_pcds = data['target'].to(device)
            target_attn_maps = data['target_attn_map'].to(device)
            distractor_pcds = data['distractor'].to(device)
            distractor_attn_maps = data['distractor_attn_map'].to(device)
            part_indicators = data['part_indicator'].to(device).repeat_interleave(3, dim=0)
            texts = data['text'].to(device)
            B, npoint = target_pcds.shape[:-1]
            pcds = torch.stack([target_pcds, distractor_pcds, distractor_pcds], dim=1).reshape(B*3, npoint, 3)
            attn_maps = torch.stack([target_attn_maps, distractor_attn_maps, distractor_attn_maps], dim=1).reshape(B*3, npoint, -1)
            input = torch.cat([pcds, attn_maps], dim=-1)
            language_enc_f, language_enc_attn = self.language_encoder(texts)
        
        
            part_latent = self.encoder(input)
            n_part, n_dim = part_latent.shape[1:]
            

            part_id = part_indicators.max(1)[1].to(torch.int)
            referenced_part_latent = gather_operation(part_latent.transpose(1,2).contiguous(), part_id.reshape(B*3, 1)).reshape(B, 3, 1, -1)
            tgt_latent, latent_for_edit, distractor_latent = referenced_part_latent.reshape(B, 3, 1, n_dim).split(1, dim=1)
            tgt_latent, latent_for_edit, distractor_latent = tgt_latent.reshape(B, n_dim), latent_for_edit.reshape(B, n_dim), distractor_latent.reshape(B, n_dim)
            modified_part_latent_res = self.update_mlp(torch.cat([language_enc_f, latent_for_edit], -1))
            modified_part_latent = latent_for_edit + modified_part_latent_res
            modified_part_latent = torch.stack([tgt_latent, modified_part_latent, distractor_latent], dim=1).reshape(B*3, 1, n_dim)
            total_modified_part_latent = (1 - part_indicators).reshape(B*3, n_part, 1) * part_latent + part_indicators.reshape(B*3, n_part, 1) * modified_part_latent
            global_feature, anchors, part_latent = self.decomposer(total_modified_part_latent)
            
            
            assert anchors.shape[1] == part_latent.shape[1] == self.num_anchors
            param_per_point = anchors.reshape(B*3, self.num_anchors, 1, 3).repeat([1, 1, self.points_per_anchor, 1]).reshape(B*3, self.npoints, 3).transpose(1,2)
            assert param_per_point.shape[1] == 3
            anchor_per_point = param_per_point

            latent_per_point = part_latent.reshape(B*3, self.num_anchors, 1, n_dim).repeat([1, 1, self.points_per_anchor, 1]).reshape(B*3, self.npoints, n_dim)
            pred = self.decode(anchor_per_point, code=global_feature, pointwise_latent=latent_per_point, device=device, variance=None) 
            pred_out = dict()
            for k, v in pred.items():
                target, edit, distractor = v.reshape(B, 3, self.npoints, 3).split(1, dim=1)
                pred_out[f"distractor_{k}"] = distractor.squeeze(1)
                pred_out[f"distractor_edit_{k}"] = edit.squeeze(1)
                pred_out[f"target_{k}"] = target.squeeze(1)
            

            pred_out['text'] = texts
            pred_out['input_distractor'] = distractor_pcds
            pred_out['input_target'] = target_pcds
            pred_out['anchors'] = anchors
            pred_out['seg_mask'] = distractor_attn_maps.max(-1)[1]
            pred_out['seg_mask_ref'] = target_attn_maps.max(-1)[1]
            pred_out['assigned_anchor'] = anchor_per_point.transpose(1, 2)
            pred_out['target_shift'] = data['target_shift']
            pred_out['target_scale'] = data['target_scale']
            pred_out['distractor_shift'] = data['distractor_shift']
            pred_out['distractor_scale'] = data['distractor_scale']
            return pred_out
        else:
            input = data['input'].to(device)
            ref = data['ref'].to(device)
            attn_map = data['attn_map'].to(device) # ->[B, npoints, n_class]
            seg_mask = data['seg_mask'].to(device)
            B = input.shape[0]
            x = torch.cat([input, attn_map], dim=-1).repeat_interleave(2, dim=0) # -> [B, npoints, n_class + 3]
            part_latent = self.encoder(x)
            input_part_latent, part_latent_for_edit = part_latent.reshape(B, 2, self.num_anchors, -1).split(1, dim=1)
            input_part_latent, part_latent_for_edit = input_part_latent.reshape(B, self.num_anchors, -1), part_latent_for_edit.reshape(B, self.num_anchors, -1)
            input_global_feature, input_anchor, input_part_latent = self.decomposer(input_part_latent)

            if self.sample_by_seg_mask:
                anchor_assignments = seg_mask.to(torch.int32)
                scale = int(self.npoints / anchor_assignments.shape[1])
                anchor_assignments = anchor_assignments.repeat_interleave(scale, dim=1)
                input_latent_per_point = gather_operation(input_part_latent.transpose(2,1).contiguous(), anchor_assignments).reshape(B, -1, self.npoints).transpose(1,2)
                input_anchor_per_point = gather_operation(input_anchor.transpose(2, 1).contiguous(), anchor_assignments).reshape(B, 3, self.npoints)
                assert input_anchor_per_point.shape[1] == 3
            else:
                input_anchor_per_point = input_anchor.reshape(B, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B, self.npoints, -1).transpose(1,2)
                assert input_anchor_per_point.shape[1] == 3
                input_latent_per_point = input_part_latent.reshape(B, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B, self.npoints, -1)
            
            pred = self.decode(input_anchor_per_point, code=input_global_feature, pointwise_latent=input_latent_per_point, device=device, variance=None) 
            pred['input'] = input
            pred['input_ref'] = ref
            pred['anchors'] = input_anchor
            pred['seg_mask'] = seg_mask 
            pred['seg_mask_ref'] = seg_mask 
            pred['assigned_anchor'] = input_anchor_per_point.transpose(1, 2)
            pred['shift'] = data['shift']
            pred['scale'] = data['scale']

            if self.latent_diffusion is not None:
                part_latent_for_edit = part_latent_for_edit.transpose(1,2)
                input_anchor = input_anchor.transpose(1,2)
            for i, prompt in enumerate(tokenized_prompts):
                print(prompts[i])
                prompt = prompt.to(device).unsqueeze(0)
                one_hot_id = torch.eye(self.num_anchors).to(device)[ids[i]].unsqueeze(0).repeat_interleave(B, dim=0)
                language_enc_f, language_enc_attn = self.language_encoder(prompt)
                language_enc_f = language_enc_f.repeat_interleave(B, dim=0)
                if self.latent_diffusion is not None:
                    
                    language_enc_f = language_enc_f.transpose(1,2)
                    # print(part_latent_for_edit.shape, language_enc_f.shape)
                    total_modified_part_latent = self.latent_diffusion.p_sample_loop(part_latent_for_edit.shape, part_latent_for_edit, language_enc_f, anchors=input_anchor)
                    # print(total_modified_part_latent.shape)
                    total_modified_part_latent = total_modified_part_latent.transpose(1, 2).unsqueeze(1)
                    K = total_modified_part_latent.shape[1]
                    # total_modified_part_latent = part_latent_for_edit + torch.rand_like(part_latent_for_edit)
                else:
                    referenced_part_latent = part_latent_for_edit[:, ids[i]]
                    if self.icmle:
                        K=10
                        conditional = torch.randn([B, K, self.conditional_dim]).cuda()
                    else:
                        K=1
                        conditional = None
                    total_modified_part_latent = self.latent_language_encoder(part_latent_for_edit, one_hot_id, language_enc_f, conditional=conditional)
                edited_global_feature, edited_anchors, edited_part_latent = self.decomposer(total_modified_part_latent.reshape(B*K, self.num_anchors, -1))
                
                
                if self.sample_by_seg_mask and not (ids[i] == 3):
                    anchor_assignments = seg_mask.to(torch.int32)
                    scale = int(self.npoints / anchor_assignments.shape[1])
                    anchor_assignments = anchor_assignments.repeat_interleave(scale, dim=1).repeat_interleave(K, dim=0)
                    edited_latent_per_point = gather_operation(edited_part_latent.reshape(B*K, self.num_anchors, -1).transpose(2,1).contiguous(), anchor_assignments).reshape(B*K, -1, self.npoints).transpose(1,2)
                    edited_anchor_per_point = gather_operation(edited_anchors.reshape(B*K, self.num_anchors, -1).transpose(2, 1).contiguous(), anchor_assignments).reshape(B*K, 3, self.npoints)
                    assert edited_anchor_per_point.shape[1] == 3
                else:
                    edited_anchor_per_point = edited_anchors.reshape(B*K, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B*K, self.npoints, -1).transpose(1,2)
                    assert edited_anchor_per_point.shape[1] == 3
                    edited_latent_per_point = edited_part_latent.reshape(B*K, self.num_anchors, 1, -1).repeat([1, 1, self.points_per_anchor, 1]).reshape(B*K, self.npoints, -1)
                

                pred_edit = self.decode(edited_anchor_per_point, code=edited_global_feature, pointwise_latent=edited_latent_per_point, device=device, variance=None)
                for k in range(K):
                    # for t in range(20, 200, 20):
                    #     pred[f'edit_pred: "{prompts[i]}" sample {k} T={t}'] = pred_edit[t].reshape(B, K, self.npoints, 3)[:, k]
                    pred[f'edit_pred: "{prompts[i]}" sample {k}'] = pred_edit['pred'].reshape(B, K, self.npoints, 3)[:, k]
        return pred
        
    def anchor_drift(self, pcds, x, global_feature, anchors, part_latent, seg_mask, device):
        out_dict = dict()
        B = global_feature.shape[0]
        direction = F.normalize(torch.randn(B, 10, self.num_anchors, 1).cuda())
        part_id = torch.tensor([1.,1., 1., 1.]).to(float).cuda().reshape(1, 1, -1, 1)
        scales = torch.linspace(0., 2, 10).to(float).cuda().reshape(1, -1, 1, 1)
        part_latent = part_latent.repeat_interleave(10, dim=0)
        global_feature = global_feature.repeat_interleave(10, dim=0)
        if self.project_latent:
            primary_part_latent = self.encoder(x)
            intrinsic, anchors = torch.split(self.down_proj(primary_part_latent), [self.part_dim, 3], dim=-1)
            intrinsic = intrinsic.repeat_interleave(10, dim=0)
        shifted_anchors = direction * scales  * part_id + anchors.unsqueeze(1) * part_id #+ anchors.unsqueeze(1)
        shifted_anchors = shifted_anchors.reshape(B*10, self.num_anchors, 3).to(torch.float32)
        if self.project_latent:
            primary_part_latent = intrinsic + self.up_proj(shifted_anchors)
            global_feature, _, part_latent = self.decomposer(primary_part_latent)
        if self.use_primary:
            part_latent = primary_part_latent 
        
        anchor_assignments = seg_mask.to(torch.int32).repeat_interleave(10, dim=0)
        scale = int(self.npoints / anchor_assignments.shape[1])
        anchor_assignments = anchor_assignments.repeat_interleave(scale, dim=1)
        latent_per_point = gather_operation(part_latent.transpose(2,1).contiguous(), anchor_assignments).reshape(B*10, -1, self.npoints).transpose(1,2)
        anchor_per_point = gather_operation(shifted_anchors.transpose(2, 1).contiguous(), anchor_assignments).reshape(B*10, 3, self.npoints)
        print(latent_per_point.shape, anchor_per_point.shape)
        assert anchor_per_point.shape[1] == 3
        pred = self.decode(anchor_per_point, code=global_feature, pointwise_latent=latent_per_point, device=device, variance=None) 
        for k in range(10):
            # for t in range(20, 200, 20):
            #     pred[f'edit_pred: "{prompts[i]}" sample {k} T={t}'] = pred_edit[t].reshape(B, K, self.npoints, 3)[:, k]
            out_dict[f'anchor sample {k}'] = shifted_anchors.reshape(B, 10, self.num_anchors, 3)[:, k]
            out_dict[f'pred anchor drift sample {k}'] = pred['pred'].reshape(B, 10, self.npoints, 3)[:, k]
        out_dict['input'] = pcds['input']
        out_dict['pred'] = pcds['input']
        out_dict['input_ref'] = pcds['input']
        
        out_dict['seg_mask'] = seg_mask 
        out_dict['seg_mask_ref'] = seg_mask 
        # out_dict['assigned_anchor'] = anchor_per_point.transpose(1, 2)
        out_dict['shift'] = pcds['shift']
        out_dict['scale'] = pcds['scale']
        return out_dict
        
    def forward(self, pcds, device='cuda', iter=-1):
        """
            input: [B, npoints, 3]
            attn_map: [B, npoints, n_class]
            seg_mask: [B, npoints]
        """

        input = pcds['input'].to(device)
        ref = pcds['ref'].to(device)
        attn_map = pcds['attn_map'].to(device) # ->[B, npoints, n_class]
        seg_mask = pcds['seg_mask'].to(device)
        loss_dict = dict()

        npoint = input.shape[-2]
        
        if self.include_attn_weight_in_encoder:
            x = torch.cat([input, attn_map], dim=-1) # -> [B, npoints, n_class + 3]
        else:
            x = input

        B = input.shape[0]

        z = self.encoder(x).reshape(B, -1)
                
        if self.training:
            
            t, _ = self.sampler.sample(B, device)



            diffusion_loss = self.diffusion.training_losses(ref.transpose(1,2), t, anchors=torch.zeros_like(ref.transpose(1,2)), code=z, pointwise_latent=None, variance=None, reduce=True)
            loss_dict['diffusion_loss'] = diffusion_loss
            return loss_dict

        else:
            
            pred = self.decode(torch.zeros_like(ref.transpose(1,2)), code=z, pointwise_latent=None, device=device, variance=None) 
            pred['input'] = input
            pred['input_ref'] = ref
            pred['seg_mask'] = seg_mask 
            pred['seg_mask_ref'] = seg_mask 
            pred['shift'] = pcds['shift']
            pred['scale'] = pcds['scale']
            pred = {k:v.detach().cpu() for k, v in pred.items()}
            return pred

python/difffacto/models/networks/point_diff.py
diffFacto/diffFacto
from .pointnet_segmentor import PointNetDenseCls
python/difffacto/models/segmentors/__init__.py
diffFacto/diffFacto
import torch
from torch import nn
import torch.nn.functional as F
from difffacto.utils.registry import ENCODERS, build_from_cfg, SEGMENTORS

@SEGMENTORS.register_module()
class PointNetDenseCls(nn.Module):
    def __init__(self, n_class = 2):
        super(PointNetDenseCls, self).__init__()
        self.k = n_class
        self.feat = build_from_cfg(dict(type="PointNet", global_feat=False, latent_dim=1024), ENCODERS)
        self.conv1 = torch.nn.Conv1d(1088, 512, 1)
        self.conv2 = torch.nn.Conv1d(512, 256, 1)
        self.conv3 = torch.nn.Conv1d(256, 128, 1)
        self.conv4 = torch.nn.Conv1d(128, self.k, 1)
        self.bn1 = nn.BatchNorm1d(512)
        self.bn2 = nn.BatchNorm1d(256)
        self.bn3 = nn.BatchNorm1d(128)

    def forward(self, x):
        batchsize = x.size()[0]
        n_pts = x.size()[1]
        x  = self.feat(x.transpose(1,2))
        global_feat = x
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = self.conv4(x)
        x = x.transpose(2,1).contiguous()
        x = F.log_softmax(x.view(-1,self.k), dim=-1)
        x = x.view(batchsize, n_pts, self.k)
        return x, global_feat
python/difffacto/models/segmentors/pointnet_segmentor.py
diffFacto/diffFacto
from .optimizers import Adam, AdamW
from .schedulers import LinearLR, StepLR, CosineAnnealingLRWithWarmUp
python/difffacto/optimizers/__init__.py
diffFacto/diffFacto
import torch
from difffacto.utils.registry import OPTIMS

@OPTIMS.register_module()
class Adam(torch.optim.Adam):
    def __init__(self, params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0):
        super().__init__(params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
        
@OPTIMS.register_module()
class AdamW(torch.optim.AdamW):
    def __init__(self, params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0):
        super().__init__(params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
        
@OPTIMS.register_module()
class Adamax(torch.optim.Adamax):
    def __init__(self, params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0):
        super().__init__(params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
python/difffacto/optimizers/optimizers.py
diffFacto/diffFacto
from torch.optim.lr_scheduler import LambdaLR, StepLR, CosineAnnealingLR, _LRScheduler
from difffacto.utils.registry import SCHEDULERS
from bisect import bisect
import warnings
import math

@SCHEDULERS.register_module()
def LinearLR(optimizer, start_epoch, end_epoch, start_lr, end_lr):
    def lr_func(epoch):
        if epoch <= start_epoch:
            return 1.0
        elif epoch <= end_epoch:
            total = end_epoch - start_epoch
            delta = epoch - start_epoch
            frac = delta / total
            return (1-frac) * 1.0 + frac * (end_lr / start_lr)
        else:
            return end_lr / start_lr
    return LambdaLR(optimizer, lr_lambda=lr_func)

@SCHEDULERS.register_module()
def CustomStepLR(optimizer, milestone, gamma):
    milestone.sort()
    def lr_func(epoch):
        idx = bisect(milestone, epoch)
        return gamma ** idx

    return LambdaLR(optimizer, lr_lambda=lr_func)

@SCHEDULERS.register_module()
class CosineAnnealingLRWithWarmUp(_LRScheduler):
    r"""Set the learning rate of each parameter group using a cosine annealing
    schedule, where :math:`\eta_{max}` is set to the initial lr and
    :math:`T_{cur}` is the number of epochs since the last restart in SGDR:

    .. math::
        \begin{aligned}
            \eta_t & = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})\left(1
            + \cos\left(\frac{T_{cur}}{T_{max}}\pi\right)\right),
            & T_{cur} \neq (2k+1)T_{max}; \\
            \eta_{t+1} & = \eta_{t} + \frac{1}{2}(\eta_{max} - \eta_{min})
            \left(1 - \cos\left(\frac{1}{T_{max}}\pi\right)\right),
            & T_{cur} = (2k+1)T_{max}.
        \end{aligned}

    When last_epoch=-1, sets initial lr as lr. Notice that because the schedule
    is defined recursively, the learning rate can be simultaneously modified
    outside this scheduler by other operators. If the learning rate is set
    solely by this scheduler, the learning rate at each step becomes:

    .. math::
        \eta_t = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})\left(1 +
        \cos\left(\frac{T_{cur}}{T_{max}}\pi\right)\right)

    It has been proposed in
    `SGDR: Stochastic Gradient Descent with Warm Restarts`_. Note that this only
    implements the cosine annealing part of SGDR, and not the restarts.

    Args:
        optimizer (Optimizer): Wrapped optimizer.
        T_max (int): Maximum number of iterations.
        eta_min (float): Minimum learning rate. Default: 0.
        last_epoch (int): The index of last epoch. Default: -1.
        verbose (bool): If ``True``, prints a message to stdout for
            each update. Default: ``False``.

    .. _SGDR\: Stochastic Gradient Descent with Warm Restarts:
        https://arxiv.org/abs/1608.03983
    """

    def __init__(self, optimizer, T_max, eta_min=0, warmup_T=10, warmup_lr=1e-6, last_epoch=-1, verbose=False):
        self.T_max = T_max - warmup_T
        self.eta_min = eta_min
        self.warmup_T=warmup_T
        self.warmup_lr=warmup_lr
        super(CosineAnnealingLRWithWarmUp, self).__init__(optimizer, last_epoch, verbose)

    def get_lr(self):
        if not self._get_lr_called_within_step:
            warnings.warn("To get the last learning rate computed by the scheduler, "
                          "please use `get_last_lr()`.", UserWarning)
        if self.last_epoch == 0:
            return [self.warmup_lr for group in self.optimizer.param_groups]
        elif self.last_epoch < self.warmup_T:
            target_lrs = [self.eta_min + (base_lr - self.eta_min) *
                    (1 + math.cos(self.warmup_T * math.pi / self.T_max)) / 2 for base_lr in self.base_lrs]
            return [self.warmup_lr + (base_lr - self.warmup_lr) * (self.last_epoch / self.warmup_T)
                for base_lr, group in zip(target_lrs, self.optimizer.param_groups)]
        elif self._step_count == 1 and self.last_epoch > 0:
            return [self.eta_min + (base_lr - self.eta_min) *
                    (1 + math.cos((self.last_epoch) * math.pi / self.T_max)) / 2
                    for base_lr, group in
                    zip(self.base_lrs, self.optimizer.param_groups)]
        elif (self.last_epoch - 1 - self.T_max) % (2 * self.T_max) == 0:
            return [group['lr'] + (base_lr - self.eta_min) *
                    (1 - math.cos(math.pi / self.T_max)) / 2
                    for base_lr, group in
                    zip(self.base_lrs, self.optimizer.param_groups)]
        return [(1 + math.cos(math.pi * self.last_epoch / self.T_max)) /
                (1 + math.cos(math.pi * (self.last_epoch - 1) / self.T_max)) *
                (group['lr'] - self.eta_min) + self.eta_min
                for group in self.optimizer.param_groups]

    def _get_closed_form_lr(self):
        return [self.eta_min + (base_lr - self.eta_min) *
                (1 + math.cos(math.pi * self.last_epoch / self.T_max)) / 2
                for base_lr in self.base_lrs]
        
        
SCHEDULERS.register_module('StepLR', StepLR)

SCHEDULERS.register_module('CosineAnnealingLR', CosineAnnealingLR)

python/difffacto/optimizers/schedulers.py
diffFacto/diffFacto
from .runner import Runner
from .gan_runner import GanRunner
from .pretrain_runner import PretrainRunner
from .language_runner import LanguageRunner
python/difffacto/runner/__init__.py
diffFacto/diffFacto
import torch 
from torch import nn
import numpy as np
from difffacto.utils.misc import fps
from difffacto.datasets.evaluation_utils import compute_all_metrics
from difffacto.utils.gan_losses import gen_loss, dis_loss, gradient_penalty
from tqdm import tqdm
from difffacto.config.config import get_cfg, save_cfg, save_args
from difffacto.utils.registry import build_from_cfg,MODELS,SAMPLERS,DATASETS,OPTIMS,SCHEDULERS,HOOKS, GENERATORS,DISCRIMINATORS 
from difffacto.utils.misc import check_file, build_file, search_ckpt, check_interval, sync, parse_losses
import pickle
import time
import datetime
from difffacto.utils import misc
from difffacto.utils import dist_utils
from einops import rearrange
class GanRunner:
    def __init__(self, device, args):
        cfg = get_cfg()

        if args.local_rank == 0:
            self.logger = build_from_cfg(cfg.logger,HOOKS,work_dir=cfg.work_dir)
        else:
            self.logger = None
        self.distributed = args.distributed
        self.no_eval=args.no_eval
        self.local_rank = args.local_rank
        self.world_size = args.world_size
        self.short_val = args.short_val
        self.eval_both = cfg.eval_both
        self.eval_whole= cfg.eval_whole
        self.prefix=args.prefix

        if args.seed is not None:
            if self.logger is not None:
                self.logger.print_log(f'Set random seed to {args.seed}, 'f'deterministic: {args.deterministic}')
            misc.set_random_seed(args.seed + args.local_rank, deterministic=args.deterministic) # seed + rank, for augmentation
        if args.distributed:
            assert args.local_rank == torch.distributed.get_rank() 
        self.cfg = cfg
        self.work_dir = cfg.work_dir
        self.save_num_batch = cfg.save_num_batch

        self.max_epoch = cfg.max_epoch 
        self.max_iter = cfg.max_iter
        assert (self.max_iter is None)^(self.max_epoch is None),"You must set max_iter or max_epoch"

        self.checkpoint_interval = cfg.checkpoint_interval
        self.eval_interval = cfg.eval_interval
        self.log_interval = cfg.log_interval
        self.resume_path = cfg.resume_path
        self.ret_traj=self.cfg.ret_traj
        self.ret_interval=self.cfg.ret_interval
        self.max_norm = cfg.max_norm
        self.disc=build_from_cfg(cfg.disc, DISCRIMINATORS)
        self.model = build_from_cfg(cfg.model,MODELS)
        if device=='cuda':
            self.model.to(self.local_rank)
            self.disc.to(self.local_rank)

        if args.distributed:
            # Sync BN
            if args.sync_bn:
                self.model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.model)
                self.disc = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.disc)
                if self.logger is not None:
                    self.logger.print_log('Using Synchronized BatchNorm ...')
            self.model = nn.parallel.DistributedDataParallel(self.model, device_ids=[args.local_rank % torch.cuda.device_count()], find_unused_parameters=True)
            self.disc = nn.parallel.DistributedDataParallel(self.disc, device_ids=[args.local_rank % torch.cuda.device_count()], find_unused_parameters=True)
            if self.logger is not None:
                self.logger.print_log('Using Distributed Data parallel ...')
        else:
            if self.logger is not None:
                self.logger.print_log('Using Data parallel ...')
            self.model = nn.DataParallel(self.model).cuda()
            self.disc = nn.DataParallel(self.disc).cuda()
        
        # optimizer & scheduler
        self.opt_gen = build_from_cfg(cfg.gen_optimizer,OPTIMS,params=self.model.module.encoder.part_aligner.parameters())
        self.gen_scheduler = build_from_cfg(cfg.gen_scheduler, SCHEDULERS, optimizer=self.opt_gen)
        
        self.opt_dis = build_from_cfg(cfg.disc_optimizer,OPTIMS,params=self.disc.parameters())
        self.disc_scheduler = build_from_cfg(cfg.disc_scheduler, SCHEDULERS, optimizer=self.opt_dis)

        if self.distributed:
            assert cfg.dataset.train.batch_size % self.world_size == 0
            cfg.dataset.train.batch_size = cfg.dataset.train.batch_size // self.world_size
        self.train_dataset, self.train_sampler = build_from_cfg(cfg.dataset.train,DATASETS, distributed=args.distributed)
        self.val_dataset, _ = build_from_cfg(cfg.dataset.val,DATASETS, distributed=False)
        self.n_class = cfg.n_class
        # only save file in the 0th process
        if  self.local_rank == 0:
            save_file = build_file(self.work_dir,prefix="config.yaml")
            save_arg = build_file(self.work_dir,prefix="args.yaml")
            save_args(save_arg, args)
            save_cfg(save_file)


        self.iter = 0
        self.total_gan_iters = 0
        self.n_critics = cfg.n_critics
        self.epoch = 0
        if self.max_epoch:
            if (self.train_dataset):
                self.total_iter = self.max_epoch * len(self.train_dataset)
            else:
                self.total_iter = 0
        else:
            self.total_iter = self.max_iter

        if (cfg.pretrained_weights):
            self.load(cfg.pretrained_weights, model_only=cfg.model_only)
        
        if self.resume_path is None:
            self.resume_path = search_ckpt(self.work_dir)
        if check_file(self.resume_path):
            self.resume()


    @property
    def finish(self):
        if self.max_epoch:
            return self.epoch>=self.max_epoch
        else:
            return self.iter>=self.max_iter
    
    def run(self):
        if self.logger is not None:
            self.logger.print_log("Start running")
        
        while not self.finish:
            if self.distributed:
                self.train_sampler.set_epoch(self.epoch)
            self.train()
            if check_interval(self.epoch,self.eval_interval) and self.local_rank == 0:
                self.val()
            if check_interval(self.epoch,self.checkpoint_interval) and self.local_rank == 0:
                self.save()
    
    
    def _update_gan_(self, data, gen=False, device='cuda'):
        self.model.eval()
        self.model.module.encoder.part_aligner.train()
        self.disc.train()
        self.opt_gen.zero_grad()
        self.opt_dis.zero_grad()


        input = data['input'].to(device)
        valid_id = data['present'].to(device)
        ref = data['ref'].to(device).transpose(1,2)
        seg_flag = data['ref_attn_map'].to(device)
        gt_shift = data.get('part_shift', torch.zeros(ref.shape[0], 3, self.n_class)).to(device)
        gt_var = data.get('part_scale', torch.ones(ref.shape[0], 3, self.n_class)).to(device) ** 2
        batch_size, C, N = ref.shape
        part_code, _ = self.model.module.encoder.get_part_code(input, seg_flag)
        part_code = part_code.transpose(1,2)
        n_class=part_code.shape[-1]
        noise = torch.randn(batch_size, self.model.module.encoder.part_aligner.noise_dim).cuda()
        fake_mean, fake_logvar = self.model.module.encoder.get_params_from_part_code(part_code, valid_id, gt_mean=gt_shift, gt_var=gt_var, ref=ref, noise=noise)
        x_fake = torch.cat([fake_mean, fake_logvar], dim=1)
        x_real = torch.cat([gt_shift, torch.log(gt_var)], dim=1)
        x_real.requires_grad=True
        x_inp = torch.cat([x_real, x_fake], dim=0)
        part_code = part_code.unsqueeze(0).repeat_interleave(2, dim=0).reshape(batch_size*2, -1, n_class)
        valid_id = valid_id.unsqueeze(0).repeat_interleave(2, dim=0).reshape(batch_size*2,  n_class)
        d_out = self.disc(x_inp, part_code, mask=valid_id)
        d_real = d_out[:batch_size, ...]
        d_fake = d_out[batch_size:, ...]
        # print('dfake mean', d_fake.mean())

        loss, loss_res = None, {}
        loss_type = "wgan"
        if gen:
            gen_loss_weight = self.cfg.gen_loss_weight
            loss_gen, gen_loss_res = gen_loss(
                d_real, d_fake, weight=gen_loss_weight, loss_type=loss_type)
            loss = loss_gen + (0. if loss is None else loss)
            loss.backward()
            # print("gen loss", loss)
            self.opt_gen.step()
            loss_res.update({
                ("train/gan_pass/gen/%s" % k): v
                for k, v in gen_loss_res.items()
            })
            loss_res['loss'] = loss.detach().cpu().item()
        else:
            # Get gradient penalty
            gp_weight = self.cfg.gp_weight
            if gp_weight > 0:
                gp_type = self.cfg.gp_type
                gp, gp_res = gradient_penalty(
                    x_real, x_fake, d_real, d_fake,
                    weight=gp_weight, gp_type=gp_type)
                loss = gp + (0. if loss is None else loss)
                # print("gp loss", gp)
                loss_res.update({
                    ("train/gan_pass/gp_loss/%s" % k): v
                    for k, v in gp_res.items()
                })

            dis_loss_weight = self.cfg.dis_loss_weight
            loss_dis, dis_loss_res = dis_loss(
                d_real, d_fake, weight=dis_loss_weight, loss_type=loss_type)
            loss = loss_dis + (0. if loss is None else loss)
            # print("dis loss", loss_dis)
            loss.backward()
            self.opt_dis.step()
            loss_res.update({
                ("train/gan_pass/dis/%s" % k): v for k, v in dis_loss_res.items()
            })
            loss_res['loss'] = loss.detach().cpu().item()
        # print("xfake mean and var", x_fake.mean(), x_fake.var())


        return loss_res
    
    def train(self):
        # torch.autograd.set_detect_anomaly(True)
        start_time = time.time()

        for batch_idx, pcds in enumerate(self.train_dataset):
            self.total_gan_iters += 1
            res = {}
            if self.total_gan_iters % self.n_critics == 0:
                gen_res = self._update_gan_(pcds, gen=True)
                res.update(gen_res)
            dis_res = self._update_gan_(pcds, gen=False)
            res.update(dis_res)
            batch_size = self.train_dataset.batch_size
            
            if check_interval(self.iter,self.log_interval) and self.iter>0:
                batch_size = self.train_dataset.batch_size
                ptime = time.time()-start_time
                eta_time = (self.total_iter-self.iter)*ptime/(batch_idx+1)
                eta_str = str(datetime.timedelta(seconds=int(eta_time)))
                data = dict(
                    name = self.cfg.name,
                    gen_lr = self.opt_gen.param_groups[0]['lr'] ,
                    dis_lr = self.opt_dis.param_groups[0]['lr'] ,
                    iter = self.iter,
                    epoch = self.epoch,
                    batch_idx = batch_idx,
                    batch_size = batch_size,
                    total_loss = dist_utils.reduce_tensor(res['loss'], self.world_size) if self.distributed else res['loss'],
                    eta=eta_str
                )
                data.update(dist_utils.reduce_tensor(res, self.world_size) if self.distributed else res)
                data = sync(data)
                if self.logger is not None:
                    self.logger.log(data)
            
            self.iter+=1
            if self.finish:
                break

        self.epoch +=1
        if self.max_epoch:
            if self.gen_scheduler is not None:
                self.gen_scheduler.step(epoch=self.epoch)
            if self.disc_scheduler is not None:
                self.disc_scheduler.step(epoch=self.epoch)


    @torch.no_grad()
    def val(self):
        if self.val_dataset is None:
            if self.logger is not None:
                self.logger.print_log("Please set Val dataset")
        else:
            if self.logger is not None:
                self.logger.print_log("Validating....")
            # TODO: need move eval into this function
            self.model.eval()
            results = dict()
            for batch_idx, pcds in tqdm(enumerate(self.val_dataset),total=len(self.val_dataset)):
                if self.short_val and batch_idx > 0: break
                _result = self.model(pcds, device=self.local_rank)
                for result, name in _result:
                    if name not in results.keys():
                        results[name] = []
                    results[name].append(sync(result, to_numpy=False))
            if self.no_eval:
                return
            for k, v in results.items():
                save_dict, eval_result = self.val_dataset.evaluate(v, self.save_num_batch, self.local_rank)
                save_file = build_file(self.work_dir, prefix=f"val/{k}_{self.prefix}_{self.epoch}.pkl")
                self.logger.print_log(f"Saving to {save_file}....")
                pickle.dump(save_dict, open(save_file, 'wb'))
                eval_result = {k:v for k, v in eval_result.items()}
                self.logger.log(eval_result,iter=self.iter)
            if self.eval_both and not self.short_val:
                results = dict()
                for batch_idx, pcds in tqdm(enumerate(self.train_dataset),total=len(self.train_dataset)):
                    if self.short_val and batch_idx > 0: break
                    _result = self.model(pcds, device=self.local_rank)
                    for result, name in _result:
                        if name not in results.keys():
                            results[name] = []
                        results[name].append(sync(result, to_numpy=False))
                if self.no_eval:
                    return
                for k, v in results.items():
                    save_dict, eval_result = self.val_dataset.evaluate(v, self.save_num_batch, self.local_rank)
                    save_file = build_file(self.work_dir, prefix=f"val/train_set_{k}_{self.prefix}_{self.epoch}.pkl")
                    self.logger.print_log(f"Saving to {save_file}....")
                    pickle.dump(save_dict, open(save_file, 'wb'))
                    eval_result = {f"train_{k}":v for k, v in eval_result.items()}
                    self.logger.log(eval_result,iter=self.iter)
                    
    def generate_samples(self, num_gen, param_sample_num):
        if self.val_dataset is None:
            if self.logger is not None:
                self.logger.print_log("Please set Val dataset")
            return
        if self.logger is not None:
            self.logger.print_log("Generating....")
        # TODO: need move eval into this function
        self.model.eval()
        results = dict(pred=[], ref=[],  seg_mask_ref=[])
        fixed_id = [0,0,0,0]
        B = self.val_dataset.batch_size
        from difffacto.datasets.dataset_utils import shapenet_chair_part_distribution
        valid_ids = []
        probs = []
        for k, v in shapenet_chair_part_distribution.items():
            valid_ids.append(torch.tensor(list(map(float, list(k)))))
            probs.append(v)
        probs = torch.tensor(probs)
        valid_ids = torch.stack(valid_ids, dim=0).cuda()
        K=param_sample_num
        for batch_idx in tqdm(range(num_gen // self.val_dataset.batch_size)):
            idx = torch.multinomial(probs, self.val_dataset.batch_size, replacement=True)
            valid_id = valid_ids[idx]
            ctx, mean_per_point, logvar_per_point, _seg_mask, _valid_id = self.model.module.sample(B, fixed_id, valid_id, 'cuda', self.epoch, K=K)
            variance_per_point = torch.exp(logvar_per_point)
            C, N = variance_per_point.shape[1:]
            pred =self.model.module.decode(mean_per_point, ctx=ctx, device=self.local_rank, variance=variance_per_point, anchor_assignments=_seg_mask.to(torch.int32), valid_id=_valid_id) 
            results['pred'].append(pred['pred'].cpu().detach().numpy())
            results['seg_mask_ref'].append(_seg_mask.cpu().detach().numpy())
        for batch_idx, pcds in tqdm(enumerate(self.val_dataset),total=len(self.val_dataset)):
            results['ref'].append(pcds['ref'].cpu().detach().numpy())
        results = {k:np.concatenate(v, axis=0) for k, v in results.items()}
        save_file = build_file(self.work_dir, prefix=f"val/gen_fixed{''.join(map(str,fixed_id))}_{self.prefix}_{self.epoch}.pkl")
        self.logger.print_log(f"Saving to {save_file}....")
        pickle.dump(results, open(save_file, 'wb'))
        self.evaluate_gen(results['pred'], results['ref'], self.val_dataset.batch_size)
    
    def evaluate_gen(self, pred, ref, batch_size):
        preds = []
        refs = []
        print(pred.shape, ref.shape)
        for i in tqdm(range(pred.shape[0])):
            _pred = pred[i]
            if _pred.shape[0] > 2048:
                _pred = fps(torch.from_numpy(_pred).cuda().unsqueeze(0).contiguous(), 2048).detach().cpu().numpy()[0]
            pred_max = _pred.max(0).reshape(1, 3) # (1, 3)
            pred_min = _pred.min(0).reshape(1, 3) # (1, 3)
            pred_shift = ((pred_min + pred_max) / 2).reshape(1, 3)
            pred_scale = (pred_max - pred_min).max(-1)[0].reshape(1, 1) / 2
            _pred = (_pred - pred_shift) / pred_scale
            preds.append(_pred)
        for i in tqdm(range(ref.shape[0])):
            _ref = ref[i]
            if _ref.shape[0] > 2048:
                _ref = fps(torch.from_numpy(_ref).cuda().unsqueeze(0).contiguous(), 2048).detach().cpu().numpy()[0]
            ref_max = _ref.max(0).reshape(1, 3) # (1, 3)
            ref_min = _ref.min(0).reshape(1, 3) # (1, 3)
            ref_shift = ((ref_min + ref_max) / 2).reshape(1, 3)
            ref_scale = (ref_max - ref_min).max(-1)[0].reshape(1, 1) / 2
            _ref = (_ref - ref_shift) / ref_scale
            refs.append(_ref)

                        
        preds = torch.from_numpy(np.stack(preds, axis=0)).cuda()
        refs = torch.from_numpy(np.stack(refs, axis=0)).cuda()
        print(preds.shape, refs.shape)
        metrics = compute_all_metrics(preds, refs, batch_size, one_way=False)
        print(metrics)
        self.logger.log(metrics,iter=self.iter)
        
    def save(self):
        save_data = {
            "meta":{
                "epoch": self.epoch,
                "iter": self.iter,
                "max_iter": self.max_iter,
                "max_epoch": self.max_epoch,
                "config": self.cfg.dump()
            },
            "model":self.model.module.state_dict() if self.distributed else self.model.state_dict(),
            "discrimator":self.disc.module.state_dict() if self.distributed else self.disc.state_dict(),
        }
        if self.model.module.encoder.part_aligner is not None:
            save_data['param_sampler'] = self.model.module.encoder.part_aligner.state_dict()
        if self.opt_gen is not None:
            save_data['opt_gen'] = self.opt_gen.state_dict()
        if self.gen_scheduler is not None:
            save_data['gen_scheduler'] = self.gen_scheduler.state_dict()
        if self.opt_dis is not None:
            save_data['opt_dis'] = self.opt_dis.state_dict()
        if self.disc_scheduler is not None:
            save_data['disc_scheduler'] = self.disc_scheduler.state_dict()
        save_file = build_file(self.work_dir,prefix=f"checkpoints/ckpt_{self.epoch}.pth")
        torch.save(save_data,save_file)
        
    
    def load(self, load_path, model_only=False):
        map_location = {'cuda:%d' % 0: 'cuda:%d' % self.local_rank}
        resume_data = torch.load(load_path, map_location=map_location)

        if (not model_only):
            meta = resume_data.get("meta",dict())
            self.epoch = meta.get("epoch",self.epoch)
            self.iter = meta.get("iter",self.iter)
            self.max_iter = meta.get("max_iter",self.max_iter)
            self.max_epoch = meta.get("max_epoch",self.max_epoch)
            self.scheduler.load_state_dict(resume_data.get("scheduler",dict()))
            # self.optimizer.load_state_dict(resume_data.get("optimizer",dict()))
        #base_ckpt = {k.replace("module.", ""): v for k, v in resume_data['model'].items()}
        state_dict = resume_data['model']
        model_state_dict = self.model.state_dict()
        # print(model_state_dict.keys())
        for k in state_dict.keys():
            if k in model_state_dict.keys():
                if state_dict[k].shape != model_state_dict[k].shape:
                    self.logger.print_log(f"Skip loading parameter: {k}, "
                                f"required shape: {model_state_dict[k].shape}, "
                                f"loaded shape: {state_dict[k].shape}")
                    state_dict[k] = model_state_dict[k]
            else:
                self.logger.print_log(f"Dropping parameter {k}")
        for k in model_state_dict.keys():
            if k not in state_dict.keys():
                print(f"Layer {k} not loaded!")
        self.model.load_state_dict(state_dict,strict=False)

        self.logger.print_log(f"Loading model parameters from {load_path}")

    def resume(self):
        self.load(self.resume_path, model_only=self.cfg.model_only)

python/difffacto/runner/gan_runner.py
diffFacto/diffFacto
import torch 
from torch import nn
import numpy as np
from difffacto.utils.misc import fps
from  difffacto.datasets.evaluation_utils import compute_all_metrics
from tqdm import tqdm
from difffacto.config.config import get_cfg, save_cfg, save_args
from difffacto.utils.registry import build_from_cfg,MODELS,SAMPLERS,DATASETS,OPTIMS,SCHEDULERS,HOOKS
from difffacto.utils.misc import check_file, build_file, search_ckpt, check_interval, sync, parse_losses
import pickle
import os
import time
import datetime
from difffacto.utils import misc
from difffacto.utils import dist_utils
from einops import rearrange
class LanguageRunner:
    def __init__(self, device, args):
        cfg = get_cfg()

        if args.local_rank == 0:
            self.logger = build_from_cfg(cfg.logger,HOOKS,work_dir=cfg.work_dir)
        else:
            self.logger = None
            
        print(cfg.work_dir)
        self.distributed = args.distributed
        self.no_eval=args.no_eval
        self.local_rank = args.local_rank
        self.world_size = args.world_size
        self.short_val = args.short_val
        self.eval_both = cfg.eval_both
        self.eval_whole= cfg.eval_whole
        self.prefix=args.prefix

        if args.seed is not None:
            if self.logger is not None:
                self.logger.print_log(f'Set random seed to {args.seed}, 'f'deterministic: {args.deterministic}')
            misc.set_random_seed(args.seed + args.local_rank, deterministic=args.deterministic) # seed + rank, for augmentation
        if args.distributed:
            assert args.local_rank == torch.distributed.get_rank() 
        self.cfg = cfg
        self.work_dir = cfg.work_dir
        self.save_num_batch = cfg.save_num_batch

        self.max_epoch = cfg.max_epoch 
        self.max_iter = cfg.max_iter
        assert (self.max_iter is None)^(self.max_epoch is None),"You must set max_iter or max_epoch"

        self.checkpoint_interval = cfg.checkpoint_interval
        self.eval_interval = cfg.eval_interval
        self.log_interval = cfg.log_interval
        self.resume_path = cfg.resume_path
        self.ret_traj=self.cfg.ret_traj
        self.ret_interval=self.cfg.ret_interval
        self.max_norm = cfg.max_norm
        self.model = build_from_cfg(cfg.model,MODELS)
        if device=='cuda':
            self.model.to(self.local_rank)

        if args.distributed:
            # Sync BN
            if args.sync_bn:
                self.model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.model)
                if self.logger is not None:
                    self.logger.print_log('Using Synchronized BatchNorm ...')
            self.model = nn.parallel.DistributedDataParallel(self.model, device_ids=[args.local_rank % torch.cuda.device_count()], find_unused_parameters=True)
            if self.logger is not None:
                self.logger.print_log('Using Distributed Data parallel ...')
        else:
            if self.logger is not None:
                self.logger.print_log('Using Data parallel ...')
            self.model = nn.DataParallel(self.model).cuda()
        
        # optimizer & scheduler
        parameters = list(self.model.module.language_encoder.parameters()) + list(self.model.module.update_mlp.parameters()) + [self.model.module.temp]
        self.optimizer = build_from_cfg(cfg.optimizer,OPTIMS,params=parameters)
        self.scheduler = build_from_cfg(cfg.scheduler, SCHEDULERS, optimizer=self.optimizer)

        if self.distributed:
            assert cfg.dataset.train.batch_size % self.world_size == 0
            cfg.dataset.train.batch_size = cfg.dataset.train.batch_size // self.world_size
        self.train_dataset, self.train_sampler = build_from_cfg(cfg.dataset.train,DATASETS, distributed=args.distributed)
        self.val_dataset, _ = build_from_cfg(cfg.dataset.val,DATASETS, distributed=False)
        self.cimle = cfg.cimle
        self.cimle_start_epoch = cfg.cimle_start_epoch if cfg.cimle_start_epoch is not None else 0
        self.cache_interval = cfg.cimle_cache_interval
        
        # only save file in the 0th process
        if  self.local_rank == 0:
            save_file = build_file(self.work_dir,prefix="config.yaml")
            save_arg = build_file(self.work_dir,prefix="args.yaml")
            save_args(save_arg, args)
            save_cfg(save_file)


        self.iter = 0
        self.epoch = 0

        if self.max_epoch:
            if (self.train_dataset):
                self.total_iter = self.max_epoch * len(self.train_dataset)
            else:
                self.total_iter = 0
        else:
            self.total_iter = self.max_iter

        if (cfg.pretrained_weights):
            self.load(cfg.pretrained_weights, model_only=cfg.model_only)
        
        if self.resume_path is None:
            self.resume_path = search_ckpt(self.work_dir)
        if check_file(self.resume_path):
            self.resume()


    @property
    def finish(self):
        if self.max_epoch:
            return self.epoch>=self.max_epoch
        else:
            return self.iter>=self.max_iter
    
    def run(self):
        if self.logger is not None:
            self.logger.print_log("Start running")
        
        while not self.finish:
            if self.distributed:
                self.train_sampler.set_epoch(self.epoch)
            # if self.cimle and self.epoch >= self.cimle_start_epoch and check_interval(self.epoch - self.cimle_start_epoch, self.cache_interval) and self.local_rank == 0:
            #     self.cache_noise()
            self.train()
            if check_interval(self.epoch,self.eval_interval) and self.local_rank == 0:
                self.val()
            if check_interval(self.epoch,self.checkpoint_interval) and self.local_rank == 0:
                self.save()
    
    @torch.no_grad()
    def cache_noise(self):
        print("########## CACHING NOISE FOR CIMLE ##########")
        self.model.eval()
        all_noises = []
        all_ids = []
        for batch_idx, pcds in tqdm(enumerate(self.train_dataset), total=len(self.train_dataset)):
            noises = self.model.module.cache_noise(pcds, self.local_rank)
            all_noises.append(noises)
            all_ids.append(pcds['id'])
        self.train_dataset.dataset.store_noise(all_noises, all_ids)
        
    @torch.no_grad()
    def cache_noise(self):
        print("########## CACHING NOISE FOR CIMLE ##########")
        self.model.eval()
        all_noises = []
        all_ids = []
        for batch_idx, pcds in tqdm(enumerate(self.train_dataset), total=len(self.train_dataset)):
            noises = self.model.module.cache_noise(pcds, self.local_rank, eval_whole=self.eval_whole)
            all_noises.append(noises)
            all_ids.append(pcds['id'])
        self.train_dataset.dataset.store_noise(all_noises, all_ids)
        
        print("########## NOISE FOR CIMLE CACHED ##########")
    @torch.no_grad()       
    def cache_part_params(self):
        print("########## CACHING PARAMS ##########")
        self.model.eval()
        all_shifts = []
        all_scales = []
        for batch_idx, pcds in tqdm(enumerate(self.train_dataset), total=len(self.train_dataset)):
            shift, scale = self.model.module.get_params(pcds, self.local_rank)
            all_shifts.append(shift)
            all_scales.append(scale)
        self.all_shifts = all_shifts
        self.all_scales = all_scales
        print("########## PARAMS CACHED ##########")
        
        
    def train(self):
        # torch.autograd.set_detect_anomaly(True)
        self.model.eval()
        self.model.module.train_language_module()
        start_time = time.time()

        for batch_idx, pcds in enumerate(self.train_dataset):
            # pickle.dump({k:v.detach().cpu() for k, v in pcds.items()}, open("/orion/u/w4756677/diffusion/anchorDIff/weights/train_data.pkl", "wb"))
            # exit()
            
            self.optimizer.zero_grad()
            batch_size = self.train_dataset.batch_size
            
            losses = self.model.module.language_train_step(pcds, self.local_rank)
            all_loss, losses = parse_losses(losses)
            all_loss.backward()
            if self.max_norm is not None:
                nn.utils.clip_grad_norm_(self.model.parameters(), self.max_norm)
            self.optimizer.step()
            if self.max_epoch is None and self.scheduler is not None:
                self.scheduler.step()
            for m in self.model.module.language_encoder.parameters():
                print(m.mean())
            
            if check_interval(self.iter,self.log_interval) and self.iter>0:
                batch_size = self.train_dataset.batch_size
                ptime = time.time()-start_time
                eta_time = (self.total_iter-self.iter)*ptime/(batch_idx+1)
                eta_str = str(datetime.timedelta(seconds=int(eta_time)))
                data = dict(
                    name = self.cfg.name,
                    lr = self.optimizer.param_groups[0]['lr'] ,
                    iter = self.iter,
                    epoch = self.epoch,
                    batch_idx = batch_idx,
                    batch_size = batch_size,
                    total_loss = dist_utils.reduce_tensor(all_loss, self.world_size) if self.distributed else all_loss,
                    eta=eta_str
                )
                data.update(dist_utils.reduce_tensor(losses, self.world_size) if self.distributed else losses)
                data = sync(data)
                if self.logger is not None:
                    self.logger.log(data)
            
            self.iter+=1
            if self.finish:
                break

        self.epoch +=1
        if self.max_epoch and self.scheduler is not None:
            self.scheduler.step()


    @torch.no_grad()
    def val(self):
        if self.val_dataset is None:
            if self.logger is not None:
                self.logger.print_log("Please set Val dataset")
        else:
            if self.logger is not None:
                self.logger.print_log("Validating....")
            # TODO: need move eval into this function
            self.model.eval()
            self.model.module.eval_language_module()
            results = dict()
            for batch_idx, pcds in tqdm(enumerate(self.val_dataset),total=len(self.val_dataset)):
                if self.short_val and batch_idx > 0: break
                _result = self.model.module.language_edit_step(pcds, device=self.local_rank)
                for result, name in _result:
                    if name not in results.keys():
                        results[name] = []
                    results[name].append(sync(result, to_numpy=False))
            if self.no_eval:
                return
            for k, v in results.items():
                save_dict, eval_result = self.val_dataset.evaluate(v, self.save_num_batch, self.local_rank)
                save_file = build_file(os.path.join("/orion/u/mikacuy/anchored_diffusion/", self.work_dir), prefix=f"val/{k}_{self.prefix}_{self.epoch}.npz")
                self.logger.print_log(f"Saving to {save_file}....")
                print(save_dict.keys())
                np.savez_compressed(save_file,  **save_dict)
                eval_result = {k:v for k, v in eval_result.items()}
                self.logger.log(eval_result,iter=self.iter)
                    
    def generate_samples(self, num_gen, param_sample_num):
        if self.val_dataset is None:
            if self.logger is not None:
                self.logger.print_log("Please set Val dataset")
            return
        if self.logger is not None:
            self.logger.print_log("Generating....")
        # TODO: need move eval into this function
        self.model.eval()
        results = dict(pred=[], ref=[],  seg_mask_ref=[])
        fixed_id = [0,0,0,0]
        B = self.val_dataset.batch_size
        from difffacto.datasets.dataset_utils import shapenet_chair_part_distribution
        valid_ids = []
        probs = []
        for k, v in shapenet_chair_part_distribution.items():
            valid_ids.append(torch.tensor(list(map(float, list(k)))))
            probs.append(v)
        probs = torch.tensor(probs)
        valid_ids = torch.stack(valid_ids, dim=0).cuda()
        K=param_sample_num
        for batch_idx in tqdm(range(num_gen // self.val_dataset.batch_size)):
            idx = torch.multinomial(probs, self.val_dataset.batch_size, replacement=True)
            valid_id = valid_ids[idx]
            ctx, mean_per_point, logvar_per_point, _seg_mask, _valid_id = self.model.module.sample(B, fixed_id, valid_id, 'cuda', self.epoch, K=K)
            variance_per_point = torch.exp(logvar_per_point)
            C, N = variance_per_point.shape[1:]
            pred =self.model.module.decode(mean_per_point, ctx=ctx, device=self.local_rank, variance=variance_per_point, anchor_assignments=_seg_mask.to(torch.int32), valid_id=_valid_id) 
            results['pred'].append(pred['pred'].cpu().detach().numpy())
            results['seg_mask_ref'].append(_seg_mask.cpu().detach().numpy())
        for batch_idx, pcds in tqdm(enumerate(self.val_dataset),total=len(self.val_dataset)):
            results['ref'].append(pcds['ref'].cpu().detach().numpy())
        results = {k:np.concatenate(v, axis=0) for k, v in results.items()}
        save_file = build_file(self.work_dir, prefix=f"val/gen_fixed{''.join(map(str,fixed_id))}_{self.prefix}_{self.epoch}.pkl")
        self.logger.print_log(f"Saving to {save_file}....")
        pickle.dump(results, open(save_file, 'wb'))
        self.evaluate_gen(results['pred'], results['ref'], self.val_dataset.batch_size)
    
    def evaluate_gen(self, pred, ref, batch_size):
        preds = []
        refs = []
        print(pred.shape, ref.shape)
        for i in tqdm(range(pred.shape[0])):
            _pred = pred[i]
            if _pred.shape[0] > 2048:
                _pred = fps(torch.from_numpy(_pred).cuda().unsqueeze(0).contiguous(), 2048).detach().cpu().numpy()[0]
            pred_max = _pred.max(0).reshape(1, 3) # (1, 3)
            pred_min = _pred.min(0).reshape(1, 3) # (1, 3)
            pred_shift = ((pred_min + pred_max) / 2).reshape(1, 3)
            pred_scale = (pred_max - pred_min).max(-1)[0].reshape(1, 1) / 2
            _pred = (_pred - pred_shift) / pred_scale
            preds.append(_pred)
        for i in tqdm(range(ref.shape[0])):
            _ref = ref[i]
            if _ref.shape[0] > 2048:
                _ref = fps(torch.from_numpy(_ref).cuda().unsqueeze(0).contiguous(), 2048).detach().cpu().numpy()[0]
            ref_max = _ref.max(0).reshape(1, 3) # (1, 3)
            ref_min = _ref.min(0).reshape(1, 3) # (1, 3)
            ref_shift = ((ref_min + ref_max) / 2).reshape(1, 3)
            ref_scale = (ref_max - ref_min).max(-1)[0].reshape(1, 1) / 2
            _ref = (_ref - ref_shift) / ref_scale
            refs.append(_ref)

                        
        preds = torch.from_numpy(np.stack(preds, axis=0)).cuda()
        refs = torch.from_numpy(np.stack(refs, axis=0)).cuda()
        print(preds.shape, refs.shape)
        metrics = compute_all_metrics(preds, refs, batch_size, one_way=False)
        print(metrics)
        self.logger.log(metrics,iter=self.iter)
        
    def save(self):
        save_data = {
            "meta":{
                "epoch": self.epoch,
                "iter": self.iter,
                "max_iter": self.max_iter,
                "max_epoch": self.max_epoch,
                "config": self.cfg.dump()
            },
            "model":self.model.module.state_dict() if self.distributed else self.model.state_dict(),
            "encoder": self.model.module.encoder.encoder.state_dict(),
            "decoder": self.model.module.diffusion.state_dict(),
            "scheduler": self.scheduler.state_dict(),
            "optimizer": self.optimizer.state_dict()
        }
        if self.model.module.encoder.part_aligner is not None:
            save_data['param_sampler'] = self.model.module.encoder.part_aligner.state_dict()

        save_file = build_file(self.work_dir,prefix=f"checkpoints/ckpt_{self.epoch}.pth")
        torch.save(save_data,save_file)
        
    
    def load(self, load_path, model_only=False):
        map_location = {'cuda:%d' % 0: 'cuda:%d' % self.local_rank}
        resume_data = torch.load(load_path, map_location=map_location)

        if (not model_only):
            meta = resume_data.get("meta",dict())
            self.epoch = meta.get("epoch",self.epoch)
            self.iter = meta.get("iter",self.iter)
            self.max_iter = meta.get("max_iter",self.max_iter)
            self.max_epoch = meta.get("max_epoch",self.max_epoch)
            self.scheduler.load_state_dict(resume_data.get("scheduler",dict()))
            # self.optimizer.load_state_dict(resume_data.get("optimizer",dict()))
        #base_ckpt = {k.replace("module.", ""): v for k, v in resume_data['model'].items()}
        state_dict = resume_data['model']
        model_state_dict = self.model.state_dict()
        # print(model_state_dict.keys())
        for k in state_dict.keys():
            if k in model_state_dict.keys():
                if state_dict[k].shape != model_state_dict[k].shape:
                    self.logger.print_log(f"Skip loading parameter: {k}, "
                                f"required shape: {model_state_dict[k].shape}, "
                                f"loaded shape: {state_dict[k].shape}")
                    state_dict[k] = model_state_dict[k]
            else:
                self.logger.print_log(f"Dropping parameter {k}")
        for k in model_state_dict.keys():
            if k not in state_dict.keys():
                print(f"Layer {k} not loaded!")
        self.model.load_state_dict(state_dict,strict=False)

        self.logger.print_log(f"Loading model parameters from {load_path}")

    def resume(self):
        self.load(self.resume_path, model_only=self.cfg.model_only)

python/difffacto/runner/language_runner.py
diffFacto/diffFacto
import torch 
from torch import nn
from difffacto.datasets.custom import CustomDataset
from difffacto.config.config import get_cfg, save_cfg, save_args
from difffacto.utils.registry import build_from_cfg,MODELS,HOOKS
from difffacto.utils.misc import check_file, build_file, search_ckpt, check_interval, sync, parse_losses
import pickle
from difffacto.utils import misc
import numpy as np
class MixingRunner:
    def __init__(self, device, args):
        cfg = get_cfg()

        if args.local_rank == 0:
            self.logger = build_from_cfg(cfg.logger,HOOKS,work_dir=cfg.work_dir)
        else:
            self.logger = None
        self.distributed = args.distributed
        self.local_rank = args.local_rank
        self.world_size = args.world_size

        if args.seed is not None:
            if self.logger is not None:
                self.logger.print_log(f'Set random seed to {args.seed}, 'f'deterministic: {args.deterministic}')
            misc.set_random_seed(args.seed + args.local_rank, deterministic=args.deterministic) # seed + rank, for augmentation
        if args.distributed:
            assert args.local_rank == torch.distributed.get_rank() 
        self.cfg = cfg
        self.work_dir = cfg.work_dir
        self.ids = cfg.ids 

        self.resume_path = cfg.resume_path
        self.model = build_from_cfg(cfg.model,MODELS)
        if device=='cuda':
            self.model.to(self.local_rank)

        if args.distributed:
            # Sync BN
            if args.sync_bn:
                self.model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.model)
                if self.logger is not None:
                    self.logger.print_log('Using Synchronized BatchNorm ...')
            self.model = nn.parallel.DistributedDataParallel(self.model, device_ids=[args.local_rank % torch.cuda.device_count()], find_unused_parameters=True)
            if self.logger is not None:
                self.logger.print_log('Using Distributed Data parallel ...')
        else:
            if self.logger is not None:
                self.logger.print_log('Using Data parallel ...')
            self.model = nn.DataParallel(self.model).cuda()
        
        # optimizer & scheduler
        self.val_dataset = CustomDataset(cfg.data_dir, npoints=cfg.npoints, scale_mode=cfg.scale_mode, part_scale_mode=cfg.part_scale_mode, clip=cfg.clip, n_class=cfg.n_class)
        self.n_class=cfg.n_class
        assert len(self.ids) == cfg.n_class
        
        # only save file in the 0th process
        if  self.local_rank == 0:
            save_file = build_file(self.work_dir,prefix="config.yaml")
            save_arg = build_file(self.work_dir,prefix="args.yaml")
            save_args(save_arg, args)
            save_cfg(save_file)



        if (cfg.pretrained_weights):
            self.load(cfg.pretrained_weights, model_only=cfg.model_only)
        
        if self.resume_path is None:
            self.resume_path = search_ckpt(self.work_dir)
        if check_file(self.resume_path):
            self.resume()


    @torch.no_grad()
    def mixing(self):
        if self.val_dataset is None:
            if self.logger is not None:
                self.logger.print_log("Please set Val dataset")
        else:
            if self.logger is not None:
                self.logger.print_log("Validating....")
            # TODO: need move eval into this function
            self.model.eval()
            datas = [self.val_dataset[i] for i in self.ids if i > 0]
            mask = [data['seg_mask'] for data in datas]
            inps = [data['input'] for data in datas]
            inps = [inps[i][mask[i] == i] if self.ids[i] > 0 else np.zeros(1) for i in range(len(self.ids))]
            inps = [torch.from_numpy(i).cuda().float() for i in inps]
            out = self.model.module.combine_latent_specific(inps, 'cuda')
            out = {k:v.cpu().numpy() for k, v in out.items()}
            save_file = build_file(self.work_dir, prefix=f"val/mixing_.pkl")
            self.logger.print_log(f"Saving to {save_file}....")
            pickle.dump(out, open(save_file, 'wb'))
                    
    def load(self, load_path, model_only=False):
        map_location = {'cuda:%d' % 0: 'cuda:%d' % self.local_rank}
        resume_data = torch.load(load_path, map_location=map_location)

        state_dict = resume_data['model']
        model_state_dict = self.model.state_dict()
        # print(model_state_dict.keys())
        for k in state_dict.keys():
            if k in model_state_dict.keys():
                if state_dict[k].shape != model_state_dict[k].shape:
                    self.logger.print_log(f"Skip loading parameter: {k}, "
                                f"required shape: {model_state_dict[k].shape}, "
                                f"loaded shape: {state_dict[k].shape}")
                    state_dict[k] = model_state_dict[k]
            else:
                self.logger.print_log(f"Dropping parameter {k}")
        for k in model_state_dict.keys():
            if k not in state_dict.keys():
                print(f"Layer {k} not loaded!")
        self.model.load_state_dict(state_dict,strict=False)

        self.logger.print_log(f"Loading model parameters from {load_path}")

    def resume(self):
        self.load(self.resume_path, model_only=self.cfg.model_only)

python/difffacto/runner/mixing_runner.py
diffFacto/diffFacto
import torch 
from torch import nn
from torch.utils.data import DataLoader
import numpy as np
from difffacto.utils.misc import fps
from  difffacto.datasets.evaluation_utils import compute_all_metrics
from tqdm import tqdm
from difffacto.config.config import get_cfg, save_cfg, save_args
from difffacto.utils.registry import build_from_cfg,MODELS,SAMPLERS,DATASETS,OPTIMS,SCHEDULERS,HOOKS
from difffacto.utils.misc import check_file, build_file, search_ckpt, check_interval, sync, parse_losses
import pickle
import time
import datetime
from difffacto.utils import misc
from difffacto.utils import dist_utils
from difffacto.datasets.shapenet_parts import ShapeNetParts
from einops import rearrange


def normalize(pcs):
    pcs_max = pcs.max(1)[0].reshape(-1, 1, 3) # (1, 3)
    pcs_min = pcs.min(1)[0].reshape(-1, 1, 3) # (1, 3)
    pcs_shift = ((pcs_min + pcs_max) / 2).reshape(-1, 1, 3)
    pcs_scale = (pcs_max - pcs_min).reshape(-1, 1, 3) / 2
    return (pcs - pcs_shift) / pcs_scale

def get_data_iterator(iterable):
    """Allows training with DataLoaders in a single infinite loop:
        for i, data in enumerate(inf_generator(train_loader)):
    """
    iterator = iterable.__iter__()
    while True:
        try:
            yield iterator.__next__()
        except StopIteration:
            iterator = iterable.__iter__()
class PretrainRunner:
    def __init__(self, device, args):
        cfg = get_cfg()

        if args.local_rank == 0:
            self.logger = build_from_cfg(cfg.logger,HOOKS,work_dir=cfg.work_dir)
        else:
            self.logger = None
        self.distributed = args.distributed
        self.local_rank = args.local_rank
        self.world_size = args.world_size
        self.eval_both = cfg.eval_both
        self.eval_whole= cfg.eval_whole

        if args.seed is not None:
            if self.logger is not None:
                self.logger.print_log(f'Set random seed to {args.seed}, 'f'deterministic: {args.deterministic}')
            misc.set_random_seed(args.seed + args.local_rank, deterministic=args.deterministic) # seed + rank, for augmentation
        if args.distributed:
            assert args.local_rank == torch.distributed.get_rank() 
        self.cfg = cfg
        self.work_dir = cfg.work_dir
        self.save_num_batch = cfg.save_num_batch

        self.max_iter = cfg.max_iter
        assert self.max_iter is not None
        self.checkpoint_interval = cfg.checkpoint_interval
        self.eval_interval = cfg.eval_interval
        self.log_interval = cfg.log_interval
        self.resume_path = cfg.resume_path
        self.ret_traj=self.cfg.ret_traj
        self.ret_interval=self.cfg.ret_interval
        self.max_norm = cfg.max_norm
        self.model = build_from_cfg(cfg.model,MODELS)
        self.n_class = cfg.n_class
        assert self.n_class
        if device=='cuda':
            self.model.to(self.local_rank)

        if args.distributed:
            # Sync BN
            if args.sync_bn:
                self.model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.model)
                if self.logger is not None:
                    self.logger.print_log('Using Synchronized BatchNorm ...')
            self.model = nn.parallel.DistributedDataParallel(self.model, device_ids=[args.local_rank % torch.cuda.device_count()], find_unused_parameters=True)
            if self.logger is not None:
                self.logger.print_log('Using Distributed Data parallel ...')
        else:
            if self.logger is not None:
                self.logger.print_log('Using Data parallel ...')
            self.model = nn.DataParallel(self.model).cuda()
        
        # optimizer & scheduler
        parameters = self.model.parameters()
            
        self.optimizers = [build_from_cfg(cfg.optimizer,OPTIMS,params=list(self.model.module.encoder.encoder[i].parameters()) + list(self.model.module.diffusion[i].parameters())) for i in range(self.n_class)]
        self.schedulers = [build_from_cfg(cfg.scheduler, SCHEDULERS, optimizer=o) for o in self.optimizers]

        if self.distributed:
            assert cfg.dataset.train.batch_size % self.world_size == 0
            cfg.dataset.train.batch_size = cfg.dataset.train.batch_size // self.world_size
        
        self.train_dsets = []
        self.val_dsets = []
        self.train_bs = cfg.train_bs
        self.val_bs = cfg.val_bs
        self.npoints = cfg.npoints
        for i in range(self.n_class):
            train_dataset = ShapeNetParts(cfg.data_path_train[i], npoints=self.npoints, scale_mode=cfg.scale_mode, eval_mode='gen')
            val_dataset = ShapeNetParts(cfg.data_path_val[i], npoints=self.npoints, scale_mode=cfg.scale_mode, eval_mode='gen')
            train_sampler = DataLoader(train_dataset, cfg.train_bs, shuffle=True, num_workers=4, drop_last=True)
            val_sampler = DataLoader(val_dataset, cfg.val_bs, shuffle=False, num_workers=0)
            train_sampler = get_data_iterator(train_sampler)
            self.train_dsets.append(train_sampler)
            self.val_dsets.append(val_sampler)
        
        # only save file in the 0th process
        if  self.local_rank == 0:
            save_file = build_file(self.work_dir,prefix="config.yaml")
            save_arg = build_file(self.work_dir,prefix="args.yaml")
            save_args(save_arg, args)
            save_cfg(save_file)


        self.iter = 0
        self.epoch = 0

        self.total_iter = self.max_iter

        if (cfg.pretrained_weights):
            self.load(cfg.pretrained_weights, model_only=cfg.model_only)
        
        if self.resume_path is None:
            self.resume_path = search_ckpt(self.work_dir)
        if check_file(self.resume_path):
            self.resume()


    @property
    def finish(self):
        return self.iter>=self.max_iter
    
    def run(self):
        if self.logger is not None:
            self.logger.print_log("Start running")
        
        while not self.finish:
            self.train()
            if check_interval(self.iter,self.eval_interval) and self.local_rank == 0:
                self.val()
            if check_interval(self.iter,self.checkpoint_interval) and self.local_rank == 0:
                self.save()
    
    def train(self):
        # torch.autograd.set_detect_anomaly(True)
        self.model.train()
        l = 0
        res = {}
        for i, iterator in enumerate(self.train_dsets):
            for j in range(self.n_class):
                self.optimizers[j].zero_grad()
                
            nt = next(iterator)
            segmask = torch.eye(self.n_class).cuda()[i].reshape(1, 1, self.n_class).expand(self.train_bs, self.npoints, -1)
            losses = self.model.module.pretrain_part(nt['input'].cuda(), segmask, part_id=i, device='cuda')
            all_loss, losses = parse_losses(losses)
            all_loss.backward()
            self.optimizers[i].step()
            self.schedulers[i].step()
            l += all_loss.detach().cpu().item()
            res.update({f'part_{i}/{k}':v for k, v in losses.items()})
        if check_interval(self.iter,self.log_interval) and self.iter>0:
            data = dict(
                name = self.cfg.name,
                lr = self.optimizers[0].param_groups[0]['lr'] ,
                iter = self.iter,
                batch_size = self.train_bs,
                total_loss = l / self.n_class,
            )
            data.update(res)
            data = sync(data)
            if self.logger is not None:
                self.logger.log(data)
        self.iter+=1
        

    @torch.no_grad()
    def val(self):
        if self.logger is not None:
            self.logger.print_log("Validating....")
        self.model.eval()
        results = [[] for _ in range(self.n_class)]
        l = max([len(d) for d in self.val_dsets])
        for i in range(l):
            r = self.model.module.pretrain_validate(self.val_bs, self.npoints, 'cuda')
            for j, rr in enumerate(r):
                results[j].append(rr.detach().cpu())
        results = [torch.cat(rrr, dim=0) for rrr in results]
        results = [normalize(result) for result in results]
        for j in range(self.n_class):
            save_file = build_file(self.work_dir, prefix=f"val/pretrain_{self.iter}_part_{j}.pt")
            torch.save(results[j], save_file)
        for j in range(self.n_class):
            val_dset = self.val_dsets[j]
            ref = []
            for batch_idx, pcds in tqdm(enumerate(val_dset),total=len(val_dset)):
                ref.append(pcds['input'])
            ref = torch.cat(ref, dim=0)
            ref = normalize(ref)
            metrics = compute_all_metrics(results[j][:ref.shape[0]].cuda(), ref.cuda(), self.val_bs, accelerated_cd=True)
            eval_result = {f"part_{j}_{k}":v for k, v in metrics.items()}
            self.logger.log(eval_result,iter=self.iter)
                    
    def save(self):
        save_data = {
            "meta":{
                "iter": self.iter,
                "max_iter": self.max_iter,
                "config": self.cfg.dump()
            },
            "model":self.model.module.state_dict() if self.distributed else self.model.state_dict(),
            "encoder": self.model.module.encoder.encoder.state_dict(),
            "scheduler": [self.schedulers[i].state_dict() for i in range(self.n_class)],
            "optimizer": [self.optimizers[i].state_dict() for i in range(self.n_class)],
        }

        save_file = build_file(self.work_dir,prefix=f"checkpoints/ckpt_{self.epoch}.pth")
        torch.save(save_data,save_file)
        
    
    def load(self, load_path, model_only=False):
        map_location = {'cuda:%d' % 0: 'cuda:%d' % self.local_rank}
        resume_data = torch.load(load_path, map_location=map_location)

        if (not model_only):
            meta = resume_data.get("meta",dict())
            self.epoch = meta.get("epoch",self.epoch)
            self.iter = meta.get("iter",self.iter)
            self.max_iter = meta.get("max_iter",self.max_iter)
            self.max_epoch = meta.get("max_epoch",self.max_epoch)
            for i in range(self.n_class):
                self.schedulers[i].load_state_dict(resume_data.get("scheduler",[dict() for _ in range(self.n_class)])[i])
                self.optimizers[i].load_state_dict(resume_data.get("optimizer",[dict() for _ in range(self.n_class)])[i])
        #base_ckpt = {k.replace("module.", ""): v for k, v in resume_data['model'].items()}
        state_dict = resume_data['model']
        model_state_dict = self.model.state_dict()
        # print(model_state_dict.keys())
        for k in state_dict.keys():
            if k in model_state_dict.keys():
                if state_dict[k].shape != model_state_dict[k].shape:
                    self.logger.print_log(f"Skip loading parameter: {k}, "
                                f"required shape: {model_state_dict[k].shape}, "
                                f"loaded shape: {state_dict[k].shape}")
                    state_dict[k] = model_state_dict[k]
            else:
                self.logger.print_log(f"Dropping parameter {k}")
        for k in model_state_dict.keys():
            if k not in state_dict.keys():
                print(f"Layer {k} not loaded!")
        self.model.load_state_dict(state_dict,strict=False)

        self.logger.print_log(f"Loading model parameters from {load_path}")

    def resume(self):
        self.load(self.resume_path, model_only=self.cfg.model_only)

python/difffacto/runner/pretrain_runner.py
diffFacto/diffFacto
import torch 
from torch import nn
import numpy as np
from difffacto.utils.misc import fps
from  difffacto.datasets.evaluation_utils import compute_all_metrics
from tqdm import tqdm
from difffacto.config.config import get_cfg, save_cfg, save_args
from difffacto.utils.registry import build_from_cfg,MODELS,SAMPLERS,DATASETS,OPTIMS,SCHEDULERS,HOOKS
from difffacto.utils.misc import check_file, build_file, search_ckpt, check_interval, sync, parse_losses
import pickle
import os
import time
import datetime
from difffacto.utils import misc
from difffacto.utils import dist_utils
from einops import rearrange
class Runner:
    def __init__(self, device, args):
        cfg = get_cfg()

        if args.local_rank == 0:
            self.logger = build_from_cfg(cfg.logger,HOOKS,work_dir=cfg.work_dir)
        else:
            self.logger = None
            
        print(cfg.work_dir)
        self.distributed = args.distributed
        self.no_eval=args.no_eval
        self.local_rank = args.local_rank
        self.world_size = args.world_size
        self.short_val = args.short_val
        self.eval_both = cfg.eval_both
        self.eval_whole= cfg.eval_whole
        self.prefix=args.prefix

        if args.seed is not None:
            if self.logger is not None:
                self.logger.print_log(f'Set random seed to {args.seed}, 'f'deterministic: {args.deterministic}')
            misc.set_random_seed(args.seed + args.local_rank, deterministic=args.deterministic) # seed + rank, for augmentation
        if args.distributed:
            assert args.local_rank == torch.distributed.get_rank() 
        self.cfg = cfg
        self.work_dir = cfg.work_dir
        self.save_num_batch = cfg.save_num_batch

        self.max_epoch = cfg.max_epoch 
        self.max_iter = cfg.max_iter
        assert (self.max_iter is None)^(self.max_epoch is None),"You must set max_iter or max_epoch"

        self.checkpoint_interval = cfg.checkpoint_interval
        self.eval_interval = cfg.eval_interval
        self.log_interval = cfg.log_interval
        self.resume_path = cfg.resume_path
        self.ret_traj=self.cfg.ret_traj
        self.ret_interval=self.cfg.ret_interval
        self.max_norm = cfg.max_norm
        self.model = build_from_cfg(cfg.model,MODELS)
        if device=='cuda':
            self.model.to(self.local_rank)

        if args.distributed:
            # Sync BN
            if args.sync_bn:
                self.model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.model)
                if self.logger is not None:
                    self.logger.print_log('Using Synchronized BatchNorm ...')
            self.model = nn.parallel.DistributedDataParallel(self.model, device_ids=[args.local_rank % torch.cuda.device_count()], find_unused_parameters=True)
            if self.logger is not None:
                self.logger.print_log('Using Distributed Data parallel ...')
        else:
            if self.logger is not None:
                self.logger.print_log('Using Data parallel ...')
            self.model = nn.DataParallel(self.model).cuda()
        
        # optimizer & scheduler
        if cfg.train_aligner:
            if cfg.joint_train:
                lr_scale = cfg.lr_scale if cfg.lr_scale is not None else 1.0
                parameters = [{'params': [], 'lr': lr_scale, 'name': 'base_group'}, {'params': [], 'name': 'aligner_group'}]
                for n, p in self.model.named_parameters():
                    if 'part_aligner' in n:
                        print(f"{n}: {cfg.optimizer.lr}")
                        parameters[1]['params'].append(p)
                    else:
                        print(f"{n}: {cfg.optimizer.lr * lr_scale}")
                        parameters[0]['params'].append(p)
            else:
                parameters = self.model.module.encoder.part_aligner.parameters()
        elif cfg.train_cvae:
            parameters = list(self.model.module.encoder.part_aligner.parameters()) + list(self.model.module.encoder.cvae_encoder.parameters())
            if hasattr(self.model.module.encoder, 'ref_encoder'):
                parameters += list(self.model.module.encoder.ref_encoder.parameters())
        else:
            parameters = self.model.parameters()
            
        self.optimizer = build_from_cfg(cfg.optimizer,OPTIMS,params=parameters)
        self.scheduler = build_from_cfg(cfg.scheduler, SCHEDULERS, optimizer=self.optimizer)

        if self.distributed:
            assert cfg.dataset.train.batch_size % self.world_size == 0
            cfg.dataset.train.batch_size = cfg.dataset.train.batch_size // self.world_size
        self.train_dataset, self.train_sampler = build_from_cfg(cfg.dataset.train,DATASETS, distributed=args.distributed)
        self.val_dataset, _ = build_from_cfg(cfg.dataset.val,DATASETS, distributed=False)
        self.cimle = cfg.cimle
        self.cimle_start_epoch = cfg.cimle_start_epoch if cfg.cimle_start_epoch is not None else 0
        self.cache_interval = cfg.cimle_cache_interval
        
        # only save file in the 0th process
        if  self.local_rank == 0:
            save_file = build_file(self.work_dir,prefix="config.yaml")
            save_arg = build_file(self.work_dir,prefix="args.yaml")
            save_args(save_arg, args)
            save_cfg(save_file)


        self.iter = 0
        self.epoch = 0

        if self.max_epoch:
            if (self.train_dataset):
                self.total_iter = self.max_epoch * len(self.train_dataset)
            else:
                self.total_iter = 0
        else:
            self.total_iter = self.max_iter

        if (cfg.pretrained_weights):
            self.load(cfg.pretrained_weights, model_only=cfg.model_only)
        
        if self.resume_path is None:
            self.resume_path = search_ckpt(self.work_dir)
        if check_file(self.resume_path):
            self.resume()


    @property
    def finish(self):
        if self.max_epoch:
            return self.epoch>=self.max_epoch
        else:
            return self.iter>=self.max_iter
    
    def run(self):
        if self.logger is not None:
            self.logger.print_log("Start running")
        
        while not self.finish:
            if self.distributed:
                self.train_sampler.set_epoch(self.epoch)
            if self.cimle and self.epoch >= self.cimle_start_epoch and check_interval(self.epoch - self.cimle_start_epoch, self.cache_interval) and self.local_rank == 0:
                self.cache_noise()
            self.train()
            if check_interval(self.epoch,self.eval_interval) and self.local_rank == 0:
                self.val()
            if check_interval(self.epoch,self.checkpoint_interval) and self.local_rank == 0:
                self.save()
    
    @torch.no_grad()
    def cache_noise(self):
        print("########## CACHING NOISE FOR CIMLE ##########")
        self.model.eval()
        all_noises = []
        all_ids = []
        for batch_idx, pcds in tqdm(enumerate(self.train_dataset), total=len(self.train_dataset)):
            noises = self.model.module.cache_noise(pcds, self.local_rank)
            all_noises.append(noises)
            all_ids.append(pcds['id'])
        self.train_dataset.dataset.store_noise(all_noises, all_ids)
        
    @torch.no_grad()
    def cache_noise(self):
        print("########## CACHING NOISE FOR CIMLE ##########")
        self.model.eval()
        all_noises = []
        all_ids = []
        for batch_idx, pcds in tqdm(enumerate(self.train_dataset), total=len(self.train_dataset)):
            noises = self.model.module.cache_noise(pcds, self.local_rank, eval_whole=self.eval_whole)
            all_noises.append(noises)
            all_ids.append(pcds['id'])
        self.train_dataset.dataset.store_noise(all_noises, all_ids)
        
        print("########## NOISE FOR CIMLE CACHED ##########")
    @torch.no_grad()       
    def cache_part_params(self):
        print("########## CACHING PARAMS ##########")
        self.model.eval()
        all_shifts = []
        all_scales = []
        for batch_idx, pcds in tqdm(enumerate(self.train_dataset), total=len(self.train_dataset)):
            shift, scale = self.model.module.get_params(pcds, self.local_rank)
            all_shifts.append(shift)
            all_scales.append(scale)
        self.all_shifts = all_shifts
        self.all_scales = all_scales
        print("########## PARAMS CACHED ##########")
        
    @torch.no_grad()       
    def interpolate_two_sets(self, data1, data2, part_id):
        assert data1.shape[0] == data2.shape[0]
        total = data1.shape[0]
        print("---Normalizing---")
        from difffacto.datasets.dataset_utils import pc_norm
        norm_data1 = data1.copy()
        norm_data2 = data2.copy()
        all_shift = np.zeros((total, 3, 4))
        all_scale  = np.ones((total, 3, 4))
        params = []
        for i in range(total):
            _data1 = data1[i]
            _data2 = data2[i]
            _norm_data1 = norm_data1[i]
            _norm_data2 = norm_data2[i]
            param_scale, param_shift = np.ones((3, 4)), np.zeros((3, 4))
            for j in range(4):
                m1 = _data1[..., 3] == j
                if np.any(m1):
                    part_point = _data1[m1, :3]
                    norm_m1, shift, scale = pc_norm(part_point, 'shape_canonical', clip=False, to_torch=False)
                    _norm_data1[m1, :3] = norm_m1
                    param_scale[:, j] = scale
                    param_shift[:, j] = shift
                m2 = _data2[..., 3] == j
                if np.any(m2):
                    part_point = _data2[m2, :3]
                    norm_m2, _, _ = pc_norm(part_point, 'shape_canonical', clip=False, to_torch=False)
                    _norm_data2[m2, :3] = norm_m2
            norm_data1[i] = _norm_data1
            norm_data2[i] = _norm_data2
            all_shift[i] = param_shift
            all_scale[i] = param_scale
        print("---DONE---")
        batch_size=8
        out_interpolation = {'pred':[], 'pred_seg_mask':[]}
        for idx in tqdm(range(200 // batch_size)):
            # save_file = build_file(self.work_dir, prefix=f"val/interpolation_part{part_id}_{self.prefix}.pkl")
            # pickle.dump(dict(set1=norm_data1[idx*batch_size:(idx+1)*batch_size][..., :3], set2=norm_data2[idx*batch_size:(idx+1)*batch_size][..., :3]), open(save_file, 'wb'))
            # exit()
            _set1 = torch.from_numpy(norm_data1[idx*batch_size:(idx+1)*batch_size]).cuda().float()
            _set2 = torch.from_numpy(norm_data2[idx*batch_size:(idx+1)*batch_size]).cuda().float()
            _param_shift = torch.from_numpy(all_shift[idx*batch_size:(idx+1)*batch_size]).cuda().float()
            _param_scale = torch.from_numpy(all_scale[idx*batch_size:(idx+1)*batch_size]).cuda().float()
            
            pred, seg_mask = self.model.module.interpolate_two_shapes(_set1, _set2, part_id, _param_shift, _param_scale.cuda(), mid_num=5)
            out_interpolation['pred'].append(pred.cpu().numpy())
            out_interpolation['pred_seg_mask'].append(seg_mask.cpu().numpy())
            # break
        out_interpolation = {k:np.concatenate(v, axis=0) for k, v in out_interpolation.items()}
        save_file = build_file(self.work_dir, prefix=f"val/interpolation_part{part_id}_{self.prefix}.pkl")
        print(f"Saving to {save_file}....")
        pickle.dump(out_interpolation, open(save_file, 'wb'))
        
    @torch.no_grad()       
    def part_mixing(self, data):
        assert isinstance(data, list)
        print("---Normalizing---")
        from difffacto.datasets.dataset_utils import pc_norm
        norm_data = []
        for k, d in enumerate(data):
            mask = torch.ones_like(d)[:,0] * k
            B, N, _ = d.shape
            mean = d.mean(1).reshape(B, 1, 3)
            std = d.std(1).reshape(B, 1, 3)
            if np.any(std == 0.0):
                std[std == 0.0] = 1.0
            normed_d = (d - mean) / std
            norm_data.append(norm_d)
        for i in range(total):
            _data1 = data1[i]
            _data2 = data2[i]
            _norm_data1 = norm_data1[i]
            _norm_data2 = norm_data2[i]
            for j in range(4):
                m1 = _data1[..., 3] == j
                if np.any(m1):
                    part_point = _data1[m1, :3]
                    norm_m1, _, _ = pc_norm(part_point, 'shape_canonical', clip=False, to_torch=False)
                    _norm_data1[m1, :3] = norm_m1
                m2 = _data2[..., 3] == j
                if np.any(m2):
                    part_point = _data1[m2, :3]
                    norm_m2, _, _ = pc_norm(part_point, 'shape_canonical', clip=False, to_torch=False)
                    _norm_data2[m2, :3] = norm_m2
            norm_data1[i] = _norm_data1
            norm_data2[i] = _norm_data2
        print("---DONE---")
        batch_size=8
        out_interpolation = {'pred':[], 'pred_seg_mask':[]}
        for idx in tqdm(range(total // batch_size)):
            _set1 = torch.from_numpy(norm_data1[idx*batch_size:(idx+1)*batch_size]).cuda().float()
            _set2 = torch.from_numpy(norm_data2[idx*batch_size:(idx+1)*batch_size]).cuda().float()
            
            pred, seg_mask = self.model.module.interpolate_two_shapes(_set1, _set2, part_id, mid_num=10)
            out_interpolation['pred'].append(pred.cpu().numpy())
            out_interpolation['pred_seg_mask'].append(seg_mask.cpu().numpy())
        out_interpolation = {k:np.concatenate(v, axis=0) for k, v in out_interpolation.items()}
        save_file = build_file(self.work_dir, prefix=f"val/interpolation_part{part_id}_{self.prefix}.pkl")
        print(f"Saving to {save_file}....")
        pickle.dump(out_interpolation, open(save_file, 'wb'))
    def train(self):
        # torch.autograd.set_detect_anomaly(True)
        self.model.train()
        start_time = time.time()

        for batch_idx, pcds in enumerate(self.train_dataset):
            # pickle.dump({k:v.detach().cpu() for k, v in pcds.items()}, open("/orion/u/w4756677/diffusion/anchorDIff/weights/train_data.pkl", "wb"))
            # exit()
            
            self.optimizer.zero_grad()
            batch_size = self.train_dataset.batch_size
            
            losses = self.model(pcds, self.local_rank, iter=self.iter, epoch=self.epoch)
            all_loss, losses = parse_losses(losses)
            all_loss.backward()
            if self.max_norm is not None:
                nn.utils.clip_grad_norm_(self.model.parameters(), self.max_norm)
            self.optimizer.step()
            if self.max_epoch is None and self.scheduler is not None:
                self.scheduler.step()
            
            
            if check_interval(self.iter,self.log_interval) and self.iter>0:
                batch_size = self.train_dataset.batch_size
                ptime = time.time()-start_time
                eta_time = (self.total_iter-self.iter)*ptime/(batch_idx+1)
                eta_str = str(datetime.timedelta(seconds=int(eta_time)))
                data = dict(
                    name = self.cfg.name,
                    lr = self.optimizer.param_groups[0]['lr'] ,
                    iter = self.iter,
                    epoch = self.epoch,
                    batch_idx = batch_idx,
                    batch_size = batch_size,
                    total_loss = dist_utils.reduce_tensor(all_loss, self.world_size) if self.distributed else all_loss,
                    eta=eta_str
                )
                data.update(dist_utils.reduce_tensor(losses, self.world_size) if self.distributed else losses)
                data = sync(data)
                if self.logger is not None:
                    self.logger.log(data)
            
            self.iter+=1
            if self.finish:
                break

        self.epoch +=1
        if self.max_epoch and self.scheduler is not None:
            self.scheduler.step()


    @torch.no_grad()
    def val(self):
        if self.val_dataset is None:
            if self.logger is not None:
                self.logger.print_log("Please set Val dataset")
        else:
            if self.logger is not None:
                self.logger.print_log("Validating....")
            # TODO: need move eval into this function
            self.model.eval()
            results = dict()
            for batch_idx, pcds in tqdm(enumerate(self.val_dataset),total=len(self.val_dataset)):
                # pickle.dump({k:v.detach().cpu() for k, v in pcds.items()}, open("/orion/u/w4756677/diffusion/anchorDIff/weights/train_data.pkl", "wb"))
                # exit()
                if self.short_val and batch_idx > 0: break
                _result = self.model(pcds, device=self.local_rank)
                for result, name in _result:
                    if name not in results.keys():
                        results[name] = []
                    results[name].append(sync(result, to_numpy=False))
            if self.no_eval:
                return
            for k, v in results.items():
                save_dict, eval_result = self.val_dataset.evaluate(v, self.save_num_batch, self.local_rank)
                save_file = build_file(self.work_dir, prefix=f"val/{k}_{self.prefix}_{self.epoch}.npz")
                self.logger.print_log(f"Saving to {save_file}....")
                np.savez_compressed(save_file,  **save_dict)
                eval_result = {k:v for k, v in eval_result.items()}
                self.logger.log(eval_result,iter=self.iter)
            if self.eval_both and not self.short_val:
                results = dict()
                for batch_idx, pcds in tqdm(enumerate(self.train_dataset),total=len(self.train_dataset)):
                    if self.short_val and batch_idx > 0: break
                    _result = self.model(pcds, device=self.local_rank)
                    for result, name in _result:
                        if name not in results.keys():
                            results[name] = []
                        results[name].append(sync(result, to_numpy=False))
                if self.no_eval:
                    return
                for k, v in results.items():
                    save_dict, eval_result = self.val_dataset.evaluate(v, self.save_num_batch, self.local_rank)
                    save_file = build_file(self.work_dir, prefix=f"val/train_set_{k}_{self.prefix}_{self.epoch}.npz")
                    self.logger.print_log(f"Saving to {save_file}....")
                    np.savez_compressed(save_file, **save_dict)
                    pickle.dump(save_dict, open(save_file, 'wb'),)
                    eval_result = {f"train_{k}":v for k, v in eval_result.items()}
                    self.logger.log(eval_result,iter=self.iter)
                    
    def generate_samples(self, num_gen, param_sample_num):
        if self.val_dataset is None:
            if self.logger is not None:
                self.logger.print_log("Please set Val dataset")
            return
        if self.logger is not None:
            self.logger.print_log("Generating....")
        # TODO: need move eval into this function
        self.model.eval()
        results = dict(pred=[], ref=[],  seg_mask_ref=[])
        fixed_id = [0,0,0,0]
        B = self.val_dataset.batch_size
        from difffacto.datasets.dataset_utils import shapenet_chair_part_distribution
        valid_ids = []
        probs = []
        for k, v in shapenet_chair_part_distribution.items():
            valid_ids.append(torch.tensor(list(map(float, list(k)))))
            probs.append(v)
        probs = torch.tensor(probs)
        valid_ids = torch.stack(valid_ids, dim=0).cuda()
        K=param_sample_num
        for batch_idx in tqdm(range(num_gen // self.val_dataset.batch_size)):
            idx = torch.multinomial(probs, self.val_dataset.batch_size, replacement=True)
            valid_id = valid_ids[idx]
            ctx, mean_per_point, logvar_per_point, _seg_mask, _valid_id = self.model.module.sample(B, fixed_id, valid_id, 'cuda', self.epoch, K=K)
            variance_per_point = torch.exp(logvar_per_point)
            C, N = variance_per_point.shape[1:]
            pred =self.model.module.decode(mean_per_point, ctx=ctx, device=self.local_rank, variance=variance_per_point, anchor_assignments=_seg_mask.to(torch.int32), valid_id=_valid_id) 
            results['pred'].append(pred['pred'].cpu().detach().numpy())
            results['seg_mask_ref'].append(_seg_mask.cpu().detach().numpy())
        for batch_idx, pcds in tqdm(enumerate(self.val_dataset),total=len(self.val_dataset)):
            results['ref'].append(pcds['ref'].cpu().detach().numpy())
        results = {k:np.concatenate(v, axis=0) for k, v in results.items()}
        save_file = build_file(self.work_dir, prefix=f"val/gen_fixed{''.join(map(str,fixed_id))}_{self.prefix}_{self.epoch}.pkl")
        self.logger.print_log(f"Saving to {save_file}....")
        pickle.dump(results, open(save_file, 'wb'))
        self.evaluate_gen(results['pred'], results['ref'], self.val_dataset.batch_size)
    
    def evaluate_gen(self, pred, ref, batch_size):
        preds = []
        refs = []
        print(pred.shape, ref.shape)
        for i in tqdm(range(pred.shape[0])):
            _pred = pred[i]
            if _pred.shape[0] > 2048:
                _pred = fps(torch.from_numpy(_pred).cuda().unsqueeze(0).contiguous(), 2048).detach().cpu().numpy()[0]
            pred_max = _pred.max(0).reshape(1, 3) # (1, 3)
            pred_min = _pred.min(0).reshape(1, 3) # (1, 3)
            pred_shift = ((pred_min + pred_max) / 2).reshape(1, 3)
            pred_scale = (pred_max - pred_min).max(-1)[0].reshape(1, 1) / 2
            _pred = (_pred - pred_shift) / pred_scale
            preds.append(_pred)
        for i in tqdm(range(ref.shape[0])):
            _ref = ref[i]
            if _ref.shape[0] > 2048:
                _ref = fps(torch.from_numpy(_ref).cuda().unsqueeze(0).contiguous(), 2048).detach().cpu().numpy()[0]
            ref_max = _ref.max(0).reshape(1, 3) # (1, 3)
            ref_min = _ref.min(0).reshape(1, 3) # (1, 3)
            ref_shift = ((ref_min + ref_max) / 2).reshape(1, 3)
            ref_scale = (ref_max - ref_min).max(-1)[0].reshape(1, 1) / 2
            _ref = (_ref - ref_shift) / ref_scale
            refs.append(_ref)

                        
        preds = torch.from_numpy(np.stack(preds, axis=0)).cuda()
        refs = torch.from_numpy(np.stack(refs, axis=0)).cuda()
        print(preds.shape, refs.shape)
        metrics = compute_all_metrics(preds, refs, batch_size, one_way=False)
        print(metrics)
        self.logger.log(metrics,iter=self.iter)
        
    def save(self):
        save_data = {
            "meta":{
                "epoch": self.epoch,
                "iter": self.iter,
                "max_iter": self.max_iter,
                "max_epoch": self.max_epoch,
                "config": self.cfg.dump()
            },
            "model":self.model.module.state_dict() if self.distributed else self.model.state_dict(),
            "encoder": self.model.module.encoder.encoder.state_dict(),
            "decoder": self.model.module.diffusion.state_dict(),
            "scheduler": self.scheduler.state_dict(),
            "optimizer": self.optimizer.state_dict()
        }
        if self.model.module.encoder.part_aligner is not None:
            save_data['param_sampler'] = self.model.module.encoder.part_aligner.state_dict()

        save_file = build_file(self.work_dir,prefix=f"checkpoints/ckpt_{self.epoch}.pth")
        torch.save(save_data,save_file)
        
    
    def load(self, load_path, model_only=False):
        map_location = {'cuda:%d' % 0: 'cuda:%d' % self.local_rank}
        resume_data = torch.load(load_path, map_location=map_location)

        if (not model_only):
            meta = resume_data.get("meta",dict())
            self.epoch = meta.get("epoch",self.epoch)
            self.iter = meta.get("iter",self.iter)
            self.max_iter = meta.get("max_iter",self.max_iter)
            self.max_epoch = meta.get("max_epoch",self.max_epoch)
            self.scheduler.load_state_dict(resume_data.get("scheduler",dict()))
            # self.optimizer.load_state_dict(resume_data.get("optimizer",dict()))
        #base_ckpt = {k.replace("module.", ""): v for k, v in resume_data['model'].items()}
        state_dict = resume_data['model']
        model_state_dict = self.model.state_dict()
        # print(model_state_dict.keys())
        for k in state_dict.keys():
            if k in model_state_dict.keys():
                if state_dict[k].shape != model_state_dict[k].shape:
                    self.logger.print_log(f"Skip loading parameter: {k}, "
                                f"required shape: {model_state_dict[k].shape}, "
                                f"loaded shape: {state_dict[k].shape}")
                    state_dict[k] = model_state_dict[k]
            else:
                self.logger.print_log(f"Dropping parameter {k}")
        for k in model_state_dict.keys():
            if k not in state_dict.keys():
                print(f"Layer {k} not loaded!")
        self.model.load_state_dict(state_dict,strict=False)

        self.logger.print_log(f"Loading model parameters from {load_path}")

    def resume(self):
        self.load(self.resume_path, model_only=self.cfg.model_only)

python/difffacto/runner/runner.py
diffFacto/diffFacto
from .sampler import Uniform
python/difffacto/samplers/__init__.py
diffFacto/diffFacto

from abc import ABC, abstractmethod

import numpy as np
import torch as th
from difffacto.utils.registry import SAMPLERS

class ScheduleSampler(ABC):
    """
    A distribution over timesteps in the diffusion process, intended to reduce
    variance of the objective.
    By default, samplers perform unbiased importance sampling, in which the
    objective's mean is unchanged.
    However, subclasses may override sample() to change how the resampled
    terms are reweighted, allowing for actual changes in the objective.
    """

    @abstractmethod
    def weights(self):
        """
        Get a numpy array of weights, one per diffusion step.
        The weights needn't be normalized, but must be positive.
        """

    def sample(self, batch_size, device):
        """
        Importance-sample timesteps for a batch.
        :param batch_size: the number of timesteps.
        :param device: the torch device to save to.
        :return: a tuple (timesteps, weights):
                 - timesteps: a tensor of timestep indices.
                 - weights: a tensor of weights to scale the resulting losses.
        """
        w = self.weights()
        p = w / np.sum(w)
        indices_np = np.random.choice(len(p), size=(batch_size,), p=p)
        indices = th.from_numpy(indices_np).long().to(device)
        weights_np = 1 / (len(p) * p[indices_np])
        weights = th.from_numpy(weights_np).float().to(device)
        return indices, weights

@SAMPLERS.register_module()
class Uniform(ScheduleSampler):
    def __init__(self, num_timesteps):
        super().__init__()
        self.num_timesteps = num_timesteps
        self.weight = np.ones([self.num_timesteps])
    def weights(self):
        return self.weight
python/difffacto/samplers/sampler.py
diffFacto/diffFacto
import enum 

class ModelMeanType(enum.Enum):
    """
    Which type of output the model predicts.
    """

    PREVIOUS_X = 1  # the model predicts x_{t-1}
    START_X = 2  # the model predicts x_0
    EPSILON = 3  # the model predicts epsilon
    DRIFTED_EPSILON1=4
    DRIFTED_EPSILON2=5
    DRIFTED_EPSILON3=6
    DRIFTED_EPSILON4=7
    EPSILON_AND_ANCHOR=8
    SCALED_EPSILON=9
    DRIFTED_EPSILON5=10


class ModelVarType(enum.Enum):
    """
    What is used as the model's output variance.
    The LEARNED_RANGE option has been added to allow the model to predict
    values between FIXED_SMALL and FIXED_LARGE, making its job easier.
    """

    LEARNED = 1
    FIXED_SMALL = 2
    FIXED_LARGE = 3
    LEARNED_RANGE = 4
    
class LossType(enum.Enum):
    MSE = 1  # use raw MSE loss (and KL when learning variances)
    RESCALED_MSE = 2  # use raw MSE loss (with RESCALED_KL when learning variances)
    KL = 3  # use the variational lower-bound
    RESCALED_KL = 4  # like KL, but rescale to estimate the full VLB

    def is_vb(self):
        return self == LossType.KL or self == LossType.RESCALED_KL


model_mean_dict = {
    'x_0':ModelMeanType.START_X, 
    'prev_x':ModelMeanType.PREVIOUS_X, 
    'epsilon':ModelMeanType.EPSILON, 
    'drifted_epsilon1':ModelMeanType.DRIFTED_EPSILON1, 
    'drifted_epsilon2':ModelMeanType.DRIFTED_EPSILON2, 
    'drifted_epsilon4':ModelMeanType.DRIFTED_EPSILON4, 
    'drifted_epsilon3':ModelMeanType.DRIFTED_EPSILON3, 
    'drifted_epsilon5':ModelMeanType.DRIFTED_EPSILON5,
    'epsilon_and_anchor':ModelMeanType.EPSILON_AND_ANCHOR, 
    'scaled_epsilon':ModelMeanType.SCALED_EPSILON
    }
model_var_dict = {'learned':ModelVarType.LEARNED, 'learned_range':ModelVarType.LEARNED_RANGE, 'fixed_large':ModelVarType.FIXED_LARGE, 'fixed_small':ModelVarType.FIXED_SMALL}
model_loss_dict = {'kl':LossType.KL, 'rescaled_kl':LossType.RESCALED_KL, 'mse':LossType.MSE, 'rescaled_mse':LossType.RESCALED_MSE}


python/difffacto/utils/constants.py
diffFacto/diffFacto
import os

import torch
import torch.multiprocessing as mp
from torch import distributed as dist



def init_dist(launcher, backend='nccl', **kwargs):
    if mp.get_start_method(allow_none=True) is None:
        mp.set_start_method('spawn')
    if launcher == 'pytorch':
        _init_dist_pytorch(backend, **kwargs)
    else:
        raise ValueError(f'Invalid launcher type: {launcher}')


def _init_dist_pytorch(backend, **kwargs):
    # TODO: use local_rank instead of rank % num_gpus
    rank = int(os.environ['RANK'])
    num_gpus = torch.cuda.device_count()
    torch.cuda.set_device(rank % num_gpus)
    dist.init_process_group(backend=backend, **kwargs)
    print(f'init distributed in rank {torch.distributed.get_rank()}')


def get_dist_info():
    if dist.is_available():
        initialized = dist.is_initialized()
    else:
        initialized = False
    if initialized:
        rank = dist.get_rank()
        world_size = dist.get_world_size()
    else:
        rank = 0
        world_size = 1
    return rank, world_size


def reduce_tensor(data, world_size):
    '''
        for acc kind, get the mean in each gpu
    '''
    def _reduce_tensor(tensor):
        if isinstance(tensor, dict):
            return {k: _reduce_tensor(v) for k, v in tensor.items()}
        elif isinstance(tensor, list):
            return [_reduce_tensor(v) for v in tensor]
        elif isinstance(tensor, torch.Tensor):
            rt = tensor.clone()
            torch.distributed.all_reduce(rt, op=torch.distributed.ReduceOp.SUM)
            rt /= world_size
            return rt

    return _reduce_tensor(data)

def gather_tensor(tensor, args):
    output_tensors = [tensor.clone() for _ in range(args.world_size)]
    torch.distributed.all_gather(output_tensors, tensor)
    concat = torch.cat(output_tensors, dim=0)
    return concat

python/difffacto/utils/dist_utils.py
diffFacto/diffFacto
from .misc import build_file, current_time
from .registry import HOOKS, build_from_cfg 
import time 
import os
from tensorboardX import SummaryWriter
from difffacto.config.config import get_cfg

@HOOKS.register_module()
class TextLogger:
    def __init__(self,work_dir):
        save_file = build_file(work_dir,prefix="textlog/log_"+time.strftime("%Y_%m_%d_%H_%M_%S", time.localtime())+".txt")
        self.log_file = open(save_file,"a")

    def log(self,data):
        msg = ",".join([f"{k}:{d}" for k,d in data.items()])
        msg= current_time()+msg+"\n"
        self.log_file.write(msg)
        self.log_file.flush()

@HOOKS.register_module()
class TensorboardLogger:
    def __init__(self,work_dir):
        self.cfg = get_cfg()
        tensorboard_dir = os.path.join(work_dir,"tensorboard")
        self.writer = SummaryWriter(tensorboard_dir,flush_secs=10)

    def log(self,data):
        step = data["iter"]
        for k,d in data.items():
            if k in ["iter","epoch","batch_idx","times","batch_size"]:
                continue
            if isinstance(d,str):
                continue
            self.writer.add_scalar(k,d,global_step=step)

@HOOKS.register_module()
class RunLogger:
    def __init__(self,work_dir,loggers=["TextLogger","TensorboardLogger"]):
        self.loggers = [build_from_cfg(l,HOOKS,work_dir=work_dir) for l in loggers]
    def log(self,data,**kwargs):
        data.update(kwargs)
        data = {k:d.item() if hasattr(d,"item") else d for k,d in data.items()}
        for logger in self.loggers:
            logger.log(data)
        self.print_log(data)
    
    def get_time(self, s):
        s = int(s)
        days = s // 60 // 60 // 24
        hours = s // 60 // 60 % 24
        minutes = s // 60 % 60
        seconds = s % 60
        return f' [{days}D:{hours}H:{minutes}M:{seconds}S] '

    def print_log(self,msg):
        if isinstance(msg,dict):
            msgs = []
            for k,d in msg.items():
                if (k == "remain_time"):
                    msgs.append(f" {k}:{self.get_time(d)}")
                else:
                    msgs.append(f" {k}:{d:.7f}" if isinstance(d,float) else f" {k}:{d}")
            msg = ",".join(msgs)
        print(current_time(),msg)
python/difffacto/utils/loggers.py
diffFacto/diffFacto
import os 
import warnings
import torch
import numpy as np
import torch.nn.functional as F
import random
from pointnet2_ops import pointnet2_utils
import time
import glob
import math

"""
Adapted from: https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/nn.py#L124
"""

from typing import Callable, Iterable, Sequence, Union


# Taken from https://discuss.pytorch.org/t/implementing-truncated-normal-initializer/4778/15
def truncated_normal(tensor, mean=0, std=1, trunc_std=2):
    size = tensor.shape
    tmp = tensor.new_empty(size + (4,)).normal_()
    valid = (tmp < trunc_std) & (tmp > -trunc_std)
    ind = valid.max(-1, keepdim=True)[1]
    tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1))
    tensor.data.mul_(std).add_(mean)
    return tensor

def str_list(argstr):
    return list(argstr.split(','))

def checkpoint(
    func: Callable[..., Union[torch.Tensor, Sequence[torch.Tensor]]],
    inputs: Sequence[torch.Tensor],
    params: Iterable[torch.Tensor],
    flag: bool,
):
    """
    Evaluate a function without caching intermediate activations, allowing for
    reduced memory at the expense of extra compute in the backward pass.
    :param func: the function to evaluate.
    :param inputs: the argument sequence to pass to `func`.
    :param params: a sequence of parameters `func` depends on but does not
                   explicitly take as arguments.
    :param flag: if False, disable gradient checkpointing.
    """
    if flag:
        args = tuple(inputs) + tuple(params)
        return CheckpointFunction.apply(func, len(inputs), *args)
    else:
        return func(*inputs)


class CheckpointFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, run_function, length, *args):
        ctx.run_function = run_function
        ctx.input_tensors = list(args[:length])
        ctx.input_params = list(args[length:])
        with torch.no_grad():
            output_tensors = ctx.run_function(*ctx.input_tensors)
        return output_tensors

    @staticmethod
    def backward(ctx, *output_grads):
        ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
        with torch.enable_grad():
            # Fixes a bug where the first op in run_function modifies the
            # Tensor storage in place, which is not allowed for detach()'d
            # Tensors.
            shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
            output_tensors = ctx.run_function(*shallow_copies)
        input_grads = torch.autograd.grad(
            output_tensors,
            ctx.input_tensors + ctx.input_params,
            output_grads,
            allow_unused=True,
        )
        del ctx.input_tensors
        del ctx.input_params
        del output_tensors
        return (None, None) + input_grads
def timestep_embedding(timesteps, dim, max_period=10000):
    """
    Create sinusoidal timestep embeddings.
    :param timesteps: a 1-D Tensor of N indices, one per batch element.
                      These may be fractional.
    :param dim: the dimension of the output.
    :param max_period: controls the minimum frequency of the embeddings.
    :return: an [N x dim] Tensor of positional embeddings.
    """
    half = dim // 2
    freqs = torch.exp(
        -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
    ).to(device=timesteps.device)
    args = timesteps[:, None].float() * freqs[None]
    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
    if dim % 2:
        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
    return embedding

def pad_unmasked_element_by_zero(tensor, mask):
    # tensor[B, n_seg]
    # mask [B, n_seg]
    B, n_seg = tensor.shape
    neg_mask = not mask.to(bool)
    zeroed_tensor = tensor.masked_fill(neg_mask, 0.)
    return zeroed_tensor
    

def get_index(coarse, fine):
    num_coarse = coarse.shape[1]
    B, num_fine = fine.shape[:2]
    num_per_anchor = num_fine // num_coarse
    dist = torch.sum(torch.pow(fine.reshape(B, num_coarse, num_per_anchor, 3).unsqueeze(2) - coarse.unsqueeze(2).unsqueeze(1), 2), -1) #[B, num_coarse, num_coarse, num_per_anchor]
    idx = torch.argmin(dist, dim=2).to(torch.int32)
    return idx.reshape(B, -1)


def parse_losses(losses):
    _losses = dict()
    for loss_name, loss_value in losses.items():
        if isinstance(loss_value, torch.Tensor):
            _losses[loss_name] = loss_value.mean()
        elif isinstance(loss_value, list):
            _losses[loss_name] = sum(_loss.mean() for _loss in loss_value)
        else:
            raise TypeError(
                '{} is not a tensor or list of tensors'.format(loss_name))

    total_loss = sum(_value for _key, _value in _losses.items() if 'loss' in _key)
    return total_loss, _losses

def current_time():
    return time.asctime( time.localtime(time.time()))

def filter_points(data, npoints):
    n = data.shape[0]
    if n <= npoints:
        return data 
    else:
        idx = np.random.RandomState(2020).permutation(n)[:npoints]
        return data[idx]

def search_ckpt(work_dir):
    files = glob.glob(os.path.join(work_dir,"checkpoints/ckpt_*.pth"))
    if len(files)==0:
        return None
    files = sorted(files,key=lambda x:int(x.split("_")[-1].split(".pth")[0]))
    return files[-1]  

def check_interval(step,step_interval):
    if step is None or step_interval is None:
        return False 
    if step % step_interval==0:
        return True 
    return False 


def sync(data, to_numpy=True):
    """
        sync data and convert data to numpy
    """
    def _sync(data):
        if isinstance(data,(list,tuple)):
            data =  [_sync(d) for d in data]
        elif isinstance(data,dict):
            data = {k:_sync(d) for k,d in data.items()}
        elif isinstance(data,torch.Tensor):
            data = data.detach().cpu()
            if to_numpy:
                data = data.numpy()
        elif not isinstance(data,(int,float,str,np.ndarray)):
            raise ValueError(f"{type(data)} is not supported")
        return data
    return _sync(data) 

def build_file(work_dir,prefix):
    """ build file and makedirs the file parent path """
    work_dir = os.path.abspath(work_dir)
    prefixes = prefix.split("/")
    file_name = prefixes[-1]
    prefix = "/".join(prefixes[:-1])
    if len(prefix)>0:
        work_dir = os.path.join(work_dir,prefix)
    os.makedirs(work_dir,exist_ok=True)
    file = os.path.join(work_dir,file_name)
    return file 

def assign_anchor(data, anchors):
    dist = torch.sum(torch.pow(data.unsqueeze(2) - anchors.unsqueeze(1), 2), dim=-1) # B, N, M
    anchor_assignments = dist.argmin(dim=-1).to(data.device)
    return anchor_assignments

def fps(data, number, ret_id=False):
    '''
        data B N 3
        number int
    '''
    fps_idx = pointnet2_utils.furthest_point_sample(data, number) 
    fps_data = pointnet2_utils.gather_operation(data.transpose(1, 2).contiguous(), fps_idx).transpose(1,2).contiguous()
    if ret_id:
        return fps_data, fps_idx
    return fps_data

def separate_point_cloud(xyz, num_points, crop, fixed_points = None, padding_zeros = False):
    '''
     seprate point cloud: usage : using to generate the incomplete point cloud with a setted number.
    '''
    _,n,c = xyz.shape

    assert n == num_points
    assert c == 3
    if crop == num_points:
        return xyz, None
        
    INPUT = []
    CROP = []
    for points in xyz:
        if isinstance(crop,list):
            num_crop = np.random.randint(crop[0],crop[1])
        else:
            num_crop = crop

        points = points.unsqueeze(0)

        if fixed_points is None:       
            center = F.normalize(torch.randn(1,1,3),p=2,dim=-1).to(xyz.device)
        else:
            if isinstance(fixed_points,list):
                fixed_point = np.random.sample(fixed_points,1)[0]
            else:
                fixed_point = fixed_points
            center = fixed_point.reshape(1,1,3).cuda()

        distance_matrix = torch.norm(center.unsqueeze(2) - points.unsqueeze(1), p =2 ,dim = -1)  # 1 1 2048

        idx = torch.argsort(distance_matrix,dim=-1, descending=False)[0,0] # 2048

        if padding_zeros:
            input_data = points.clone()
            input_data[0, idx[:num_crop]] =  input_data[0,idx[:num_crop]] * 0

        else:
            input_data = points.clone()[0, idx[num_crop:]].unsqueeze(0) # 1 N 3

        crop_data =  points.clone()[0, idx[:num_crop]].unsqueeze(0)

        if isinstance(crop,list):
            INPUT.append(fps(input_data,2048))
            CROP.append(fps(crop_data,2048))
        else:
            INPUT.append(input_data)
            CROP.append(crop_data)

    input_data = torch.cat(INPUT,dim=0)# B N 3
    crop_data = torch.cat(CROP,dim=0)# B M 3

    return input_data.contiguous(), crop_data.contiguous()

def check_file(file,ext=None):
    if file is None:
        return False
    if not os.path.exists(file):
        warnings.warn(f"{file} is not exists")
        return False
    if not os.path.isfile(file):
        warnings.warn(f"{file} must be a file")
        return False
    if ext:
        if not os.path.splitext(file)[1] in ext:
            # warnings.warn(f"the type of {file} must be in {ext}")
            return False
    return True

def mean_flat(tensor):
    """
    Take the mean over all non-batch dimensions.
    """
    return tensor.mean(dim=list(range(1, len(tensor.shape))))

def reparameterize_gaussian(mean, logvar):
    std = torch.exp(0.5 * logvar)
    eps = torch.randn(std.size()).to(mean)
    return mean + std * eps

def gaussian_logprob(z, mean, logvar=None ):
    log_z = -0.5 * np.log(2 * np.pi)
    std_squared = torch.exp(logvar)
    return -logvar + log_z - (z - mean).pow(2) / (2. * std_squared)

def gaussian_entropy(logvar, dim=1):
    const = 0.5 * float(logvar.size(dim)) * (1. + np.log(np.pi * 2))
    ent = 0.5 * logvar.sum(dim=dim, keepdim=False) + const
    return ent

def standard_normal_logprob(z):
    dim = z.size(-1)
    log_z = -0.5 * dim * np.log(2 * np.pi)
    return log_z - z.pow(2) / 2
def gaussian_log_likelihood(z, mean=None, logvar=None, dim=3):
    """
    Compute the log probability of x belonging to the gaussian with mean and log variance
    
    param x: [B, 3, N] tensor
    param mean: [B, 3, N] tensor 
    param log_var: [B, 3, N] tensor 
    """
    log_z = -0.5 * np.log(2 * np.pi) * dim
    if logvar is None:
        std_squared = 1.
        logvar = 0.
    else:
        std_squared = torch.exp(logvar) if isinstance(logvar, torch.Tensor) else math.exp(logvar)
    if mean is None:
        mean = 0.
    return -logvar + log_z - (z - mean).pow(2) / (2. * std_squared)
    
    
def set_random_seed(seed, deterministic=False):
    """Set random seed.
    Args:
        seed (int): Seed to be used.
        deterministic (bool): Whether to set the deterministic option for
            CUDNN backend, i.e., set `torch.backends.cudnn.deterministic`
            to True and `torch.backends.cudnn.benchmark` to False.
            Default: False.

    # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
    if cuda_deterministic:  # slower, more reproducible
        cudnn.deterministic = True
        cudnn.benchmark = False
    else:  # faster, less reproducible
        cudnn.deterministic = False
        cudnn.benchmark = True

    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
"""
def get_part_loss(anchors, gt_pcds, loss_func):
    loss = 0.
    actual_num=0
    num_classes=anchors.shape[1]
    for i in range(num_classes):
        batch_idx = part_means[:, i, 0] == i 
        if batch_idx.any():
            actual_num += 1
            batch_means = part_means[batch_idx, i, 1:]
            anchor = anchors[batch_idx, i]
            loss += loss_func(anchor, batch_means)

    return loss / actual_num

"""
python/difffacto/utils/misc.py
diffFacto/diffFacto
class Registry:
    def __init__(self):
        self._modules = {}

    def register_module(self, name=None,module=None):
        def _register_module(module):
            key = name
            if key is None:
                key = module.__name__
            assert key not in self._modules,f"{key} is already registered."
            self._modules[key]=module
            return module

        if module is not None:
            return _register_module(module)

        return _register_module

    def get(self,name):
        assert name in self._modules,f"{name} is not registered."
        return self._modules[name]


def build_from_cfg(cfg,registry,**kwargs):
    if isinstance(cfg,str):
        return registry.get(cfg)(**kwargs)
    elif isinstance(cfg,dict):
        args = cfg.copy()
        args.update(kwargs)
        obj_type = args.pop('type')
        obj_cls = registry.get(obj_type)
        try:
            module = obj_cls(**args)
        except TypeError as e:
            if "<class" not in str(e):
                e = f"{obj_cls}.{e}"
            raise TypeError(e)

        return module
    elif isinstance(cfg,list):
        from jittor import nn
        return nn.Sequential([build_from_cfg(c,registry,**kwargs) for c in cfg])
    elif cfg is None:
        return None
    else:
        raise TypeError(f"type {type(cfg)} not support")


DATASETS = Registry()
MODELS = Registry()
ENCODERS = Registry()
DECOMPOSERS = Registry()
DIFFUSIONS = Registry()
NETS = Registry()
SCHEDULERS = Registry()
HOOKS = Registry()
LOSSES = Registry()
OPTIMS = Registry()
SAMPLERS = Registry()
METRICS = Registry()
SEGMENTORS = Registry()
GENERATORS = Registry()
DISCRIMINATORS = Registry()
python/difffacto/utils/registry.py
diffFacto/diffFacto
from setuptools import setup, find_packages
import os
path = os.path.join(os.path.dirname(__file__),"python")

setup(
    name="difffacto",
    packages=find_packages(path),
    package_dir={'': "python"},
)
setup.py
diffFacto/diffFacto
import argparse 
import torch
from anchor_diff.datasets.evaluation_utils import _pairwise_EMD_CD_ 
from anchor_diff.datasets.dataset_utils import load_ply_withNormals, shapenet_part_normal_cat_to_id
import numpy as np
from anchor_diff.utils.misc import fps
import pickle
import json
import os
from tqdm import tqdm
from pointnet2_ops import pointnet2_utils


seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
                'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46],
                'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27],
                'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40],
                'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--datadir',
        type=str) 
    parser.add_argument(
        '--prefix',
        type=str) 
    args = parser.parse_args()
    # import torch
    # data1 = torch.load("/orion/u/w4756677/diffusion/anchorDIff/weights/chair-separate-v2.pt")
    # data2 = torch.load("/orion/u/w4756677/diffusion/anchorDIff/weights/chair-separate.pt")
    # ratio = [700*4, 700*4, 690*4, 140*4]
    # data = torch.cat([data1, data2], dim=0)[:2800]
    # print(data.shape)
    # out_data = []
    # out_mask = []
    # for i in tqdm(range(data.shape[0])):
    #     _data = data[i][..., :3].cuda()
    #     _mask = data[i][..., 3].cuda()
    #     if i > ratio[3]:
    #         wo_arm = _mask != 3 
    #         _data = _data[wo_arm]
    #         _mask = _mask[wo_arm]
    #     if i > ratio[2]:
    #         wo_leg = _mask != 2
    #         _data = _data[wo_leg]
    #         _mask = _mask[wo_leg]
    #     _data, ids = fps(_data[None], 2048, True)
    #     _mask = pointnet2_utils.gather_operation(_mask.reshape(1, 1, -1).contiguous(), ids).reshape(2048)
    #     print(_data.shape, _mask.shape)
    #     out_data.append(_data[0].cpu().numpy())
    #     out_mask.append(_mask.cpu().numpy())
    # out_data = np.stack(out_data, axis=0)
    # out_mask = np.stack(out_mask, axis=0)
    # pickle.load(open())
    # print(out_data.shape, out_mask.shape)
    # pickle.dump(dict(pred=out_data, pred_seg_mask=out_mask), open("/orion/u/w4756677/diffusion/anchorDIff/weights/lion_part/lion_2800_chair.pkl", "wb"))
    # exit()
    # data1 = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/weights/lion_Chair_2048.pkl", 'rb'))
    # print(data.keys(), data1.keys())
    # n = data1['pred'].shape[0]
    # pickle.dump( dict(pred=np.concatenate([data1['pred'], data['pred'][:2800 - n]], axis=0), pred_seg_mask=np.concatenate([data1['pred_seg_mask'], data['pred_seg_mask'][:2800 - n]], axis=0)), open(f"/orion/u/w4756677/diffusion/anchorDIff/weights/lion_2800_chair.pkl", 'wb'))
    # exit()
    # data_out = []
    # for i in range(704):
    #     _data = data[i]
    #     perm = torch.randperm(8192)
    #     idx = perm[:2048]
    #     data_out.append(_data[idx])
    # data_out = torch.stack(data_out, dim=0)
    # print(data_out.shape)
    # pickle.dump(dict(pred=data_out[..., :3].numpy().astype(float), pred_seg_mask=data_out[..., 3].numpy().astype(int)), open("/orion/u/w4756677/diffusion/anchorDIff/weights/lion_chair_v2.pkl", 'wb'))
    # exit()
    # print(data['parts'].shape)
    # data = data['parts'].cpu().numpy()
    # np.save("/orion/u/w4756677/diffusion/anchorDIff/weights/pruned_part_2000_part_2.npy", data)
    # num = 2800 if args.cat == "Chair" else 1360
    data = pickle.load(open(args.datadir, 'rb'))
    print(data.keys())
    # print(data.keys())
    pred = [v[:700] for k, v in data.items() if f"pred_sample" in k][:4]
    pred = np.concatenate(pred, axis=0)
    # print(len(pred))
    # print(pred.shape)
    # del data['pred']
    mask = data['pred_seg_mask'][:700, None].repeat(4, axis=1).reshape(2800, 2048)

    print(pred.shape, mask.shape)
    pickle.dump(dict(pred=pred, pred_seg_mask=mask), open(f"/orion/u/w4756677/diffusion/anchorDIff/weights/variant_data/{args.prefix}.pkl", 'wb'))
    # print(pred.shape)
    exit()
    # pickle.dump( dict(pred=pred[:2800], pred_seg_mask=mask[:2800]), open(f"/orion/u/w4756677/diffusion/anchorDIff/weights/v2_dpm_{num}_{args.cat}.pkl", 'wb'))
    # exit()
    preds = [v for k, v in data.items() if 'pred' in k]
    # print(preds)
    print(len(preds))
    N = 704 if args.cat == 'Chair' else 340
    # factors = [(N, 1), (N, 2), (N, 3), (N, 4), (N, 5)]
    factors = [(N, 4)]
    # pred = torch.stack([-pred[..., 2], pred[..., 1], pred[..., 0]], dim=-1)
    
    for size, m in factors:
        _preds = [preds[p] for p in range(m)]
        _preds = [p[:size] for p in _preds]
        _mask = mask[None, :size].repeat(m, axis=0).reshape(-1, 2048)
        _preds = np.concatenate(_preds, axis=0)
        print(_preds.shape)
        pickle.dump( dict(pred=_preds, seg_mask_ref=_mask), open(f"/orion/u/w4756677/diffusion/anchorDIff/weights/624_n100_{size}_{m}.pkl", 'wb'))
                        

tools/assemble_data.py
diffFacto/diffFacto
import numpy as np 
import os 
from anchor_diff.datasets.dataset_utils import load_ply_withNormals
from tqdm import tqdm
def save_data_as_text(data_dir, save_dir):
    fns = sorted(os.listdir(data_dir))  # Get each txt file under the folder according to the path and put it in the list
    # print(fns[0][0:-4])
    ids = [f[:-8] for f in fns if (f[-7:] == 'ori.ply')]
    print(len(ids))
    for i in tqdm(range(len(ids))):
        fn = ids[i]
        point_set = load_ply_withNormals(os.path.join(data_dir, fn + '_ori.ply')).astype(np.float32) # size 20488,7 read in this txt file, a total of 20488 points, each point xyz rgb + sub category label
        
        seg = np.loadtxt(os.path.join(data_dir, fn + '_ori.labels')).astype(np.long) - 1 + 24# Get the label of the sub category
        mask = np.abs(point_set).sum(1) != 0
        valid_point_set = point_set[mask]
        valid_seg = seg[mask]
        _data = np.concatenate([valid_point_set, valid_seg[:, None]], axis=-1)
        np.savetxt(os.path.join(save_dir, f"{fn}.txt"), _data)
    
save_data_as_text('/orion/u/w4756677/datasets/coalace_data_processed/03636649', '/orion/u/w4756677/datasets/colasce_data_txt/03636649')
tools/colasce_to_txt.py
diffFacto/diffFacto
from typing import Dict
from collections import defaultdict
import numpy as np
from pandas import DataFrame, Series
from six.moves import cPickle



part_names = ["back", "seat", "leg", "arm"]

part_semantic_groups = {
    "back": ["back"],
    "seat": ["seat"],
    "leg": ["leg", "wheel", "base"],
    "arm": ["arm"],
}

def unpickle_data(file_name, python2_to_3=False):
    """Restore data previously saved with pickle_data().
    :param file_name: file holding the pickled data.
    :param python2_to_3: (boolean), if True, pickle happened under python2x, unpickling under python3x.
    :return: a generator over the un-pickled items.
    Note, about implementing the python2_to_3 see
        https://stackoverflow.com/questions/28218466/unpickling-a-python-2-object-with-python-3
    """

    in_file = open(file_name, "rb")
    if python2_to_3:
        size = cPickle.load(in_file, encoding="latin1")
    else:
        size = cPickle.load(in_file)

    for _ in range(size):
        if python2_to_3:
            yield cPickle.load(in_file, encoding="latin1")
        else:
            yield cPickle.load(in_file)
    in_file.close()

def convert_labels_to_one_hot(labels, n_classes=3):
    # Convert labels to indicators.
    n_utter = len(labels)
    targets = np.array(labels, dtype=np.int64)
    target_oh = np.zeros(shape=(n_utter, n_classes))
    for i in range(n_utter):
        target_oh[i, targets[i]] = 1
    assert np.all(np.sum(target_oh, axis=1) == 1)
    return target_oh


def get_part_indicator(text: Series, word2int: Dict):
    part_indicator = np.zeros((len(text), len(part_names)), dtype=np.float32)

    part_idx_set = dict()
    for k, v in part_semantic_groups.items():
        part_idx_set[k] = set(list(map(lambda x:word2int[x], v)))
        
    for i in range(len(text)):
        token_set = set(text[i])
        for j, part_name in enumerate(part_names):
            current_set = part_idx_set[part_name]
            other_set = set.union(*part_idx_set.values()) - current_set
            if len(set.intersection(token_set, current_set)) and not len(
                set.intersection(token_set, other_set)
            ):
                part_indicator[i, j] = 1

    part_mask = np.sum(part_indicator, 1) > 0

    return part_indicator, part_mask


def get_mask_of_game_data(
    game_data: DataFrame,
    word2int: Dict,
    only_correct: bool,
    only_easy_context: bool,
    max_seq_len: int,
    only_one_part_name: bool,
):
    """
    only_correct (if True): mask will be 1 in location iff human listener predicted correctly.
    only_easy (if True): uses only easy context examples (more dissimilar triplet chairs)
    max_seq_len: drops examples with len(utterance) > max_seq_len
    only_one_part_name (if True): uses only utterances describing only one part in the give set.
    """
    mask = np.array(game_data.correct)
    if not only_correct:
        mask = np.ones_like(mask, dtype=np.bool)

    if only_easy_context:
        context_mask = np.array(game_data.context_condition == "easy", dtype=np.bool)
        mask = np.logical_and(mask, context_mask)

    short_mask = np.array(
        game_data.text.apply(lambda x: len(x)) <= max_seq_len, dtype=np.bool
    )
    mask = np.logical_and(mask, short_mask)

    part_indicator, part_mask = get_part_indicator(game_data.text, word2int)
    if only_one_part_name:
        mask = np.logical_and(mask, part_mask)

    return mask, part_indicator


def shuffle_game_geometries(geo_ids, labels, parts=None, random_seed=None):
    """e.g. if [a, b, c] with label 1 makes it [b, a, c] with label 0."""
    if random_seed is not None:
        np.random.seed(random_seed)
    shuffle = np.random.shuffle
    for i in range(len(geo_ids)):
        idx = [0, 1, 2]
        shuffle(idx)
        geo_ids[i] = geo_ids[i][idx]
        labels[i] = labels[i][idx]
        if parts is not None:
            parts[i] = parts[i][idx]

    if parts is not None:
        return geo_ids, labels, parts
    else:
        return geo_ids, labels


def pad_text_symbols_with_zeros(
    text, max_seq_len, dtype=np.int64, force_zero_end=False
):
    """
    force_zero_end (bool) if True every sequence will end with zero, alternatively,
        sequences with equal or more elements than max_seq_len will end with the element at that max_seq_len position.
    """
    text_padded = []
    seq_len = []

    if force_zero_end:
        last_slot = 1
    else:
        last_slot = 0

    for sentence in text:
        pad_many = max_seq_len - len(sentence) + last_slot
        if pad_many > 0:
            text_padded.append(
                np.pad(sentence, (0, pad_many), "constant", constant_values=0)
            )
            seq_len.append(len(sentence))
        else:
            keep_same = min(max_seq_len, len(sentence))
            kept_text = sentence[:keep_same]
            if force_zero_end:
                kept_text.append(0)
            text_padded.append(kept_text)
            seq_len.append(keep_same)

    text_padded = np.array(text_padded, dtype=dtype)
    seq_len = np.array(seq_len, dtype)
    return text_padded, seq_len


def group_geometries(geo_ids, target_indicators=None):
    """If the geometries associated with an utterance are the same, group them together.
    Input:
        geo_ids (N x 3): N triplets of integers.
        target_indicators: (N x 3) indicator of which was the target geometry for each of the geo_ids.
    Returns:
         A dictionary where each key maps to the rows of the ``geo_ids`` that are comprised by the same set of integers, and same ``target_indicators``, if the latter is not None.
    """
    if target_indicators is not None and not np.all(
        np.unique(target_indicators) == [0, 1]
    ):
        raise ValueError("provide one-hot indicators.")

    n_triplets = len(geo_ids)
    if target_indicators is not None:
        if n_triplets != len(target_indicators):
            raise ValueError()

    groups = defaultdict(list)
    for i in range(n_triplets):
        g = geo_ids[i].astype(np.int)
        if target_indicators is not None:
            t = g[np.where(target_indicators[i])]
            key = tuple(np.hstack([sorted(g), t]))
        else:
            key = tuple(sorted(g))
        groups[key].append(i)
    return groups


def group_target_geometries(geo_ids, indicators):
    """Returns a dictionary mapping each geo_id at the rows of geo_ids of which it was
    used as a target. I.e., the values of each key in the result are an equivalence class.
    """

    if not np.all(np.unique(indicators) == [0, 1]):
        raise ValueError("provide one-hot indicators.")

    groups = defaultdict(list)  # Hold for each target_geometry the rows which is used.
    n_triplets = len(geo_ids)
    for i in range(n_triplets):
        target_geo_i = geo_ids[i][np.where(indicators[i])][0]
        groups[target_geo_i].append(i)
    return groups



def rotation(pc):
    rot = np.array([[0, 0, -1], [0, 1, 0], [-1, 0, 0.0]])
    rot_pc = pc @ rot
    return rot_pc


def get_pc_norm_info_each_sample(geos):
    k, l, _ = geos.shape

    normed_geos = geos.copy()
    centroids = np.mean(normed_geos, axis=1, keepdims=True)  # [3, 1, 3]

    normed_geos = normed_geos - centroids

    m = np.max(np.sqrt(np.sum(normed_geos ** 2, axis=-1)))

    return centroids, m



tools/data_utils.py
diffFacto/diffFacto
import argparse 
import torch
from anchor_diff.datasets.evaluation_utils import _pairwise_EMD_CD_, EMD_CD
from anchor_diff.datasets.dataset_utils import load_ply_withNormals, shapenet_part_normal_cat_to_id
import numpy as np
from anchor_diff.utils.misc import fps
import pickle
import json
import os
from tqdm import tqdm


seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
                'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46],
                'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27],
                'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40],
                'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--cat',
        type=str) 
    parser.add_argument(
        '--npoint',
        type=int,
        default=2048) 
    args = parser.parse_args()
    cat_id = shapenet_part_normal_cat_to_id[args.cat]
    
    ref_part_pointcloud = [[], [], [], []]
    root = '/orion/u/w4756677/datasets/coalace_data_processed' if args.cat != 'Car' else '/orion/u/w4756677/datasets/shapenetcore_partanno_segmentation_benchmark_v0_normal'
    with open(f"{root}/train_test_split/shuffled_test_file_list.json", 'r') as f:
        train_ids = set([str(d.split('/')[2]) for d in json.load(f)])
    
    fns = sorted(os.listdir(f"{root}/{cat_id}"))  # Get each txt file under the folder according to the path and put it in the list
    if args.cat != 'Car':
        densePoints = [f for f in fns if (f[-7:] == 'ori.ply')]
        labels = [f for f in fns if (f[-10:] == 'ori.labels')]
        
        labels = [os.path.join(f"{root}/{cat_id}", fn) for fn in labels if ((fn[0:-11] in train_ids))]
        densePoints = [os.path.join(f"{root}/{cat_id}", fn) for fn in densePoints if ((fn[0:-8] in train_ids)  )]
        print(len(labels), len(densePoints))
    else:
        densePoints = [f for f in fns if (f[-3:] == 'txt')]
        densePoints = [os.path.join(f"{root}/{cat_id}", fn) for fn in densePoints if ((fn[0:-4] in train_ids)  )]
        print(len(densePoints))
        
    for i in tqdm(range(len(densePoints))):
        if args.cat != 'Car':
            _ref = torch.from_numpy(load_ply_withNormals(densePoints[i]).astype(np.float32)).cuda()
            _ref_seg_mask = torch.from_numpy(np.loadtxt(labels[i]).astype(np.int32)).cuda() - 1
        else:
            data = torch.from_numpy(np.loadtxt(densePoints[i]).astype(np.float32)).cuda()
            _ref_seg_mask = data[:, -1].to(torch.int32) - seg_classes[args.cat][0] # Get the label of the sub category
            _ref = data[:, 0:3]
            
        for j in range(4):
            ref_idx = _ref_seg_mask == j
            if args.cat != "Car":
                ref_points = _ref[ref_idx].unsqueeze(0)
                if ref_points.abs().sum() != 0:
                    if ref_points.shape[1] > args.npoint:
                        ref_points = fps(ref_points, args.npoint)
                    ref_part_pointcloud[j].append(ref_points)
            else:
                if torch.any(ref_idx):
                    ref_points = _ref[ref_idx].unsqueeze(0)
                    if ref_points.shape[1] > 100:
                        while ref_points.shape[1] < args.npoint:
                            ref_points = torch.cat([ref_points, ref_points], dim=1)
                        if ref_points.shape[1] > args.npoint:
                            ref_points = fps(ref_points, args.npoint)
                        ref_part_pointcloud[j].append(ref_points)
    thresh = 2000
    out = []
    t = 0
    out = []
    l = ['back', 'seat', 'leg', 'arm']
    ref_part_pointcloud = [torch.cat(l, dim=0) for l in ref_part_pointcloud]
    for i in range(4):
        part_pcd = ref_part_pointcloud[i]
        part_pcd_max = part_pcd.max(1)[0].reshape(-1, 1,3) # (1, 3)
        part_pcd_min = part_pcd.min(1)[0].reshape(-1, 1,3) # (1, 3)
        part_pcd_shift = ((part_pcd_min + part_pcd_max) / 2).reshape(-1, 1, 3)
        part_pcd_scale = (part_pcd_max - part_pcd_min).reshape(-1, 1, 3) / 2
        part_pcd = (part_pcd - part_pcd_shift) / part_pcd_scale
        out_part_s = []
        out_part_m = []
        out_part_l = []
        torch.save(part_pcd.cpu(), f"/orion/u/w4756677/diffusion/anchorDIff/weights/test_{l[i]}.pt")
    exit()
    for j in range(n):
        _points = part_pcd[j].unsqueeze(0)
        metrics = EMD_CD(_points.expand(n, -1, -1), part_pcd, 32, reduced=False)
        close_percent_s = (metrics['MMD-CD'] < 0.02).sum()
        close_percent_m = (metrics['MMD-CD'] < 0.03).sum()
        close_percent_l = (metrics['MMD-CD'] < 0.04).sum()
        print(close_percent_s, close_percent_m, close_percent_l)
        if close_percent_s < 10:
            out_part_s.append(_points)
        if close_percent_m < 10:
            out_part_m.append(_points)
        if close_percent_l < 10:
            out_part_l.append(_points)
    print(len(out_part_s), len(out_part_m), len(out_part_l))
    out_part_s = torch.cat(out_part_s, dim=0)
    out_part_m = torch.cat(out_part_m, dim=0)
    out_part_l = torch.cat(out_part_l, dim=0)
    torch.save(out_part_s.cpu(), "/orion/u/w4756677/diffusion/anchorDIff/weights/pruned_seat_s.pt")
    torch.save(out_part_m.cpu(), "/orion/u/w4756677/diffusion/anchorDIff/weights/pruned_seat_m.pt")
    torch.save(out_part_l.cpu(), "/orion/u/w4756677/diffusion/anchorDIff/weights/pruned_seat_l.pt")
    
    
    # print(f'saved at /orion/u/w4756677/diffusion/anchorDIff/weights/pruned_part_{thresh}_part_{args.part_id}.pkl')
    # pickle.dump(dict(parts=out), open(f"/orion/u/w4756677/diffusion/anchorDIff/weights/pruned_part_{thresh}_part_{args.part_id}.pkl", "wb"))

tools/fps_shape.py
diffFacto/diffFacto
import argparse 
import torch
from anchor_diff.datasets.evaluation_utils import _pairwise_EMD_CD_ 
from anchor_diff.datasets.dataset_utils import load_ply_withNormals, shapenet_part_normal_cat_to_id
import numpy as np
from anchor_diff.utils.misc import fps
import pickle
import json
import os
from tqdm import tqdm


seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
                'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46],
                'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27],
                'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40],
                'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--cat',
        type=str) 
    args = parser.parse_args()
    cat_id = shapenet_part_normal_cat_to_id[args.cat]
    
    ref_part_pointcloud = [[], [], [], []]
    root = '/orion/u/w4756677/datasets/coalace_data_processed' if args.cat != 'Car' else '/orion/u/w4756677/datasets/shapenetcore_partanno_segmentation_benchmark_v0_normal'
    with open(f"{root}/train_test_split/shuffled_test_file_list.json", 'r') as f:
        test_ids = set([str(d.split('/')[2]) for d in json.load(f)])
    with open(os.path.join(root, 'train_test_split', 'shuffled_train_file_list.json'), 'r') as f:
        train_ids = set([str(d.split('/')[2]) for d in json.load(f)]) # '928c86eabc0be624c2bf2dc31ba1713' this is the first value
    with open(os.path.join(root, 'train_test_split', 'shuffled_val_file_list.json'), 'r') as f:
        val_ids = set([str(d.split('/')[2]) for d in json.load(f)])
    fns = sorted(os.listdir(f"{root}/{cat_id}"))  # Get each txt file under the folder according to the path and put it in the list
    if args.cat != 'Car':
        densePoints = [f for f in fns if (f[-7:] == 'ori.ply')]
        labels = [f for f in fns if (f[-10:] == 'ori.labels')]
        
        labels = [os.path.join(f"{root}/{cat_id}", fn) for fn in labels if ((fn[0:-11] in train_ids) or (fn[0:-11] in val_ids))]
        densePoints = [os.path.join(f"{root}/{cat_id}", fn) for fn in densePoints if ((fn[0:-8] in train_ids) or (fn[0:-8] in val_ids))]
        print(len(labels), len(densePoints))
    else:
        densePoints = [f for f in fns if (f[-3:] == 'txt')]
        densePoints = [os.path.join(f"{root}/{cat_id}", fn) for fn in densePoints if ((fn[0:-4] in test_ids) )]
        print(len(densePoints))
        
    for i in tqdm(range(len(densePoints))):
        if args.cat != 'Car':
            _ref = torch.from_numpy(load_ply_withNormals(densePoints[i]).astype(np.float32)).cuda()
            _ref_seg_mask = torch.from_numpy(np.loadtxt(labels[i]).astype(np.int32)).cuda() - 1
        else:
            data = torch.from_numpy(np.loadtxt(densePoints[i]).astype(np.float32)).cuda()
            _ref_seg_mask = data[:, -1].to(torch.int32) - seg_classes[args.cat][0] # Get the label of the sub category
            _ref = data[:, 0:3]
            
        for j in range(4):
            ref_idx = _ref_seg_mask == j
            if args.cat != "Car":
                ref_points = _ref[ref_idx].unsqueeze(0)
                if ref_points.abs().sum() != 0:
                    if ref_points.shape[1] > 2048:
                        ref_points = fps(ref_points, 2048)
                    ref_part_pointcloud[j].append(ref_points)
            else:
                if torch.any(ref_idx):
                    ref_points = _ref[ref_idx].unsqueeze(0)
                    if ref_points.shape[1] > 100:
                        while ref_points.shape[1] < 2048:
                            ref_points = torch.cat([ref_points, ref_points], dim=1)
                        if ref_points.shape[1] > 2048:
                            ref_points = fps(ref_points, 2048)
                        ref_part_pointcloud[j].append(ref_points)
            
    print(len(ref_part_pointcloud))
    ref_part_pointcloud =  [torch.cat(ll, dim=0) for ll in ref_part_pointcloud]
    for s in ref_part_pointcloud:
        print(s.shape)
    print('Normalizing------------------')
    for i in range(4):
        _ref = ref_part_pointcloud[i]
        _ref_shift = _ref.mean(1).reshape(-1, 1, 3)
        _ref_scale = _ref.std(1).reshape(-1, 1, 3)
        _ref = (_ref - _ref_shift) / _ref_scale
        ref_part_pointcloud[i] = _ref.detach().cpu().numpy()
    print(f'saved at /orion/u/w4756677/diffusion/anchorDIff/weights/ref_part_{args.cat}.pkl')
    pickle.dump(dict(ref=ref_part_pointcloud), open(f"/orion/u/w4756677/diffusion/anchorDIff/weights/ref_part_{args.cat}.pkl", "wb"))

tools/get_trained_parts.py
diffFacto/diffFacto
import torch 
from torch import nn
from torch.utils.data import DataLoader
from torch.optim import SGD, Adam
from torch.autograd import grad
import argparse
from tqdm import tqdm
from anchor_diff.config.config import init_cfg
from anchor_diff.config.config import get_cfg, save_cfg, save_args
from anchor_diff.utils.registry import build_from_cfg,MODELS,HOOKS
from anchor_diff.utils.misc import check_file, build_file, search_ckpt, sync, parse_losses
import pickle
import time
import datetime
from anchor_diff.utils import misc
from anchor_diff.datasets.shapenet_seg import _ShapeNetSegParts

class NoiseOptimizer:
    def __init__(self, device, args):
        cfg = get_cfg()

        self.logger = build_from_cfg(cfg.logger,HOOKS,work_dir=cfg.work_dir)

        self.local_rank = args.local_rank
        self.prefix=args.prefix

        if args.seed is not None:
            if self.logger is not None:
                self.logger.print_log(f'Set random seed to {args.seed}, 'f'deterministic: {args.deterministic}')
            misc.set_random_seed(args.seed + args.local_rank, deterministic=args.deterministic) # seed + rank, for augmentation

        self.cfg = cfg
        self.work_dir = cfg.work_dir
        self.save_num_batch = cfg.save_num_batch

        self.max_iter = cfg.max_iter

        self.resume_path = cfg.resume_path
        self.max_norm = cfg.max_norm
        self.model = build_from_cfg(cfg.model,MODELS)
        if device=='cuda':
            self.model.to(self.local_rank)

        
        


        self.dataset = _ShapeNetSegParts(**cfg.dataset)
        self.loader = DataLoader(self.dataset, batch_size=args.batch_size, shuffle=False, num_workers=0)
            
        # only save file in the 0th process
        if  self.local_rank == 0:
            save_file = build_file(self.work_dir,prefix="config.yaml")
            save_arg = build_file(self.work_dir,prefix="args.yaml")
            save_args(save_arg, args)
            save_cfg(save_file)


        self.iter = 0

        self.total_iter = self.max_iter

        if (cfg.pretrained_weights):
            self.load(cfg.pretrained_weights, model_only=cfg.model_only)
        
        if self.resume_path is None:
            self.resume_path = search_ckpt(self.work_dir)
        if check_file(self.resume_path):
            self.resume()


    
    def optimize(self, pcds, batch_idx, start_time):
        if self.logger is not None:
            self.logger.print_log("Start running")
        noise, _ = self.model.encoder.sample_noise(pcds, 'cuda', 1)
        z = nn.Parameter(noise)
        prev = torch.zeros(1).cuda()
        lr = 1
        optimizer = Adam([z], lr=lr)
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.5, patience=10, verbose=True, min_lr=5e-2)
        for it in range(self.max_iter):
            optimizer.zero_grad()
            losses = self.model.optimize_latent(pcds, z, device=self.local_rank)
            all_loss, losses = parse_losses(losses)
            # all_loss =  (z ** 2).sum(1)
            # print(grad(outputs=all_loss, inputs=z))
            all_loss.backward()
            # print(z.grad, z.requires_grad)
            optimizer.step()
            scheduler.step(all_loss)
            with torch.no_grad():
                ptime = time.time()-start_time
                eta_time = (self.max_iter-it)*ptime/ (it + 1)
                eta_str = str(datetime.timedelta(seconds=int(eta_time)))
                data = dict(
                    name = self.cfg.name,
                    lr = optimizer.param_groups[0]['lr'],
                    iter = it,
                    batch_idx = batch_idx,
                    total_loss = all_loss,
                    eta=eta_str,
                    z_mean = z.mean(),
                    
                )
                data.update(losses)
                data = sync(data)
                if self.logger is not None:
                    self.logger.log(data)
            if torch.allclose(all_loss, prev):
                break 
            prev = all_loss
        return z
                
    def run(self):

        self.model.eval()
        start_time = time.time()
        results = []
        for batch_idx, pcds in enumerate(self.loader):
            z = self.optimize(pcds, batch_idx, start_time)
            with torch.no_grad():
                _result = self.model.cimle_forward(pcds, noise=z, device=self.local_rank)
                _result['latent_code'] = z.detach().cpu()
                results.append(sync(_result, to_numpy=False))
        save_dict, eval_result = self.dataset.evaluate(results, self.save_num_batch, self.local_rank)
        save_file = build_file(self.work_dir, prefix=f"val/sample_{self.prefix}_cimle_optimized.pkl")
        self.logger.print_log(f"Saving to {save_file}....")
        pickle.dump(save_dict, open(save_file, 'wb'))
        eval_result = {k:v for k, v in eval_result.items()}
        self.logger.log(eval_result,iter=self.iter)
            
        
    
    def load(self, load_path, model_only=False):
        map_location = {'cuda:%d' % 0: 'cuda:%d' % self.local_rank}
        resume_data = torch.load(load_path, map_location=map_location)

        if (not model_only):
            meta = resume_data.get("meta",dict())
            self.epoch = meta.get("epoch",self.epoch)
            self.iter = meta.get("iter",self.iter)
            self.max_iter = meta.get("max_iter",self.max_iter)
            self.max_epoch = meta.get("max_epoch",self.max_epoch)
            self.scheduler.load_state_dict(resume_data.get("scheduler",dict()))
            # self.optimizer.load_state_dict(resume_data.get("optimizer",dict()))
        base_ckpt = {k.replace("module.", ""): v for k, v in resume_data['model'].items()}
        self.model.load_state_dict(base_ckpt,strict=True)

        self.logger.print_log(f"Loading model parameters from {load_path}")

    def resume(self):
        self.load(self.resume_path, model_only=self.cfg.model_only)


def main():
    parser = argparse.ArgumentParser(description="Jittor Object Detection Training")

    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch'],
        default='none',
        help='job launcher') 
    parser.add_argument('--local_rank', type=int, default=0)

    parser.add_argument('--seed', type=int, default=0, help='random seed')

    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')

    parser.add_argument(
        "--config-file",
        default="",
        metavar="FILE",
        help="path to config file",
        type=str,
    )

    parser.add_argument(
        "--prefix",
        default="exp",
        type=str
    )

    parser.add_argument(
        "--save_dir",
        default=".",
        type=str,
    )
    parser.add_argument(
        "--batch_size",
        default=32,
        type=int,
    )
    
    args = parser.parse_args()

    device='cuda'

    if args.config_file:
        init_cfg(args.config_file)

    if device == 'cuda':
        torch.backends.cudnn.benchmark = True
    # init distributed env first, since logger depends on the dist info.
    args.world_size = 1
    

    runner = NoiseOptimizer(device, args)

    runner.run()


if __name__ == "__main__":
    main()
tools/optimize_noise.py
diffFacto/diffFacto
import argparse 
import torch
from anchor_diff.datasets.evaluation_utils import _pairwise_EMD_CD_ 
from anchor_diff.datasets.dataset_utils import load_ply_withNormals, shapenet_part_normal_cat_to_id
import numpy as np
from anchor_diff.utils.misc import fps
import pickle
import json
import os
from tqdm import tqdm
from pointnet2_ops import pointnet2_utils
import h5py

seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
                'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46],
                'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27],
                'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40],
                'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--datadir',
        type=str) 
    parser.add_argument(
        '--cat',
        type=str) 
    args = parser.parse_args()
    
    data = h5py.File("/orion/u/w4756677/datasets/partglot_data/subsampled_partglot_data_triplet.h5", 'r')
    distractors = data['distractor_pcds']
    targets = data['target_pcds']
    print(distractors.shape, targets.shape)
    print(data.keys())
    exit()
    
    # import torch
    # data1 = torch.load("/orion/u/w4756677/diffusion/anchorDIff/weights/chair-separate-v2.pt")
    # data2 = torch.load("/orion/u/w4756677/diffusion/anchorDIff/weights/chair-separate.pt")
    # ratio = [700*4, 700*4, 690*4, 140*4]
    # data = torch.cat([data1, data2], dim=0)[:2800]
    # print(data.shape)
    # out_data = []
    # out_mask = []
    # for i in tqdm(range(data.shape[0])):
    #     _data = data[i][..., :3].cuda()
    #     _mask = data[i][..., 3].cuda()
    #     if i > ratio[3]:
    #         wo_arm = _mask != 3 
    #         _data = _data[wo_arm]
    #         _mask = _mask[wo_arm]
    #     if i > ratio[2]:
    #         wo_leg = _mask != 2
    #         _data = _data[wo_leg]
    #         _mask = _mask[wo_leg]
    #     _data, ids = fps(_data[None], 2048, True)
    #     _mask = pointnet2_utils.gather_operation(_mask.reshape(1, 1, -1).contiguous(), ids).reshape(2048)
    #     print(_data.shape, _mask.shape)
    #     out_data.append(_data[0].cpu().numpy())
    #     out_mask.append(_mask.cpu().numpy())
    # out_data = np.stack(out_data, axis=0)
    # out_mask = np.stack(out_mask, axis=0)
    # pickle.load(open())
    # print(out_data.shape, out_mask.shape)
    # pickle.dump(dict(pred=out_data, pred_seg_mask=out_mask), open("/orion/u/w4756677/diffusion/anchorDIff/weights/lion_part/lion_2800_chair.pkl", "wb"))
    # exit()
    # data1 = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/weights/lion_Chair_2048.pkl", 'rb'))
    # print(data.keys(), data1.keys())
    # n = data1['pred'].shape[0]
    # pickle.dump( dict(pred=np.concatenate([data1['pred'], data['pred'][:2800 - n]], axis=0), pred_seg_mask=np.concatenate([data1['pred_seg_mask'], data['pred_seg_mask'][:2800 - n]], axis=0)), open(f"/orion/u/w4756677/diffusion/anchorDIff/weights/lion_2800_chair.pkl", 'wb'))
    # exit()
    # data_out = []
    # for i in range(704):
    #     _data = data[i]
    #     perm = torch.randperm(8192)
    #     idx = perm[:2048]
    #     data_out.append(_data[idx])
    # data_out = torch.stack(data_out, dim=0)
    # print(data_out.shape)
    # pickle.dump(dict(pred=data_out[..., :3].numpy().astype(float), pred_seg_mask=data_out[..., 3].numpy().astype(int)), open("/orion/u/w4756677/diffusion/anchorDIff/weights/lion_chair_v2.pkl", 'wb'))
    # exit()
    # print(data['parts'].shape)
    # data = data['parts'].cpu().numpy()
    # np.save("/orion/u/w4756677/diffusion/anchorDIff/weights/pruned_part_2000_part_2.npy", data)
    # num = 2800 if args.cat == "Chair" else 1360
    seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
                            'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46],
                            'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27],
                            'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40],
                            'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}
    # dir_to_ds = "/orion/u/w4756677/datasets/colasce_data_txt/02691156"
    # dirs = os.listdir(dir_to_ds)
    # parts = [[], [], [], []]
    # parts_fns = [[], [], [], []]
    # permed_parts = []
    # permed_parts_fns = []
    # for d in tqdm(dirs):
    #     print(d[:-4])        
    #     out = np.loadtxt(os.path.join(dir_to_ds, d)).astype(np.float32)
    #     # out[..., 3] -= 12
    #     mask = out[..., 3]
    #     for k in range(4):
    #         part_mask = mask == k
    #         if np.any(part_mask):
    #             part_pts = out[part_mask, :3]
    #             parts[k].append(part_pts)
    #             parts_fns[k].append(d[:-4])
    # for k in range(4):
    #     assert (len(parts[k]) == len(parts_fns[k]))
    #     print(len(parts[k]) )
    #     parts[k] = np.stack(parts[k], axis=0)
    # for k in range(4):
    #     part = parts[k]
    #     part_fn = parts_fns[k]
    #     l = part.shape[0]
    #     perm = np.random.permutation(np.arange(l))
    #     part_fn = [part_fn[i] for i in np.arange(l).tolist()]
    #     part = part[perm]
    #     assert part.shape[0] == l
    #     if l > 2000:
    #         part = part[:2000]
    #         part_fn = part_fn[:2000]
    #     permed_parts.append(part)
    #     permed_parts_fns.append(part_fn)
    # for i in range(4):
    #     print(permed_parts[i].shape, len(permed_parts_fns[i]))
    # pickle.dump(dict(parts=permed_parts, permed_parts_fns=permed_parts_fns), open("/orion/u/w4756677/diffusion/anchorDIff/weights/ref_part/permuted_parts_airplane.pkl", "wb"))
    # exit()
    part0 = np.load("/orion/u/w4756677/diffusion/anchorDIff/weights/sgf_part/mixing_chair_two/mixing0.npy")
    part1 = np.load("/orion/u/w4756677/diffusion/anchorDIff/weights/sgf_part/mixing_chair_two/mixing1.npy")
    part2 = np.load("/orion/u/w4756677/diffusion/anchorDIff/weights/sgf_part/mixing_chair_two/mixing2.npy")
    part3 = np.load("/orion/u/w4756677/diffusion/anchorDIff/weights/sgf_part/mixing_chair_two/mixing3.npy")
    mask = np.array([0, 1, 2, 3]).repeat(2048, axis=0)
    allpart = np.concatenate([part0, part1, part2, part3], axis=0).reshape(-1, 3)
    print(mask.shape, allpart.shape)
    np.save("/orion/u/w4756677/diffusion/anchorDIff/weights/sgf_part/mixing_chair_two/total.npy", np.concatenate([allpart, mask[...,None]], axis=-1))
    exit()
    from anchor_diff.utils.misc import fps
    from pointnet2_ops import pointnet2_utils
    sgf_leg_interp = torch.from_numpy(np.load("/orion/u/w4756677/diffusion/anchorDIff/weights/interpolation/sgf_leg_interp.npy"))
    sgf_leg_back = torch.from_numpy(np.load("/orion/u/w4756677/diffusion/anchorDIff/weights/interpolation/sgf_back_for_leg_interp.npy"))
    sgf_leg_seat = torch.from_numpy(np.load("/orion/u/w4756677/diffusion/anchorDIff/weights/interpolation/sgf_seat_for_leg_interp.npy"))
    sgf_leg_arm = torch.from_numpy(np.load("/orion/u/w4756677/diffusion/anchorDIff/weights/interpolation/sgf_arm_for_leg_interp.npy"))
    data = pickle.load(open("/orion/u/w4756677/diffusion/anchorDIff/weights/interpolation/leg_set.pkl", "rb"))
    set1 = torch.from_numpy(data['from']).cuda().float()
    out = []
    mask = []
    arm_id = 0
    for i in tqdm(range(set1.shape[0])):
        _out = []
        m = []
        _data1 = set1[i]
        pts1 = _data1[:, :3]
        mask1 = _data1[:, 3]
        em1 = mask1 == 3
        _out.append(sgf_leg_back[i])
        _out.append(sgf_leg_seat[i])
        m = [torch.zeros([2048]), torch.ones([2048])]
        if torch.any(em1): 
            _out.append(sgf_leg_arm[arm_id])
            arm_id += 1
            m.append(torch.ones([2048]) * 3)
        _out = torch.cat(_out, dim=1).repeat_interleave(10, dim=0)
        print(_out.shape, sgf_leg_interp.shape)
        _out = torch.cat([_out, sgf_leg_interp[i]], dim=1)
        # m.append(2)
        m.append(torch.ones([2048]) * 2)
        _mask = torch.cat(m).reshape(1, -1).repeat_interleave(10, dim=0)
        print(m)
        print(_mask.shape, _out.shape)
        print(_out.shape)
        if _out.shape[1] > 2048:
            _out, idx = fps(_out.cuda(), 2048, True)
            print(_out.shape)
            _mask = pointnet2_utils.gather_operation(_mask.reshape(10, 1, -1).contiguous().float().cuda(), idx.int()).view(10, 2048, 1)
            print(_out.shape, _mask.shape)
        _out = np.concatenate([_out.cpu().numpy(), _mask.cpu().numpy()], axis=-1)
        out.append(_out)
    out = np.stack(out, axis=0)
    print(out.shape)
    np.save("/orion/u/w4756677/diffusion/anchorDIff/weights/interpolation/sgf_interp.npy", out)
    exit()
    
    for p in path:
        data = np.loadtxt(fn[1]).astype(np.float32) # size 20488,7 read in this txt file, a total of 20488 points, each point xyz rgb + sub category label
        token = (os.path.splitext(os.path.basename(fn[1]))[0])
        data.append()
        
    data = pickle.load(open(args.datadir, 'rb'))
    
    # print(data.keys())
    pred = data['ref']
    # print(len(pred))
    # print(pred.shape)
    # del data['pred']
    mask = data['ref_seg_mask']
    out = np.concatenate([pred, mask[..., None]], axis=-1)
    # del data['pred_seg_mask']
    perm_out = np.random.permutation(out)
    print(out.shape, perm_out.shape)
    pickle.dump(dict(set1=out, set2=perm_out), open("/orion/u/w4756677/diffusion/anchorDIff/weights/interpolation/interpolation_set_airplanes.pkl", 'wb'))
    # print(pred.shape)
    exit()
    # pickle.dump( dict(pred=pred[:2800], pred_seg_mask=mask[:2800]), open(f"/orion/u/w4756677/diffusion/anchorDIff/weights/v2_dpm_{num}_{args.cat}.pkl", 'wb'))
    # exit()
    preds = [v for k, v in data.items() if 'pred' in k]
    # print(preds)
    print(len(preds))
    N = 704 if args.cat == 'Chair' else 340
    # factors = [(N, 1), (N, 2), (N, 3), (N, 4), (N, 5)]
    factors = [(N, 4)]
    # pred = torch.stack([-pred[..., 2], pred[..., 1], pred[..., 0]], dim=-1)
    
    for size, m in factors:
        _preds = [preds[p] for p in range(m)]
        _preds = [p[:size] for p in _preds]
        _mask = mask[None, :size].repeat(m, axis=0).reshape(-1, 2048)
        _preds = np.concatenate(_preds, axis=0)
        print(_preds.shape)
        pickle.dump( dict(pred=_preds, seg_mask_ref=_mask), open(f"/orion/u/w4756677/diffusion/anchorDIff/weights/624_n100_{size}_{m}.pkl", 'wb'))
                        

tools/playground.py
diffFacto/diffFacto
import torch 
from pointnet2_ops.pointnet2_utils import furthest_point_sample, gather_operation
import numpy as np
import pickle
from tqdm import tqdm
import os.path as osp 
import h5py
from data_utils import *

def preprocess():
    path = "/mnt/disk/wang/diffusion/PartGlot/logs/pre_trained/pn_agnostic/10-22_02-28-36/pred_label/final/out.pkl"
    data = pickle.load(open(path, 'rb'))
    segs_data = torch.from_numpy(data["geos"].astype(np.float32)).to('cuda')
    segs_mask = torch.from_numpy(data["geos_masks"].astype(np.float32)).to('cuda')
    attn_maps = torch.from_numpy(data['attn_maps'].astype(np.float32)).to('cuda')
    out_segs_data = []
    out_attn_maps = []
    for i in tqdm(range(segs_data.shape[0])):
        seg_data, mask, attn_map = segs_data[i], segs_mask[i], attn_maps[i]
        used_points, used_attn_map = seg_data[mask == 1].reshape(-1, 3), attn_map[:, mask==1].unsqueeze(2).expand(-1, -1, 512).reshape(4, -1).transpose(1,0)
        if used_points.shape[0] > 2048:
            idx = furthest_point_sample(used_points.unsqueeze(0), 2048)
            used_points = gather_operation(used_points.unsqueeze(0).transpose(1,2).contiguous(), idx).squeeze(0).transpose(1,0)
            used_attn_map = gather_operation(used_attn_map.unsqueeze(0).transpose(1,2).contiguous(), idx).squeeze(0).transpose(1,0)
        if used_points.shape[0] < 2048:
            continue
        
        out_segs_data.append(used_points)
        out_attn_maps.append(used_attn_map)
    out_segs_data = torch.stack(out_segs_data, dim=0)
    out_attn_maps = torch.stack(out_attn_maps, dim=0)
    print(out_segs_data.shape, out_attn_maps.shape)
    pickle.dump(out_segs_data.cpu().numpy(), open("/mnt/disk/wang/diffusion/datasets/partglot_data/shapenet_pointcloud_pn_agnostic.pkl", 'wb'))
    pickle.dump(out_attn_maps.cpu().numpy(), open("/mnt/disk/wang/diffusion/datasets/partglot_data/shapenet_label_pn_agnostic.pkl", "wb"))

def pre_process_partglot_data():
    (game_data,
    word2int,
    int2word,
    int2sn,
    sn2int,
    sorted_sn
    ) = unpickle_data(osp.join("/mnt/disk3/wang/diffusion/datasets/partglot_data", "game_data.pkl"))
    gt_attn_maps = pickle.load(open("/mnt/disk3/wang/diffusion/PartGlot/logs/pre_trained/pn_aware/11-01_10-46-46/pred_label/final/ground_truth_attnmaps.pkl", "rb"))['gt_attn_maps']
    distractor_attn_maps = pickle.load(open("/mnt/disk3/wang/diffusion/PartGlot/logs/pre_trained/pn_aware/11-01_11-06-43/pred_label/final/distractor_attnmaps.pkl", "rb"))['gt_attn_maps']
    h5_data = h5py.File(osp.join("/mnt/disk3/wang/diffusion/datasets/partglot_data", "cic_bsp.h5"), "r")
    segs_data = h5_data["data"][:].astype(np.float32)
    print(type(segs_data))
    segs_mask = h5_data["mask"][:].astype(np.float32)
    padded_text, seq_len = pad_text_symbols_with_zeros(
            game_data["text"], 33, force_zero_end=True
        )
    labels = convert_labels_to_one_hot(game_data["target_chair"])
    geo_ids = np.array(game_data[["chair_a", "chair_b", "chair_c"]], dtype=np.int32)
    mask, part_indicator = get_mask_of_game_data(
            game_data,
            word2int,
            True,
            True,
            33,
            True,
        )
    masked_geo_ids = geo_ids[mask]
    masked_labels = labels[mask]
    masked_padded_text = padded_text[mask]
    masked_text = game_data["text"][mask]
    masked_part_indicator = part_indicator[mask]
    l = len(masked_geo_ids)
    print(l)
    out_gt = []
    out_distractor = []
    out_gt_mask = []
    out_distractor_mask = []
    out_texts = []
    out_part_indicator = []
    for idx in tqdm(range(l)):
        geo_id = torch.from_numpy(masked_geo_ids[idx])
        target = torch.tensor(np.argmax(masked_labels[idx]))
        text = [torch.from_numpy(masked_padded_text[idx])] * 2
        part_ind = [torch.from_numpy(masked_part_indicator[idx])] * 2

        geos = torch.from_numpy(segs_data[geo_id])
        geos_mask = torch.from_numpy(segs_mask[geo_id])
        gt = [geos[target], geos[target]]
        gt_mask = [geos_mask[target]] * 2
        distractor_mask = [geos_mask[i] for i in range(3) if i != target]
        distractors = [geos[i] for i in range(3) if i != target]
        out_gt += gt
        out_gt_mask += gt_mask
        out_distractor += distractors
        out_distractor_mask += distractor_mask
        out_texts += text
        out_part_indicator += part_ind
    
    out_gt, out_distractor, out_gt_mask, out_distractor_mask, out_texts, out_part_indicator = map(lambda p: torch.stack(p, dim=0).numpy(), [out_gt, out_distractor, out_gt_mask, out_distractor_mask, out_texts, out_part_indicator])
    print(out_gt.shape, out_distractor.shape, out_gt_mask.shape, out_distractor_mask.shape, out_texts.shape, out_part_indicator.shape)
    pickle.dump(dict(gt_geos=out_gt, distractor_geos=out_distractor, gt_mask=out_gt_mask, distractor_mask=out_distractor_mask, texts=out_texts, part_indicator=out_part_indicator), open("/mnt/disk3/wang/diffusion/datasets/partglot_data/processed_partglot_data.pkl", "wb"))

def pre_process_partglot_data_triplet():
    (game_data,
    word2int,
    int2word,
    int2sn,
    sn2int,
    sorted_sn
    ) = unpickle_data(osp.join("/mnt/disk3/wang/diffusion/datasets/partglot_data", "game_data.pkl"))
    gt_attn_maps = pickle.load(open("/mnt/disk3/wang/diffusion/PartGlot/logs/pre_trained/pn_aware/11-01_10-46-46/pred_label/final/ground_truth_attnmaps.pkl", "rb"))['gt_attn_maps']
    distractor_attn_maps = pickle.load(open("/mnt/disk3/wang/diffusion/PartGlot/logs/pre_trained/pn_aware/11-01_11-06-43/pred_label/final/distractor_attnmaps.pkl", "rb"))['gt_attn_maps']
    h5_data = h5py.File(osp.join("/mnt/disk3/wang/diffusion/datasets/partglot_data", "cic_bsp.h5"), "r")
    segs_data = h5_data["data"][:].astype(np.float32)
    print(type(segs_data))
    segs_mask = h5_data["mask"][:].astype(np.float32)
    padded_text, seq_len = pad_text_symbols_with_zeros(
            game_data["text"], 33, force_zero_end=True
        )
    labels = convert_labels_to_one_hot(game_data["target_chair"])
    geo_ids = np.array(game_data[["chair_a", "chair_b", "chair_c"]], dtype=np.int32)
    mask, part_indicator = get_mask_of_game_data(
            game_data,
            word2int,
            True,
            True,
            33,
            True,
        )
    masked_geo_ids = geo_ids[mask]
    masked_labels = labels[mask]
    masked_padded_text = padded_text[mask]
    masked_text = game_data["text"][mask]
    masked_part_indicator = part_indicator[mask]
    l = len(masked_geo_ids)
    print(l)
    out_gt = []
    out_distractor = []
    out_gt_mask = []
    out_distractor_mask = []
    out_texts = []
    out_part_indicator = []
    for idx in tqdm(range(l)):
        geo_id = torch.from_numpy(masked_geo_ids[idx])
        target = torch.tensor(np.argmax(masked_labels[idx]))
        text = torch.from_numpy(masked_padded_text[idx])
        part_ind = torch.from_numpy(masked_part_indicator[idx])
        used_gt_points, used_gt_attn_map = gt[gt_mask == 1].reshape(-1, 3), gt_attn_map[:, gt_mask==1].unsqueeze(2).expand(-1, -1, 512).reshape(4, -1).transpose(1,0)
        used_distractor_points, used_distractor_attn_map = distractor[distractor_mask == 1].reshape(-1, 3), distractor_attn_map[:, distractor_mask==1].unsqueeze(2).expand(-1, -1, 512).reshape(4, -1).transpose(1,0)
        

        geos = torch.from_numpy(segs_data[geo_id])
        geos_mask = torch.from_numpy(segs_mask[geo_id])
        gt = geos[target]
        gt_mask = geos_mask[target]
        distractor_mask = [geos_mask[i] for i in range(3) if i != target]
        distractors = [geos[i] for i in range(3) if i != target]
        out_gt += gt
        out_gt_mask += gt_mask
        out_distractor += distractors
        out_distractor_mask += distractor_mask
        out_texts += text
        out_part_indicator += part_ind
    
    out_gt, out_distractor, out_gt_mask, out_distractor_mask, out_texts, out_part_indicator = map(lambda p: torch.stack(p, dim=0).numpy(), [out_gt, out_distractor, out_gt_mask, out_distractor_mask, out_texts, out_part_indicator])
    print(out_gt.shape, out_distractor.shape, out_gt_mask.shape, out_distractor_mask.shape, out_texts.shape, out_part_indicator.shape)
    pickle.dump(dict(gt_geos=out_gt, distractor_geos=out_distractor, gt_mask=out_gt_mask, distractor_mask=out_distractor_mask, texts=out_texts, part_indicator=out_part_indicator), open("/mnt/disk3/wang/diffusion/datasets/partglot_data/processed_partglot_data.pkl", "wb"))


def process_partglot_data():
    out = pickle.load(open("/mnt/disk3/wang/diffusion/datasets/partglot_data/processed_partglot_data.pkl", "rb"))
    gt_attn_maps = pickle.load(open("/mnt/disk3/wang/diffusion/PartGlot/logs/pre_trained/pn_aware/11-01_10-46-46/pred_label/final/ground_truth_attnmaps.pkl", "rb"))['gt_attn_maps']
    distractor_attn_maps = pickle.load(open("/mnt/disk3/wang/diffusion/PartGlot/logs/pre_trained/pn_aware/11-01_11-06-43/pred_label/final/distractor_attnmaps.pkl", "rb"))['gt_attn_maps']
    gt = out['gt_geos']
    distractor = out['distractor_geos'] 
    gt_mask = out['gt_mask']
    distractor_mask = out['distractor_mask']
    texts = out['texts']
    part_indicator = out['part_indicator']
    print(gt.shape, distractor.shape, gt_mask.shape, distractor_mask.shape, texts.shape, part_indicator.shape, gt_attn_maps.shape, distractor_attn_maps.shape)
    hf = h5py.File('/mnt/disk3/wang/diffusion/datasets/partglot_data/processed_partglot_data.h5', 'w')
    hf.create_dataset("target_geos", data=gt)
    hf.create_dataset("distractor_geos", data=distractor)
    hf.create_dataset("target_geo_masks", data=gt_mask)
    hf.create_dataset("distractor_geo_masks", data=distractor_mask)
    hf.create_dataset("texts", data=texts)
    hf.create_dataset("part_indicator", data=part_indicator)
    hf.create_dataset("target_attn_maps", data=gt_attn_maps)
    hf.create_dataset("distractor_attn_maps", data=distractor_attn_maps)
    hf.close()

def subsample_partglot_data():
    data = h5py.File("/mnt/disk3/wang/diffusion/datasets/partglot_data/processed_partglot_data.h5", "r")
    
    gts = data['target_geos'][:]
    l = gts.shape[0]
    assert l % 2 == 0
    B = l //2
    print(B)
    gt_attn_maps = data['target_attn_maps'][:].reshape(B, 2, 4, 50)
    distractor_attn_maps = data['distractor_attn_maps'][:].reshape(B, 2, 4, 50)
    gts = gts.reshape(B, 2, 50, 512, 3)
    distractors = data['distractor_geos'][:].reshape(B, 2, 50, 512, 3)
    gt_masks = data['target_geo_masks'][:].reshape(B, 2, 50)
    distractor_masks = data['distractor_geo_masks'][:].reshape(B, 2, 50)
    texts = data['texts'][:].reshape(B, 2, -1)
    print(texts.shape)
    part_indicators = data['part_indicator'][:].reshape(B, 2, 4)
    
    subsampled_part_indicators = []
    subsampled_texts = []
    subsampled_gts = []
    subsampled_distractors = []
    subsampled_gt_attn_maps = []
    subsampled_distractor_attn_maps = []
    for i in tqdm(range(B)):
        gt = torch.from_numpy(gts[i]).cuda()
        assert torch.all(gt[0] == gt[1])
        gt = gt[0]
        distractor0, distractor1 = torch.split(torch.from_numpy(distractors[i]).cuda(), 1, dim=0)
        distractor0, distractor1 = distractor0.squeeze(0), distractor1.squeeze(0)
        gt_mask = torch.from_numpy(gt_masks[i, 0]).cuda()
        distractor0_mask, distractor1_mask =  torch.split(torch.from_numpy(distractor_masks[i]).cuda(), 1, dim=0)
        distractor0_mask, distractor1_mask = distractor0_mask.squeeze(0), distractor1_mask.squeeze(0)
        gt_attn_map = torch.from_numpy(gt_attn_maps[i, 0]).cuda()
        distractor0_attn_map, distractor1_attn_map = torch.split(torch.from_numpy(distractor_attn_maps[i]).cuda(), 1, dim=0)
        distractor0_attn_map, distractor1_attn_map = distractor0_attn_map.squeeze(0), distractor1_attn_map.squeeze(0)
        used_gt_points, used_gt_attn_map = gt[gt_mask == 1].reshape(-1, 3), gt_attn_map[:, gt_mask==1].unsqueeze(2).expand(-1, -1, 512).reshape(4, -1).transpose(1,0)
        used_distractor0_points, used_distractor0_attn_map = distractor0[distractor0_mask == 1].reshape(-1, 3), distractor0_attn_map[:, distractor0_mask==1].unsqueeze(2).expand(-1, -1, 512).reshape(4, -1).transpose(1,0)
        used_distractor1_points, used_distractor1_attn_map = distractor1[distractor1_mask == 1].reshape(-1, 3), distractor1_attn_map[:, distractor1_mask==1].unsqueeze(2).expand(-1, -1, 512).reshape(4, -1).transpose(1,0)
        
        if used_gt_points.shape[0] > 2048:
            idx = furthest_point_sample(used_gt_points.unsqueeze(0), 2048)
            used_gt_points = gather_operation(used_gt_points.unsqueeze(0).transpose(1,2).contiguous(), idx).squeeze(0).transpose(1,0)
            used_gt_attn_map = gather_operation(used_gt_attn_map.unsqueeze(0).transpose(1,2).contiguous(), idx).squeeze(0).transpose(1,0)
        elif used_gt_points.shape[0] < 2048:
            continue
        
        if used_distractor0_points.shape[0] > 2048:
            idxx = furthest_point_sample(used_distractor0_points.unsqueeze(0), 2048)
            used_distractor0_points = gather_operation(used_distractor0_points.unsqueeze(0).transpose(1,2).contiguous(), idxx).squeeze(0).transpose(1,0)
            used_distractor0_attn_map = gather_operation(used_distractor0_attn_map.unsqueeze(0).transpose(1,2).contiguous(), idxx).squeeze(0).transpose(1,0)
        elif used_distractor0_points.shape[0] < 2048:
            continue
        
        if used_distractor1_points.shape[0] > 2048:
            idxxx = furthest_point_sample(used_distractor1_points.unsqueeze(0), 2048)
            used_distractor1_points = gather_operation(used_distractor1_points.unsqueeze(0).transpose(1,2).contiguous(), idxxx).squeeze(0).transpose(1,0)
            used_distractor1_attn_map = gather_operation(used_distractor1_attn_map.unsqueeze(0).transpose(1,2).contiguous(), idxxx).squeeze(0).transpose(1,0)
        elif used_distractor1_points.shape[0] < 2048:
            continue

        subsampled_gts.append(used_gt_points.cpu().numpy())
        subsampled_gt_attn_maps.append(used_gt_attn_map.cpu().numpy())
        subsampled_distractors.append(torch.stack([used_distractor0_points, used_distractor1_points], dim=0).cpu().numpy())
        subsampled_distractor_attn_maps.append(torch.stack([used_distractor0_attn_map, used_distractor1_attn_map], dim=0).cpu().numpy())   
        subsampled_part_indicators.append(part_indicators[i, 0])
        subsampled_texts.append(texts[i, 0])         

    subsampled_gts = np.stack(subsampled_gts, axis=0)
    subsampled_gt_attn_maps = np.stack(subsampled_gt_attn_maps, axis=0)
    subsampled_distractors = np.stack(subsampled_distractors, axis=0)
    subsampled_distractor_attn_maps = np.stack(subsampled_distractor_attn_maps, axis=0)
    subsampled_part_indicators = np.stack(subsampled_part_indicators, axis=0)
    subsampled_texts = np.stack(subsampled_texts, axis=0)
    
    print(subsampled_gts.shape, subsampled_distractors.shape, subsampled_texts.shape, subsampled_part_indicators.shape, subsampled_gt_attn_maps.shape, subsampled_distractor_attn_maps.shape)
    
    hf = h5py.File('/mnt/disk3/wang/diffusion/datasets/partglot_data/subsampled_partglot_data_triplet.h5', 'w')
    hf.create_dataset("target_pcds", data=subsampled_gts)
    hf.create_dataset("distractor_pcds", data=subsampled_distractors)
    hf.create_dataset("texts", data=texts)
    hf.create_dataset("part_indicators", data=subsampled_part_indicators)
    hf.create_dataset("target_attn_maps", data=subsampled_gt_attn_maps)
    hf.create_dataset("distractor_attn_maps", data=subsampled_distractor_attn_maps)
    hf.close()


# preprocess()


subsample_partglot_data()

tools/process_data.py
diffFacto/diffFacto
import argparse 
import torch
from anchor_diff.datasets.evaluation_utils import _pairwise_EMD_CD_ 
from anchor_diff.datasets.dataset_utils import load_ply_withNormals, shapenet_part_normal_cat_to_id
import numpy as np
from anchor_diff.utils.misc import fps
import pickle
import json
import os
from tqdm import tqdm


seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
                'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46],
                'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27],
                'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40],
                'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--datadir',
        type=str) 
    parser.add_argument(
        '--cat',
        type=str) 
    parser.add_argument(
        '--canonical',
        action='store_true'
    )
    parser.add_argument(
        '--npoint',
        type=int,
        default=2048
    )
    args = parser.parse_args()
    data = dict(np.load(args.datadir))
    cat_id = shapenet_part_normal_cat_to_id[args.cat]
    print(data.keys())
    assert 'pred' in data.keys() and ('pred_seg_mask' in data.keys() or 'seg_mask_ref' in data.keys())
    
    pred = torch.from_numpy(data['pred']).cuda()
    # pred = torch.stack([-pred[..., 2], pred[..., 1], pred[..., 0]], dim=-1)
    pred_seg_mask = data.get('pred_seg_mask', None)
    if pred_seg_mask is None:
        pred_seg_mask = data['seg_mask_ref']
    pred_seg_mask = torch.from_numpy(pred_seg_mask).cuda()
    size = pred.shape[0]
    pred_part_pointcloud = [[], [], [], []]
    pred_part_mask = [[], [], [], []]
    ref_part_pointcloud = [[], [], [], []]
    root = '/orion/u/w4756677/datasets/coalace_data_processed' if args.cat != 'Car' else '/orion/u/w4756677/datasets/shapenetcore_partanno_segmentation_benchmark_v0_normal'
    with open(f"{root}/train_test_split/shuffled_test_file_list.json", 'r') as f:
        test_ids = set([str(d.split('/')[2]) for d in json.load(f)])
    
    fns = sorted(os.listdir(f"{root}/{cat_id}"))  # Get each txt file under the folder according to the path and put it in the list
    if args.cat != 'Car':
        densePoints = [f for f in fns if (f[-7:] == 'ori.ply')]
        labels = [f for f in fns if (f[-10:] == 'ori.labels')]
        
        labels = [os.path.join(f"{root}/{cat_id}", fn) for fn in labels if ((fn[0:-11] in test_ids))]
        densePoints = [os.path.join(f"{root}/{cat_id}", fn) for fn in densePoints if ((fn[0:-8] in test_ids) )]
        print(len(labels), len(densePoints))
    else:
        densePoints = [f for f in fns if (f[-3:] == 'txt')]
        densePoints = [os.path.join(f"{root}/{cat_id}", fn) for fn in densePoints if ((fn[0:-4] in test_ids) )]
        print(len(densePoints))
        
    print(size)
    for i in tqdm(range(size)):
        _pred = pred[i]
        _pred_seg_mask = pred_seg_mask[i]
        for j in range(4):
            idx = _pred_seg_mask == j 
            if torch.any(idx):
                points = _pred[idx].unsqueeze(0)
                if points.shape[1] > 100:
                    mask = torch.ones([1, args.npoint])
                    if points.shape[1] < args.npoint:
                        mask[:, points.shape[1]:] = 0
                    while points.shape[1] < args.npoint:
                        points = torch.cat([points, points], dim=1)
                    if points.shape[1] > args.npoint:
                        points = points[:, :args.npoint]
                    pred_part_pointcloud[j].append(points)
                    assert (mask.shape == (1, args.npoint))
                    pred_part_mask[j].append(mask)
    for i in tqdm(range(len(densePoints))):
        if args.cat != 'Car':
            _ref = torch.from_numpy(load_ply_withNormals(densePoints[i]).astype(np.float32)).cuda()
            _ref_seg_mask = torch.from_numpy(np.loadtxt(labels[i]).astype(np.int32)).cuda() - 1
        else:
            data = torch.from_numpy(np.loadtxt(densePoints[i]).astype(np.float32)).cuda()
            _ref_seg_mask = data[:, -1].to(torch.int32) - seg_classes[args.cat][0] # Get the label of the sub category
            _ref = data[:, 0:3]
            
        for j in range(4):
            ref_idx = _ref_seg_mask == j
            if args.cat != "Car":
                ref_points = _ref[ref_idx].unsqueeze(0)
                if ref_points.abs().sum() != 0:
                    if ref_points.shape[1] > args.npoint:
                        ref_points = fps(ref_points, args.npoint)
                    ref_part_pointcloud[j].append(ref_points)
            else:
                if torch.any(ref_idx):
                    ref_points = _ref[ref_idx].unsqueeze(0)
                    if ref_points.shape[1] > 100:
                        while ref_points.shape[1] < args.npoint:
                            ref_points = torch.cat([ref_points, ref_points], dim=1)
                        if ref_points.shape[1] > args.npoint:
                            ref_points = fps(ref_points, args.npoint)
                        ref_part_pointcloud[j].append(ref_points)
            
    print(len(pred_part_pointcloud), len(ref_part_pointcloud), len(pred_part_mask))
    pred_part_pointcloud, ref_part_pointcloud, pred_part_mask = map(lambda l: [torch.cat(ll, dim=0) for ll in l], [pred_part_pointcloud, ref_part_pointcloud, pred_part_mask])
    for l, s, k in zip(pred_part_pointcloud, ref_part_pointcloud, pred_part_mask):
        print(s.shape, l.shape, k.shape)
    print('Normalizing------------------')
    if args.canonical:
        for i in range(4):
            _pred = pred_part_pointcloud[i]
            _pred_max = _pred.max(1)[0].reshape(-1, 1,3) # (1, 3)
            _pred_min = _pred.min(1)[0].reshape(-1, 1,3) # (1, 3)
            _pred_shift = ((_pred_min + _pred_max) / 2).reshape(-1, 1, 3)
            _pred_scale = (_pred_max - _pred_min).reshape(-1, 1, 3) / 2
            _pred = (_pred - _pred_shift) / _pred_scale
            pred_part_pointcloud[i] = _pred.detach().cpu().numpy()
            print(_pred.max(), _pred.min())
            _ref = ref_part_pointcloud[i]
            _ref_max = _ref.max(1)[0].reshape(-1, 1,3) # (1, 3)
            _ref_min = _ref.min(1)[0].reshape(-1, 1,3) # (1, 3)
            _ref_shift = ((_ref_min + _ref_max) / 2).reshape(-1, 1, 3)
            _ref_scale = (_ref_max - _ref_min).reshape(-1, 1, 3) / 2
            _ref = (_ref - _ref_shift) / _ref_scale
            ref_part_pointcloud[i] = _ref.detach().cpu().numpy()
    else:
        for i in range(4):
            _pred = pred_part_pointcloud[i]
            _pred_max = _pred.max(1)[0].reshape(-1, 1,3) # (1, 3)
            _pred_min = _pred.min(1)[0].reshape(-1, 1,3) # (1, 3)
            _pred_shift = ((_pred_min + _pred_max) / 2).reshape(-1, 1, 3)
            _pred_scale = (_pred_max - _pred_min).max(-1)[0].reshape(-1, 1, 1) / 2
            _pred = (_pred - _pred_shift) / _pred_scale
            pred_part_pointcloud[i] = _pred.detach().cpu().numpy()
            _ref = ref_part_pointcloud[i]
            _ref_max = _ref.max(1)[0].reshape(-1, 1,3) # (1, 3)
            _ref_min = _ref.min(1)[0].reshape(-1, 1,3) # (1, 3)
            _ref_shift = ((_ref_min + _ref_max) / 2).reshape(-1, 1, 3)
            _ref_scale = (_ref_max - _ref_min).max(-1)[0].reshape(-1, 1, 1) / 2
            _ref = (_ref - _ref_shift) / _ref_scale
            ref_part_pointcloud[i] = _ref.detach().cpu().numpy()
        
    filename = os.path.basename(args.datadir)[:-4]
    s = 'uncan' if not args.canonical else 'can'
    print(f'saved at /orion/u/w4756677/diffusion/anchorDIff/weights/{filename}_part_{args.cat}_{s}.pkl')
    pickle.dump(dict(ref=ref_part_pointcloud, pred=pred_part_pointcloud, pred_mask=pred_part_mask), open(f"/orion/u/w4756677/diffusion/anchorDIff/weights/{filename}_part_{args.cat}_{s}.pkl", "wb"))

tools/process_data_for_parts.py
diffFacto/diffFacto
import argparse 
import torch
from anchor_diff.datasets.evaluation_utils import compute_all_metrics 
from anchor_diff.datasets.dataset_utils import load_ply_withNormals
import numpy as np
from anchor_diff.utils.misc import fps
import pickle
import json
import os
from tqdm import tqdm

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--datadir',
        type=str) 
    args = parser.parse_args()
    data = pickle.load(open(args.datadir, 'rb'))
    print(data.keys())
    metrics = []
    pred = data['pred']
    pred_mask = data.get('pred_mask', None)
    ref = data['ref']
    for i in range(4):
        print(f"Part {i} --------------------------------")
        _pred = torch.from_numpy(pred[i]).cuda().float()
        _ref = torch.from_numpy(ref[i]).cuda().float()
        if pred_mask is not None:
            _pred_mask = pred_mask[i].cuda()
        metric = compute_all_metrics(_pred, _ref, 32, mask=_pred_mask)
        print(metric)
        metrics.append(metric)
        
    avg_m = {k:0 for k in metrics[0].keys()}
    weights = [0.311, 0.312, 0.306, 0.057]
    # weights = [0.2788225675, 0.2755519215, 0.2722812756, 0.1733442355]
    for i, m in enumerate(metrics):
        for k, v in m.items():
            avg_m[k] += v * weights[i]
    print(avg_m)
tools/process_data_for_parts_eval.py
diffFacto/diffFacto
import pickle
import numpy as np 
import os 
from tqdm import tqdm


seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
               'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46], 'Mug': [36, 37],
               'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27], 'Table': [47, 48, 49],
               'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40], 'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}

def save_data_as_text(data_dir, save_dir):
    data = pickle.load(open(data_dir, 'rb'))
    pred = data['pred']
    del data['pred']
    idx = [178, 278, 388, 523, 697, 711]
    # remove_top_idx = [3, 24, 43, 114, 123, 131, 135, 139, 195, 201, 209, 219, 235, 282, 335, 366, 384, 385, 390, 403, 420, 
    #                   457, 463, 489, 513, 535, 545, 574, 612, 622, 623, 644, 667, 678, 698, 711, 713, 717, 724, 732]
    
    # preds = [v for k, v in data.items() if 'pred' in k]
    m = 1
    # pred = np.concatenate(preds, axis=0)
    N = pred.shape[0]
    print(N)
    seg = data['pred_seg_mask'] + 12
    # seg = seg[None].repeat(m, axis=0)
    save_data = np.concatenate([pred, seg[..., None]], axis=-1)
    print(save_data.shape)
    # exit()
    save_data = save_data.reshape(N, -1, 4)
    for i in tqdm(range(save_data.shape[0])):
        # if i not in idx: continue
        # print(i)
        # for k, v in data.items():
        #     if 'pred' in k:
        #         print(k)
        #         _data = data[k][i]
        #         _seg = seg[i]
        #         _data = _data[_seg != 8]
        #         _data = np.concatenate([_data, _seg[_seg != 8][..., None]], axis=-1)
        #         print(_data.shape)
        #         np.savetxt(os.path.join(save_dir, f"{i}_sample_{k[-1]}.txt"), _data)
                
        # if np.all(_seg != 0) or np.all(_seg != 1):
        #     print(i)
        #     continue
        np.savetxt(os.path.join(save_dir, f"{i}.txt"), save_data[i])
                
        
    
save_data_as_text('/orion/u/w4756677/diffusion/anchorDIff/weights/ae_lion.pkl', '/orion/u/w4756677/datasets/synthesized_data/03001627')
tools/process_synthesized_data.py
diffFacto/diffFacto
import argparse
from anchor_diff.runner.gan_runner import GanRunner 
from anchor_diff.config.config import init_cfg
from anchor_diff.utils.misc import str_list
from anchor_diff.utils import dist_utils
import torch

def main():
    parser = argparse.ArgumentParser(description="Jittor Object Detection Training")

    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch'],
        default='none',
        help='job launcher') 
    parser.add_argument('--local_rank', type=int, default=0)
    parser.add_argument('--seed', type=int, default=0, help='random seed')
    parser.add_argument(
    '--sync_bn', 
    action='store_true', 
    default=False, 
    help='whether to use sync bn')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')      
    parser.add_argument(
        "--config-file",
        default="",
        metavar="FILE",
        help="path to config file",
        type=str,
    )
    parser.add_argument(
        "--task",
        default="train",
        help="train,val,val_gen",
        type=str,
    )
    parser.add_argument(
        "--prefix",
        default="exp",
        type=str
    )

    parser.add_argument(
        "--no_cuda",
        action='store_true'
    )
    parser.add_argument(
        "--no_eval",
        action='store_true'
    )
    parser.add_argument(
        "--short_val",
        action='store_true'
    )

    parser.add_argument(
        "--gen_num",
        default=400,
        type=int,
    )
    parser.add_argument(
        "--param_sample_num",
        default=10,
        type=int,
    )
    parser.add_argument(
        "--save_dir",
        default=".",
        type=str,
    )
    
    args = parser.parse_args()

    if args.no_cuda:
        device='cpu'
    else:
        device='cuda'

    assert args.task in ["train","val", "val_gen"],f"{args.task} not support, please choose [train,val,test,vis_test]"
    
    if args.config_file:
        init_cfg(args.config_file)

    if device == 'cuda':
        torch.backends.cudnn.benchmark = True
    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        args.distributed = False
        args.world_size = 1
    else:
        args.distributed = True
        dist_utils.init_dist(args.launcher)
        # re-set gpu_ids with distributed training mode
        _, world_size = dist_utils.get_dist_info()
        args.world_size = world_size

    runner = GanRunner(device, args)

    if args.task == "train":
        runner.run()
    elif args.task == "val":
        runner.val()
    elif args.task == "val_gen":
        runner.generate_samples(args.gen_num, args.param_sample_num)

if __name__ == "__main__":
    main()
tools/run_cgan.py
diffFacto/diffFacto
import argparse
from anchor_diff.runner.language_runner import LanguageRunner 
from anchor_diff.config.config import init_cfg
from anchor_diff.utils.misc import str_list
from anchor_diff.utils import dist_utils
import torch

def main():
    parser = argparse.ArgumentParser(description="Jittor Object Detection Training")

    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch'],
        default='none',
        help='job launcher') 
    parser.add_argument('--local_rank', type=int, default=0)
    parser.add_argument('--seed', type=int, default=0, help='random seed')
    parser.add_argument(
    '--sync_bn', 
    action='store_true', 
    default=False, 
    help='whether to use sync bn')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')      
    parser.add_argument(
        "--config-file",
        default="",
        metavar="FILE",
        help="path to config file",
        type=str,
    )
    parser.add_argument(
        "--task",
        default="train",
        help="train,val,val_gen",
        type=str,
    )
    parser.add_argument(
        "--prefix",
        default="exp",
        type=str
    )

    parser.add_argument(
        "--no_cuda",
        action='store_true'
    )
    parser.add_argument(
        "--no_eval",
        action='store_true'
    )
    parser.add_argument(
        "--short_val",
        action='store_true'
    )

    parser.add_argument(
        "--gen_num",
        default=400,
        type=int,
    )
    parser.add_argument(
        "--part_id",
        default=2,
        type=int,
    )
    parser.add_argument(
        '--interpolation_dir',
        default="./",
        type=str
    )
    parser.add_argument(
        "--param_sample_num",
        default=10,
        type=int,
    )
    parser.add_argument(
        "--save_dir",
        default=".",
        type=str,
    )
    
    args = parser.parse_args()

    if args.no_cuda:
        device='cpu'
    else:
        device='cuda'

    assert args.task in ["train","val", "val_gen", 'interpolation'],f"{args.task} not support, please choose [train,val,test,vis_test]"
    
    if args.config_file:
        init_cfg(args.config_file)

    if device == 'cuda':
        torch.backends.cudnn.benchmark = True
    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        args.distributed = False
        args.world_size = 1
    else:
        args.distributed = True
        dist_utils.init_dist(args.launcher)
        # re-set gpu_ids with distributed training mode
        _, world_size = dist_utils.get_dist_info()
        args.world_size = world_size

    runner = LanguageRunner(device, args)

    if args.task == "train":
        runner.run()
    if args.task == "interpolation":
        import pickle 
        data = pickle.load(open(args.interpolation_dir, "rb"))
        runner.interpolate_two_sets(data['set1'], data['set2'], args.part_id)
    elif args.task == "val":
        runner.val()
    elif args.task == "val_gen":
        runner.generate_samples(args.gen_num, args.param_sample_num)

if __name__ == "__main__":
    main()
tools/run_language.py
diffFacto/diffFacto
import argparse
from anchor_diff.runner.mixing_runner import MixingRunner 
from anchor_diff.config.config import init_cfg
from anchor_diff.utils.misc import str_list
from anchor_diff.utils import dist_utils
import torch

def main():
    parser = argparse.ArgumentParser(description="Jittor Object Detection Training")

    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch'],
        default='none',
        help='job launcher') 
    parser.add_argument('--local_rank', type=int, default=0)
    parser.add_argument('--seed', type=int, default=0, help='random seed')
    parser.add_argument(
    '--sync_bn', 
    action='store_true', 
    default=False, 
    help='whether to use sync bn')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')      
    parser.add_argument(
        "--config-file",
        default="",
        metavar="FILE",
        help="path to config file",
        type=str,
    )

    parser.add_argument(
        "--no_cuda",
        action='store_true'
    )

    parser.add_argument(
        "--save_dir",
        default=".",
        type=str,
    )
    
    args = parser.parse_args()

    if args.no_cuda:
        device='cpu'
    else:
        device='cuda'

    
    if args.config_file:
        init_cfg(args.config_file)

    if device == 'cuda':
        torch.backends.cudnn.benchmark = True
    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        args.distributed = False
        args.world_size = 1
    else:
        args.distributed = True
        dist_utils.init_dist(args.launcher)
        # re-set gpu_ids with distributed training mode
        _, world_size = dist_utils.get_dist_info()
        args.world_size = world_size

    runner = MixingRunner(device, args)

    runner.mixing()


if __name__ == "__main__":
    main()
tools/run_mixing.py
diffFacto/diffFacto
import argparse
from difffacto.runner.runner import Runner 
from difffacto.config.config import init_cfg
from difffacto.utils.misc import str_list
from difffacto.utils import dist_utils
import torch

def main():
    parser = argparse.ArgumentParser(description="Jittor Object Detection Training")

    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch'],
        default='none',
        help='job launcher') 
    parser.add_argument('--local_rank', type=int, default=0)
    parser.add_argument('--seed', type=int, default=0, help='random seed')
    parser.add_argument(
    '--sync_bn', 
    action='store_true', 
    default=False, 
    help='whether to use sync bn')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')      
    parser.add_argument(
        "--config-file",
        default="",
        metavar="FILE",
        help="path to config file",
        type=str,
    )
    parser.add_argument(
        "--task",
        default="train",
        help="train,val,val_gen",
        type=str,
    )
    parser.add_argument(
        "--prefix",
        default="exp",
        type=str
    )

    parser.add_argument(
        "--no_cuda",
        action='store_true'
    )
    parser.add_argument(
        "--no_eval",
        action='store_true'
    )
    parser.add_argument(
        "--short_val",
        action='store_true'
    )

    parser.add_argument(
        "--gen_num",
        default=400,
        type=int,
    )
    parser.add_argument(
        "--part_id",
        default=2,
        type=int,
    )
    parser.add_argument(
        '--interpolation_dir',
        default="./",
        type=str
    )
    parser.add_argument(
        "--param_sample_num",
        default=10,
        type=int,
    )
    parser.add_argument(
        "--save_dir",
        default=".",
        type=str,
    )
    
    args = parser.parse_args()

    if args.no_cuda:
        device='cpu'
    else:
        device='cuda'

    assert args.task in ["train","val", "val_gen", 'interpolation'],f"{args.task} not support, please choose [train,val,test,vis_test]"
    
    if args.config_file:
        init_cfg(args.config_file)

    if device == 'cuda':
        torch.backends.cudnn.benchmark = True
    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        args.distributed = False
        args.world_size = 1
    else:
        args.distributed = True
        dist_utils.init_dist(args.launcher)
        # re-set gpu_ids with distributed training mode
        _, world_size = dist_utils.get_dist_info()
        args.world_size = world_size

    runner = Runner(device, args)

    if args.task == "train":
        runner.run()
    if args.task == "interpolation":
        import pickle 
        data = pickle.load(open(args.interpolation_dir, "rb"))
        runner.interpolate_two_sets(data['set1'], data['set2'], args.part_id)
    elif args.task == "val":
        runner.val()
    elif args.task == "val_gen":
        runner.generate_samples(args.gen_num, args.param_sample_num)

if __name__ == "__main__":
    main()
tools/run_net.py
diffFacto/diffFacto
import argparse
from anchor_diff.runner.pretrain_runner import PretrainRunner
from anchor_diff.config.config import init_cfg
from anchor_diff.utils.misc import str_list
from anchor_diff.utils import dist_utils
import torch

def main():
    parser = argparse.ArgumentParser(description="Jittor Object Detection Training")

    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch'],
        default='none',
        help='job launcher') 
    parser.add_argument('--local_rank', type=int, default=0)
    parser.add_argument('--seed', type=int, default=0, help='random seed')
    parser.add_argument(
    '--sync_bn', 
    action='store_true', 
    default=False, 
    help='whether to use sync bn')
    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')      
    parser.add_argument(
        "--config-file",
        default="",
        metavar="FILE",
        help="path to config file",
        type=str,
    )
    parser.add_argument(
        "--task",
        default="train",
        help="train,val,val_gen",
        type=str,
    )

    parser.add_argument(
        "--no_cuda",
        action='store_true'
    )

    parser.add_argument(
        "--save_dir",
        default=".",
        type=str,
    )
    
    args = parser.parse_args()

    if args.no_cuda:
        device='cpu'
    else:
        device='cuda'

    assert args.task in ["train","val"],f"{args.task} not support, please choose [train,val,test,vis_test]"
    
    if args.config_file:
        init_cfg(args.config_file)

    if device == 'cuda':
        torch.backends.cudnn.benchmark = True
    # init distributed env first, since logger depends on the dist info.
    if args.launcher == 'none':
        args.distributed = False
        args.world_size = 1
    else:
        args.distributed = True
        dist_utils.init_dist(args.launcher)
        # re-set gpu_ids with distributed training mode
        _, world_size = dist_utils.get_dist_info()
        args.world_size = world_size

    runner = PretrainRunner(device, args)

    if args.task == "train":
        runner.run()
    elif args.task == "val":
        runner.val()

if __name__ == "__main__":
    main()
tools/run_pretrain.py
diffFacto/diffFacto
import torch 
from torch import nn
import os
import numpy as np 
from torch.utils.data import DataLoader
from torch.optim import SGD, Adam
from torch.autograd import grad
import argparse
from tqdm import tqdm
from anchor_diff.config.config import init_cfg
from anchor_diff.config.config import get_cfg, save_cfg, save_args
from anchor_diff.utils.registry import build_from_cfg,MODELS,HOOKS
from anchor_diff.utils.misc import check_file, build_file, search_ckpt, sync, parse_losses
import pickle
import time
import datetime
from anchor_diff.utils import misc
from anchor_diff.datasets.shapenet_seg import _ShapeNetSegParts
16
point_dist_dict = {
    'Chair': {'1111': [5, 5, 5, 1], '1101':[6, 6, 0, 4], '1110':[6,5,5,0]},
    'Airplane': {'1111':[8, 4, 3, 1], '1110':[8, 5, 3, 0]},
    'Lamp': {'1101':[4, 8, 0, 4], '0111':[0, 8, 4, 4]},
    'Car': {'0111':[0, 3, 3, 10], '1111':[3, 3, 3, 7]},
}

class SampleOnePart:
    def __init__(self, device, args):
        cfg = get_cfg()
        cfg.work_dir = os.path.join('/orion/u/mikacuy/anchored_diffusion', cfg.work_dir) 
        self.logger = build_from_cfg(cfg.logger,HOOKS,work_dir=cfg.work_dir)
        self.args=args
        self.local_rank = args.local_rank

        if args.seed is not None:
            if self.logger is not None:
                self.logger.print_log(f'Set random seed to {args.seed}')
            misc.set_random_seed(args.seed + args.local_rank, deterministic=False) # seed + rank, for augmentation

        self.cfg = cfg
        self.work_dir = cfg.work_dir
        self.save_num_batch = cfg.save_num_batch

        self.max_iter = cfg.max_iter
        self.sample_points=args.sample_points
        self.resume_path = cfg.resume_path
        self.max_norm = cfg.max_norm
        self.model = build_from_cfg(cfg.model,MODELS)
        if device=='cuda':
            self.model.to(self.local_rank)
        self.cat = args.cat  
        

        self.batch_size=args.batch_size
        data = pickle.load(open(args.latent_dir, 'rb'))    
        self.codes = data['part_latents']
        self.valid_id = data['valid_id']
        self.gt_mean = data['sample 0 mean']
        self.gt_logvar = data['sample 0 logvar']     
        self.token = data['token']     
        self.length = self.codes.shape[0]  
        # only save file in the 0th process
        if  self.local_rank == 0:
            save_file = build_file(self.work_dir,prefix="config.yaml")
            save_arg = build_file(self.work_dir,prefix="args.yaml")
            save_args(save_arg, args)
            save_cfg(save_file)

        self.total_num=args.total_num if args.total_num >0 else  self.length
        if args.total_num >0:
            choice = np.random.choice(self.length, self.total_num, False)
            self.codes = self.codes[choice]
            self.valid_id = self.valid_id[choice]
            self.gt_mean = self.gt_mean[choice]
            self.gt_logvar = self.gt_logvar[choice]
            self.token = self.token[choice]
        self.length = self.codes.shape[0] 



        if (cfg.pretrained_weights):
            self.load(cfg.pretrained_weights, model_only=cfg.model_only)
        
        if self.resume_path is None:
            self.resume_path = search_ckpt(self.work_dir)
        if check_file(self.resume_path):
            self.resume()


    @torch.no_grad()
    def sample_one_part(self):
        self.model.eval()
        if self.logger is not None:
            self.logger.print_log("Start running")
        out_dump = {'pred':[], 'pred_seg_mask':[], 'valid_id':[], 'codes':[], 'noise':[], 'means':[], 'logvars':[]}
        for it in tqdm(range(self.length // self.batch_size)):
            if (it + 1) * self.batch_size > self.length:
                end = self.length 
            else:
                end = (it + 1) * self.batch_size
            codes = torch.from_numpy(self.codes[it * self.batch_size:end])
            valid_id = torch.from_numpy(self.valid_id[it * self.batch_size:end])
            gt_mean = torch.from_numpy(self.gt_mean[it * self.batch_size:end])
            gt_logvar = torch.from_numpy(self.gt_logvar[it * self.batch_size:end])    
            seg_mask = []
            for i in range(valid_id.shape[0]):
                _valid = ''.join(map(str, valid_id[i].cpu().tolist()))
                mask = point_dist_dict[self.cat].get(_valid, None)
                if mask is None:
                    _seg_mask = torch.arange(4) * valid_id[i] + torch.argmax(valid_id[i]) * (1 - valid_id[i])
                    _seg_mask = _seg_mask.reshape(4, 1).expand(-1, self.sample_points // 4).reshape(self.sample_points)
                else:
                    ratio = self.sample_points // 16
                    _seg_mask = torch.arange(4).unsqueeze(-1).repeat_interleave(torch.tensor(mask) * ratio, dim=1).reshape(self.sample_points)
                seg_mask.append(_seg_mask)
            seg_mask = torch.stack(seg_mask, dim=0).cuda().int()
            print(seg_mask.shape)
            codes = codes.cuda()
            valid_id = valid_id.cuda()
            gt_mean=gt_mean.cuda()
            gt_logvar=gt_logvar.cuda()
            pred, seg_mask, valid_id, codes, noise_latents, means, logvars = self.model.sample_one_part(codes, valid_id, gt_mean, gt_logvar, seg_mask, self.args.sample_part, self.args.how_many_each, self.args.fix_size, self.args.param_sample_num, self.args.selective_param_sample)
            out_dump['pred'].append(pred.cpu().numpy())
            out_dump['pred_seg_mask'].append(seg_mask.cpu().numpy())
            out_dump['valid_id'].append(valid_id.cpu().numpy())
            out_dump['codes'].append(codes.cpu().numpy())
            out_dump['noise'].append(noise_latents.cpu().numpy())
            out_dump['means'].append(means.cpu().numpy())
            out_dump['logvars'].append(logvars.cpu().numpy())
        out_dump = {k:np.concatenate(v, axis=0) for k, v in out_dump.items()}
        save_file = build_file(self.work_dir, prefix=f"{self.args.cat}_sample{self.args.sample_part}_{self.args.how_many_each}_{self.args.param_sample_num}.pkl")
        self.logger.print_log(f"Saving to {save_file}....")
        pickle.dump(out_dump, open(save_file, 'wb'))
        # return out_dump
                
    def run(self):

        self.model.eval()
        start_time = time.time()
        results = []
        for batch_idx, pcds in enumerate(self.loader):
            z = self.optimize(pcds, batch_idx, start_time)
            with torch.no_grad():
                _result = self.model.cimle_forward(pcds, noise=z, device=self.local_rank)
                _result['latent_code'] = z.detach().cpu()
                results.append(sync(_result, to_numpy=False))
        save_dict, eval_result = self.dataset.evaluate(results, self.save_num_batch, self.local_rank)
        save_file = build_file(self.work_dir, prefix=f"val/sample_{self.prefix}_cimle_optimized.pkl")
        self.logger.print_log(f"Saving to {save_file}....")
        pickle.dump(save_dict, open(save_file, 'wb'))
        eval_result = {k:v for k, v in eval_result.items()}
        self.logger.log(eval_result,iter=self.iter)
            
        
    
    def load(self, load_path, model_only=True):
        map_location = {'cuda:%d' % 0: 'cuda:%d' % self.local_rank}
        resume_data = torch.load(load_path, map_location=map_location)

        base_ckpt = {k.replace("module.", ""): v for k, v in resume_data['model'].items()}
        self.model.load_state_dict(base_ckpt,strict=True)

        self.logger.print_log(f"Loading model parameters from {load_path}")

    def resume(self):
        self.load(self.resume_path, model_only=self.cfg.model_only)


def main():
    parser = argparse.ArgumentParser(description="Jittor Object Detection Training")

    parser.add_argument('--local_rank', type=int, default=0)

    parser.add_argument('--seed', type=int, default=0, help='random seed')


    parser.add_argument(
        "--config-file",
        default="",
        metavar="FILE",
        help="path to config file",
        type=str,
    )

    parser.add_argument(
        "--total_num",
        default=-1,
        type=int
    )

    parser.add_argument(
        "--save_dir",
        default=".",
        type=str,
    )
    parser.add_argument(
        "--batch_size",
        default=32,
        type=int,
    )
    parser.add_argument(
        "--sample_points",
        default=8192,
        type=int,
    )
    parser.add_argument(
        "--cat",
        default='Chair',
        type=str,
    )
    parser.add_argument(
        "--latent_dir", 
        type=str
    )
    parser.add_argument(
        "--sample_part", 
        type=int
    )
    parser.add_argument(
        "--how_many_each",
        default=50, 
        type=int
    )
    parser.add_argument(
        "--param_sample_num",
        default=1, 
        type=int
    )
    parser.add_argument(
        "--selective_param_sample",
        action='store_true', 
    )
    parser.add_argument(
        "--fix_size",
        action='store_true', 
    )
    
    args = parser.parse_args()

    device='cuda'

    if args.config_file:
        init_cfg(args.config_file)

    if device == 'cuda':
        torch.backends.cudnn.benchmark = True
    # init distributed env first, since logger depends on the dist info.
    args.world_size = 1
    

    runner = SampleOnePart(device, args)

    runner.sample_one_part()


if __name__ == "__main__":
    main()
tools/run_sample_one_part.py
diffFacto/diffFacto
import torch 
from torch import nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.optim import SGD, Adam
from torch.autograd import grad
import argparse
from tqdm import tqdm
from anchor_diff.config.config import init_cfg
from anchor_diff.config.config import get_cfg, save_cfg, save_args
from anchor_diff.utils.registry import build_from_cfg,MODELS,HOOKS
from anchor_diff.utils.misc import check_file, build_file, search_ckpt, sync, parse_losses
import pickle
import time
import math
import datetime
from anchor_diff.utils import misc
from anchor_diff.datasets.shapenet_seg import _ShapeNetSegParts

class ShapeEditor:
    def __init__(self, device, args):
        cfg = get_cfg()

        self.logger = build_from_cfg(cfg.logger,HOOKS,work_dir=cfg.work_dir)

        self.local_rank = args.local_rank
        self.prefix=args.prefix

        if args.seed is not None:
            if self.logger is not None:
                self.logger.print_log(f'Set random seed to {args.seed}, 'f'deterministic: {args.deterministic}')
            misc.set_random_seed(args.seed + args.local_rank, deterministic=args.deterministic) # seed + rank, for augmentation

        self.cfg = cfg
        self.work_dir = cfg.work_dir
        self.save_num_batch = cfg.save_num_batch

        self.max_iter = cfg.max_iter
        self.edit_id=args.edit_id

        self.resume_path = cfg.resume_path
        self.max_norm = cfg.max_norm
        self.model = build_from_cfg(cfg.model,MODELS)
        if device=='cuda':
            self.model.to(self.local_rank)

        
        


        self.data = pickle.load(open(args.datadir, 'rb'))
        self.idx = args.idx
        self.n_class = args.n_class
            
        # only save file in the 0th process
        if  self.local_rank == 0:
            save_file = build_file(self.work_dir,prefix="config.yaml")
            save_arg = build_file(self.work_dir,prefix="args.yaml")
            save_args(save_arg, args)
            save_cfg(save_file)


        self.iter = 0

        self.total_iter = self.max_iter

        if (cfg.pretrained_weights):
            self.load(cfg.pretrained_weights, model_only=cfg.model_only)
        
        if self.resume_path is None:
            self.resume_path = search_ckpt(self.work_dir)
        if check_file(self.resume_path):
            self.resume()


    
    def optimize(self, start_time):
        if self.logger is not None:
            self.logger.print_log("Start running")
        noise = torch.randn(1, self.model.encoder.part_aligner.noise_dim).cuda()
        z = nn.Parameter(noise)
        prev = torch.zeros(1).cuda()
        lr = 1
        optimizer = Adam([z], lr=lr)
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.5, patience=10, verbose=True, min_lr=5e-2)
        for it in range(self.max_iter):
            optimizer.zero_grad()
            input = self.data['input'][self.idx].cuda().unsqueeze(0)
            seg_mask = self.data['seg_mask'][self.idx].cuda().unsqueeze(0)
            ref_means = self.data['part_shift'][self.idx].cuda().unsqueeze(0)
            ref_vars = self.data['part_scale'][self.idx].cuda().unsqueeze(0) ** 2
            valid_id = self.data['present'][self.idx].cuda().unsqueeze(0)
            # print(valid_id.shape)
            edit_part_mean = None
            edit_part_var = ref_vars[..., self.edit_id]
            print(edit_part_var)
            edit_part_var[:, [2]] *= 1.2
            print(edit_part_var)
            # print(edit_part_mean.shape, edit_part_var.shape)
            fix_ids = torch.tensor([0, 1, 1, 1]).cuda()
            seg_flag = F.one_hot(seg_mask, num_classes=self.n_class)
            losses = self.model.edit_latent(z, input, seg_flag,valid_id, ref_means, ref_vars, fix_ids, self.edit_id, edit_part_mean, edit_part_var, fit_weight=0.05)
            all_loss, losses = parse_losses(losses)
            # all_loss =  (z ** 2).sum(1)
            # print(grad(outputs=all_loss, inputs=z))
            all_loss.backward()
            # print(z.grad, z.requires_grad)
            optimizer.step()
            scheduler.step(all_loss)
            with torch.no_grad():
                ptime = time.time()-start_time
                eta_time = (self.max_iter-it)*ptime/ (it + 1)
                eta_str = str(datetime.timedelta(seconds=int(eta_time)))
                data = dict(
                    name = self.cfg.name,
                    lr = optimizer.param_groups[0]['lr'],
                    iter = it,
                    total_loss = all_loss,
                    eta=eta_str,
                    z_mean = z.mean(),
                )
                data.update(losses)
                data = sync(data)
                if self.logger is not None:
                    self.logger.log(data)
            if torch.allclose(all_loss, prev):
                break 
            prev = all_loss
        return z
                
    def run(self):

        self.model.eval()
        start_time = time.time()
        results = []
        z = self.optimize(start_time)
        with torch.no_grad():
            pcds = {k:v[self.idx].unsqueeze(0) for k, v in self.data.items()}
            for k, v in pcds.items():
                print(k, v.shape)
            _result = self.model.cimle_forward(pcds, noise=z.unsqueeze(1), device=self.local_rank)
            _result['latent_code'] = z.detach().cpu()
        save_dict=dict(pred=_result['pred'].detach().cpu().numpy(), ref=_result['input_ref'].detach().cpu().numpy(), seg_mask_ref=pcds['ref_seg_mask'].detach().cpu().numpy())
        save_file = build_file(self.work_dir, prefix=f"val/sample_{self.prefix}_edited_{self.idx}.pkl")
        self.logger.print_log(f"Saving to {save_file}....")
        pickle.dump(save_dict, open(save_file, 'wb'))
            
        
    
    def load(self, load_path, model_only=False):
        map_location = {'cuda:%d' % 0: 'cuda:%d' % self.local_rank}
        resume_data = torch.load(load_path, map_location=map_location)

        if (not model_only):
            meta = resume_data.get("meta",dict())
            self.epoch = meta.get("epoch",self.epoch)
            self.iter = meta.get("iter",self.iter)
            self.max_iter = meta.get("max_iter",self.max_iter)
            self.max_epoch = meta.get("max_epoch",self.max_epoch)
            self.scheduler.load_state_dict(resume_data.get("scheduler",dict()))
            # self.optimizer.load_state_dict(resume_data.get("optimizer",dict()))
        base_ckpt = {k.replace("module.", ""): v for k, v in resume_data['model'].items()}
        self.model.load_state_dict(base_ckpt,strict=True)

        self.logger.print_log(f"Loading model parameters from {load_path}")

    def resume(self):
        self.load(self.resume_path, model_only=self.cfg.model_only)


def main():
    parser = argparse.ArgumentParser(description="Jittor Object Detection Training")

    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch'],
        default='none',
        help='job launcher') 
    parser.add_argument('--local_rank', type=int, default=0)
    parser.add_argument('--edit_id', type=int, default=0)
    parser.add_argument('--idx', type=int, default=21)
    parser.add_argument('--n_class', type=int, default=4)
    parser.add_argument('--datadir', type=str)

    parser.add_argument('--seed', type=int, default=0, help='random seed')

    parser.add_argument(
        '--deterministic',
        action='store_true',
        help='whether to set deterministic options for CUDNN backend.')

    parser.add_argument(
        "--config-file",
        default="",
        metavar="FILE",
        help="path to config file",
        type=str,
    )

    parser.add_argument(
        "--prefix",
        default="exp",
        type=str
    )

    parser.add_argument(
        "--save_dir",
        default=".",
        type=str,
    )
    parser.add_argument(
        "--batch_size",
        default=32,
        type=int,
    )
    
    args = parser.parse_args()

    device='cuda'

    if args.config_file:
        init_cfg(args.config_file)

    if device == 'cuda':
        torch.backends.cudnn.benchmark = True
    # init distributed env first, since logger depends on the dist info.
    args.world_size = 1
    

    runner = ShapeEditor(device, args)

    runner.run()


if __name__ == "__main__":
    main()
tools/shape_edit.py
diffFacto/diffFacto
import argparse 
import torch
from anchor_diff.datasets.evaluation_utils import _pairwise_EMD_CD_ 
from anchor_diff.datasets.dataset_utils import load_ply_withNormals, shapenet_part_normal_cat_to_id
import numpy as np
from anchor_diff.utils.misc import fps
import pickle
import json
import os
from tqdm import tqdm
from anchor_diff.datasets.evaluation_utils import compute_all_metrics 


seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
                'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46],
                'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27],
                'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40],
                'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--datadir',
        type=str) 
    parser.add_argument(
        '--cat',
        type=str) 
    parser.add_argument(
        '--num_points_sample',
        type=int,
        default=2048) 
    parser.add_argument(
        '--num_points_compute',
        type=int,
        default=2048) 
    parser.add_argument(
        '--bs',
        type=int,
        default=32) 
    
    args = parser.parse_args()
    data = pickle.load(open(args.datadir, 'rb'))
    cat_id = shapenet_part_normal_cat_to_id[args.cat]
    print(data.keys())
    assert 'pred' in data.keys() and ('pred_seg_mask' in data.keys() or 'seg_mask_ref' in data.keys())
    pred = torch.from_numpy(data['pred']).cuda()
    # pred = torch.stack([-pred[..., 2], pred[..., 1], pred[..., 0]], dim=-1)
    pred_seg_mask = data.get('pred_seg_mask', None)
    if pred_seg_mask is None:
        pred_seg_mask = data['seg_mask_ref']
    pred_seg_mask = torch.from_numpy(pred_seg_mask).cuda()
    size = pred.shape[0]
    pred_part_pointcloud = [[], [], [], []]
    ref_part_pointcloud = [[], [], [], []]
    root = '/orion/u/w4756677/datasets/coalace_data_processed' if args.cat != 'Car' else '/orion/u/w4756677/datasets/shapenetcore_partanno_segmentation_benchmark_v0_normal'
    with open(f"{root}/train_test_split/shuffled_test_file_list.json", 'r') as f:
        test_ids = set([str(d.split('/')[2]) for d in json.load(f)])
    
    fns = sorted(os.listdir(f"{root}/{cat_id}"))  # Get each txt file under the folder according to the path and put it in the list
    if args.cat != 'Car':
        densePoints = [f for f in fns if (f[-7:] == 'ori.ply')]
        labels = [f for f in fns if (f[-10:] == 'ori.labels')]
        
        labels = [os.path.join(f"{root}/{cat_id}", fn) for fn in labels if ((fn[0:-11] in test_ids))]
        densePoints = [os.path.join(f"{root}/{cat_id}", fn) for fn in densePoints if ((fn[0:-8] in test_ids) )]
        print(len(labels), len(densePoints))
    else:
        densePoints = [f for f in fns if (f[-3:] == 'txt')]
        densePoints = [os.path.join(f"{root}/{cat_id}", fn) for fn in densePoints if ((fn[0:-4] in test_ids) )]
        print(len(densePoints))
        
    print(size)
    pred_params = []
    ref_params = []
    for i in tqdm(range(size)):
        _pred = pred[i]
        _pred_seg_mask = pred_seg_mask[i]
            
        _pred_max = _pred.max(0)[0].reshape(1,3) # (1, 3)
        _pred_min = _pred.min(0)[0].reshape(1,3) # (1, 3)
        _pred_shift = ((_pred_min + _pred_max) / 2).reshape(1, 3)
        _pred_scale = (_pred_max - _pred_min).max(-1)[0].reshape(1, 1) / 2
        _pred = (_pred - _pred_shift) / _pred_scale
        one_pred_params = []
        for j in range(4):
            idx = _pred_seg_mask == j 
            if torch.any(idx):
                points = _pred[idx]
                if points.shape[0] > 100:
                    part_mean = (points.max(0, keepdim=True)[0] - points.min(0, keepdim=True)[0]) / 2.0
                    part_std = (points.max(0, keepdim=True)[0] - points.min(0, keepdim=True)[0]) / 2.0
                    one_pred_params.append(torch.rand(args.num_points_sample, 3).cuda() * part_std + part_mean)
        one_pred_params = torch.cat(one_pred_params, dim=0).unsqueeze(0)
        if one_pred_params.shape[1] > args.num_points_compute:
            one_pred_params = fps(one_pred_params, args.num_points_compute)
        pred_params.append(one_pred_params)
        
    for i in tqdm(range(len(densePoints))):
        if args.cat != 'Car':
            _ref = torch.from_numpy(load_ply_withNormals(densePoints[i]).astype(np.float32)).cuda()
            _ref_seg_mask = torch.from_numpy(np.loadtxt(labels[i]).astype(np.int32)).cuda() - 1
        else:
            data = torch.from_numpy(np.loadtxt(densePoints[i]).astype(np.float32)).cuda()
            _ref_seg_mask = data[:, -1].to(torch.int32) - seg_classes[args.cat][0] # Get the label of the sub category
            _ref = data[:, 0:3]
            
        _ref_max = _ref.max(0)[0].reshape(1,3) # (1, 3)
        _ref_min = _ref.min(0)[0].reshape(1,3) # (1, 3)
        _ref_shift = ((_ref_min + _ref_max) / 2).reshape(1, 3)
        _ref_scale = (_ref_max - _ref_min).max(-1)[0].reshape(1, 1) / 2
        _ref = (_ref - _ref_shift) / _ref_scale
        one_ref_params = []
        for j in range(4):
            ref_idx = _ref_seg_mask == j
            if args.cat != "Car":
                ref_points = _ref[ref_idx]
                if ref_points.abs().sum() != 0:
                    ref_part_mean = (ref_points.max(0, keepdim=True)[0] - ref_points.min(0, keepdim=True)[0]) / 2.0
                    ref_part_std = (ref_points.max(0, keepdim=True)[0] - ref_points.min(0, keepdim=True)[0]) / 2.0
                    one_ref_params.append(torch.rand(args.num_points_sample, 3).cuda() * ref_part_std + ref_part_mean)
            else:
                if torch.any(ref_idx):
                    ref_points = _ref[ref_idx]
                    if ref_points.shape[0] > 100:
                        ref_part_mean = (ref_points.max(0, keepdim=True)[0] - ref_points.min(0, keepdim=True)[0]) / 2.0
                        ref_part_std = (ref_points.max(0, keepdim=True)[0] - ref_points.min(0, keepdim=True)[0]) / 2.0
                        one_ref_params.append(torch.rand(args.num_points_sample, 3).cuda() * ref_part_std + ref_part_mean)
        one_ref_params = torch.cat(one_ref_params, dim=0).unsqueeze(0)
        if one_ref_params.shape[1] > args.num_points_compute:
            one_ref_params = fps(one_ref_params, args.num_points_compute)
        ref_params.append(one_ref_params)
    
    pred_params = torch.cat(pred_params,dim=0)
    ref_params = torch.cat(ref_params,dim=0)
            
    print(pred_params.shape, ref_params.shape)
    metric = compute_all_metrics(pred_params, ref_params, 32)
    print(metric)
    # filename = os.path.basename(args.datadir)[:-4]
    # print(f'saved at /orion/u/w4756677/diffusion/anchorDIff/weights/{filename}_param_{args.cat}.pkl')
    # pickle.dump(dict(ref=ref_params, pred=pred_params, metric=metric), open(f"/orion/u/w4756677/diffusion/anchorDIff/weights/{filename}_param_{args.cat}.pkl", "wb"))

tools/test_bbox_metric.py
diffFacto/diffFacto
import argparse 
import torch
import os
from anchor_diff.datasets.evaluation_utils import _pairwise_EMD_CD_ 
from anchor_diff.datasets.dataset_utils import load_ply_withNormals, shapenet_part_normal_cat_to_id
import numpy as np
from anchor_diff.utils.misc import fps
import pickle
import json
import os
from tqdm import tqdm
from anchor_diff.datasets.evaluation_utils import compute_bbox_metric 


seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
                'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46],
                'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27],
                'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40],
                'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--data_dir',
        type=str) 
    parser.add_argument(
        '--ref_dir',
        type=str) 
    parser.add_argument(
        '--save_dir',
        type=str
    )
    parser.add_argument(
        '--prefix',
        type=str
    )
    parser.add_argument(
        '--cat',
        type=str) 
    parser.add_argument(
        '--num_points_sample',
        type=int,
        default=2048) 
    parser.add_argument(
        '--num_points_compute',
        type=int,
        default=2048) 
    parser.add_argument(
        '--bs',
        type=int,
        default=32) 
    parser.add_argument(
        '--cov_thresh',
        type=int,
        default=4) 
    parser.add_argument(
        '--thresh',
        type=int,
        default=95) 
    parser.add_argument(
        '--chamfer',
        action='store_true') 
    parser.add_argument(
        '--no_nn',
        action='store_true') 
    args = parser.parse_args()
    thresh = args.thresh / 100
    cov_thresh = [0.03, 0.04, 0.05, 0.06, 100][args.cov_thresh]
    data = pickle.load(open(args.data_dir, 'rb'))
    refs = pickle.load(open(args.ref_dir, 'rb'))
    cat_id = shapenet_part_normal_cat_to_id[args.cat]
    print(data.keys())
    assert 'pred' in data.keys() and ('pred_seg_mask' in data.keys() or 'seg_mask_ref' in data.keys())
    pred = torch.from_numpy(data['pred']).cuda()[:2800]
    ref = torch.from_numpy(refs['ref']).cuda()
    # pred = torch.stack([-pred[..., 2], pred[..., 1], pred[..., 0]], dim=-1)
    pred_seg_mask = data.get('pred_seg_mask', None)
    ref_seg_mask = refs.get('ref_seg_mask', None)
    if pred_seg_mask is None:
        pred_seg_mask = data['seg_mask_ref']
    pred_seg_mask = torch.from_numpy(pred_seg_mask).cuda()[:2800]
    ref_seg_mask = torch.from_numpy(ref_seg_mask).cuda()
    metric = 'chamfer' if args.chamfer else 'iou'
    bbox_metrics = compute_bbox_metric(pred, pred_seg_mask, ref, ref_seg_mask, 32, n_class=4, thresh=thresh, metric=metric, no_nn=args.no_nn, cov_thresh=cov_thresh)
    for k, v in bbox_metrics.items():
        print(f"{k}: {v}")
    with open(os.path.join(args.save_dir, f"{args.prefix}_{args.cat}_robust_thresh{args.thresh}_cov_thresh_{args.cov_thresh}.txt"), 'w') as f:
        for k, v in bbox_metrics.items():
            f.write(f"{k}: {v}")
            f.write('\n')
    
tools/test_bbox_metric_robust.py
diffFacto/diffFacto
import argparse 
import torch
from anchor_diff.datasets.evaluation_utils import _pairwise_EMD_CD_ 
from anchor_diff.datasets.dataset_utils import load_ply_withNormals, shapenet_part_normal_cat_to_id
import numpy as np
from anchor_diff.utils.misc import fps
import pickle
import json
import os
from tqdm import tqdm
from anchor_diff.datasets.evaluation_utils import compute_all_metrics 


seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
                'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46],
                'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27],
                'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40],
                'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--cat',
        type=str) 
    parser.add_argument(
        '--num_points_sample',
        type=int,
        default=2048) 
    parser.add_argument(
        '--num_points_compute',
        type=int,
        default=2048) 
    parser.add_argument(
        '--bs',
        type=int,
        default=32) 
    parser.add_argument(
        '--thresh',
        type=int,
        default=95) 
    args = parser.parse_args()
    cat_id = shapenet_part_normal_cat_to_id[args.cat]
    root = '/orion/u/w4756677/datasets/coalace_data_processed' if args.cat != 'Car' else '/orion/u/w4756677/datasets/shapenetcore_partanno_segmentation_benchmark_v0_normal'
    with open(f"{root}/train_test_split/shuffled_test_file_list.json", 'r') as f:
        test_ids = set([str(d.split('/')[2]) for d in json.load(f)])
    
    fns = sorted(os.listdir(f"{root}/{cat_id}"))  # Get each txt file under the folder according to the path and put it in the list
    if args.cat != 'Car':
        densePoints = [f for f in fns if (f[-7:] == 'ori.ply')]
        labels = [f for f in fns if (f[-10:] == 'ori.labels')]
        
        labels = [os.path.join(f"{root}/{cat_id}", fn) for fn in labels if ((fn[0:-11] in test_ids))]
        densePoints = [os.path.join(f"{root}/{cat_id}", fn) for fn in densePoints if ((fn[0:-8] in test_ids) )]
        print(len(labels), len(densePoints))
    else:
        densePoints = [f for f in fns if (f[-3:] == 'txt')]
        densePoints = [os.path.join(f"{root}/{cat_id}", fn) for fn in densePoints if ((fn[0:-4] in test_ids) )]
        print(len(densePoints))
        
    ref = []
        
    # pred_params = torch.cat(pred_params,dim=0)
    # pickle.dump(dict(pred=torch.cat(pred_params,dim=0).cpu().numpy(), ref=pred.cpu().numpy()), open("/orion/u/w4756677/diffusion/anchorDIff/weights/out.pkl", 'wb'))
    refs_out = []
    masks_out = []
    for i in tqdm(range(len(densePoints))):
        if args.cat != 'Car':
            _ref = torch.from_numpy(load_ply_withNormals(densePoints[i]).astype(np.float32)).cuda()
            _ref_seg_mask = torch.from_numpy(np.loadtxt(labels[i]).astype(np.int32)).cuda() - 1
        else:
            data = torch.from_numpy(np.loadtxt(densePoints[i]).astype(np.float32)).cuda()
            _ref_seg_mask = data[:, -1].to(torch.int32) - seg_classes[args.cat][0] # Get the label of the sub category
            _ref = data[:, 0:3]
            
        
        ref_out = torch.zeros(4, 512, 3)
        mask_out = torch.zeros(4, 512)
        valid_id = [0 for _ in range(4)]
        for j in range(4):
            ref_idx = _ref_seg_mask == j
            ref_points = _ref[ref_idx]
            if ref_points.abs().sum() != 0:
                valid_id[j] = 1
                ref_out[j] = ref_points[:512]
                mask_out[j] = j
        valid_part = torch.tensor(valid_id).argmax()
        valid_point = ref_out[valid_part]
        for j in range(4):
            if valid_id[j] == 0:
                mask_out[j] = valid_part
                ref_out[j] = valid_point
        ref_out = ref_out.reshape(2048, 3)
        mask_out = mask_out.reshape(2048)
        ref_out_max = ref_out.max(0)[0].reshape(1,3) # (1, 3)
        ref_out_min = ref_out.min(0)[0].reshape(1,3) # (1, 3)
        ref_out_shift = ((ref_out_min + ref_out_max) / 2).reshape(1, 3)
        ref_out_scale = (ref_out_max - ref_out_min).max(-1)[0].reshape(1, 1) / 2
        ref_out = (ref_out - ref_out_shift) / ref_out_scale
        refs_out.append(ref_out)
        masks_out.append(mask_out)
    
    refs_out = torch.cat(refs_out, dim=0).reshape(-1, 2048, 3)
    masks_out = torch.cat(masks_out, dim=0).reshape(-1, 2048).int()
    print(refs_out.shape[0], masks_out.shape[0])
    assert (masks_out.shape[0] == refs_out.shape[0])
    pickle.dump(dict(ref=refs_out.cpu().numpy(), ref_seg_mask=masks_out.cpu().numpy()), open("/orion/u/w4756677/diffusion/anchorDIff/weights/out_ref_airplane.pkl", 'wb'))
            
    # filename = os.path.basename(args.datadir)[:-4]
    # print(f'saved at /orion/u/w4756677/diffusion/anchorDIff/weights/{filename}_param_{args.cat}.pkl')
    # pickle.dump(dict(ref=ref_params, pred=pred_params, metric=metric), open(f"/orion/u/w4756677/diffusion/anchorDIff/weights/{filename}_param_{args.cat}.pkl", "wb"))

tools/test_bbox_metric_robust_ copy.py
diffFacto/diffFacto
import argparse 
import torch
from anchor_diff.datasets.evaluation_utils import _pairwise_EMD_CD_ 
from anchor_diff.datasets.dataset_utils import load_ply_withNormals, shapenet_part_normal_cat_to_id
import numpy as np
from anchor_diff.utils.misc import fps
import pickle
import json
import os
from tqdm import tqdm
from anchor_diff.datasets.evaluation_utils import compute_all_metrics 


seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
                'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46],
                'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27],
                'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40],
                'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--datadir',
        type=str) 
    parser.add_argument(
        '--cat',
        type=str) 
    parser.add_argument(
        '--num_points_sample',
        type=int,
        default=2048) 
    parser.add_argument(
        '--num_points_compute',
        type=int,
        default=2048) 
    parser.add_argument(
        '--bs',
        type=int,
        default=32) 
    parser.add_argument(
        '--thresh',
        type=int,
        default=95) 
    args = parser.parse_args()
    thresh = args.thresh / 100
    data = pickle.load(open(args.datadir, 'rb'))
    cat_id = shapenet_part_normal_cat_to_id[args.cat]
    print(data.keys())
    assert 'pred' in data.keys() and ('pred_seg_mask' in data.keys() or 'seg_mask_ref' in data.keys())
    
    pred = torch.from_numpy(data['pred']).cuda()
    # pred = torch.stack([-pred[..., 2], pred[..., 1], pred[..., 0]], dim=-1)
    pred_seg_mask = data.get('pred_seg_mask', None)
    if pred_seg_mask is None:
        pred_seg_mask = data['seg_mask_ref']
    pred_seg_mask = torch.from_numpy(pred_seg_mask).cuda()
    size = pred.shape[0]
    pred_part_pointcloud = [[], [], [], []]
    ref_part_pointcloud = [[], [], [], []]
    root = '/orion/u/w4756677/datasets/coalace_data_processed' if args.cat != 'Car' else '/orion/u/w4756677/datasets/shapenetcore_partanno_segmentation_benchmark_v0_normal'
    with open(f"{root}/train_test_split/shuffled_test_file_list.json", 'r') as f:
        test_ids = set([str(d.split('/')[2]) for d in json.load(f)])
    
    fns = sorted(os.listdir(f"{root}/{cat_id}"))  # Get each txt file under the folder according to the path and put it in the list
    if args.cat != 'Car':
        densePoints = [f for f in fns if (f[-7:] == 'ori.ply')]
        labels = [f for f in fns if (f[-10:] == 'ori.labels')]
        
        labels = [os.path.join(f"{root}/{cat_id}", fn) for fn in labels if ((fn[0:-11] in test_ids))]
        densePoints = [os.path.join(f"{root}/{cat_id}", fn) for fn in densePoints if ((fn[0:-8] in test_ids) )]
        print(len(labels), len(densePoints))
    else:
        densePoints = [f for f in fns if (f[-3:] == 'txt')]
        densePoints = [os.path.join(f"{root}/{cat_id}", fn) for fn in densePoints if ((fn[0:-4] in test_ids) )]
        print(len(densePoints))
        
    print(size)
    pred_params = []
    ref_params = []
    ref = []
    for i in tqdm(range(size)):
        _pred = pred[i]
        _pred_seg_mask = pred_seg_mask[i]
            
        _pred_max = _pred.max(0)[0].reshape(1,3) # (1, 3)
        _pred_min = _pred.min(0)[0].reshape(1,3) # (1, 3)
        _pred_shift = ((_pred_min + _pred_max) / 2).reshape(1, 3)
        _pred_scale = (_pred_max - _pred_min).max(-1)[0].reshape(1, 1) / 2
        _pred = (_pred - _pred_shift) / _pred_scale
        one_pred_params = []
        for j in range(4):
            idx = _pred_seg_mask == j 
            if torch.any(idx):
                points = _pred[idx]
                if points.shape[0] > 100:
                    
                    part_max = torch.quantile(points, thresh, dim=0, keepdim=True) 
                    part_min =  torch.quantile(points, 1 - thresh, dim=0, keepdim=True)
                    # part_std = (torch.quantile(points, thresh, dim=0, keepdim=True) - torch.quantile(points, 1 - thresh, dim=0, keepdim=True)) / 2.0
                    one_pred_params.append(torch.rand(args.num_points_sample, 3).cuda() * (part_max - part_min) + part_min)
        one_pred_params = torch.cat(one_pred_params, dim=0).unsqueeze(0)
        if one_pred_params.shape[1] > args.num_points_compute:
            one_pred_params = fps(one_pred_params, args.num_points_compute)
        pred_params.append(one_pred_params)
        
    # pred_params = torch.cat(pred_params,dim=0)
    # pickle.dump(dict(pred=torch.cat(pred_params,dim=0).cpu().numpy(), ref=pred.cpu().numpy()), open("/orion/u/w4756677/diffusion/anchorDIff/weights/out.pkl", 'wb'))
    for i in tqdm(range(len(densePoints))):
        if args.cat != 'Car':
            _ref = torch.from_numpy(load_ply_withNormals(densePoints[i]).astype(np.float32)).cuda()
            _ref_seg_mask = torch.from_numpy(np.loadtxt(labels[i]).astype(np.int32)).cuda() - 1
        else:
            data = torch.from_numpy(np.loadtxt(densePoints[i]).astype(np.float32)).cuda()
            _ref_seg_mask = data[:, -1].to(torch.int32) - seg_classes[args.cat][0] # Get the label of the sub category
            _ref = data[:, 0:3]
            
        _ref_max = _ref.max(0)[0].reshape(1,3) # (1, 3)
        _ref_min = _ref.min(0)[0].reshape(1,3) # (1, 3)
        _ref_shift = ((_ref_min + _ref_max) / 2).reshape(1, 3)
        _ref_scale = (_ref_max - _ref_min).max(-1)[0].reshape(1, 1) / 2
        _ref = (_ref - _ref_shift) / _ref_scale
        ref.append(_ref.unsqueeze(0))
        one_ref_params = []
        for j in range(4):
            ref_idx = _ref_seg_mask == j
            if args.cat != "Car":
                ref_points = _ref[ref_idx]
                if ref_points.abs().sum() != 0:
                    ref_part_max = torch.quantile(ref_points, thresh, dim=0, keepdim=True) 
                    ref_part_min =  torch.quantile(ref_points, 1 - thresh, dim=0, keepdim=True)
                    one_ref_params.append(torch.rand(args.num_points_sample, 3).cuda() * (ref_part_max - ref_part_min) + ref_part_min)
            else:
                if torch.any(ref_idx):
                    ref_points = _ref[ref_idx]
                    if ref_points.shape[0] > 100:
                        ref_part_max = torch.quantile(ref_points, thresh, dim=0, keepdim=True) 
                        ref_part_min =  torch.quantile(ref_points, 1 - thresh, dim=0, keepdim=True)
                        one_ref_params.append(torch.rand(args.num_points_sample, 3).cuda() * (ref_part_max - ref_part_min) + ref_part_min)
        one_ref_params = torch.cat(one_ref_params, dim=0).unsqueeze(0)
        if one_ref_params.shape[1] > args.num_points_compute:
            one_ref_params = fps(one_ref_params, args.num_points_compute)
        ref_params.append(one_ref_params)
    
    pred_params = torch.cat(pred_params,dim=0)
    ref_params = torch.cat(ref_params,dim=0)
    ref = torch.cat(ref, dim=0)
    pickle.dump(dict(pred=ref_params.cpu().numpy(), pred_ref=ref.cpu().numpy()), open("/orion/u/w4756677/diffusion/anchorDIff/weights/out_ref.pkl", 'wb'))
    pickle.dump(dict(pred=pred_params.cpu().numpy(), pred_ref=pred.cpu().numpy()), open("/orion/u/w4756677/diffusion/anchorDIff/weights/out.pkl", 'wb'))
            
    print(pred_params.shape, ref_params.shape)
    metric = compute_all_metrics(pred_params, ref_params, 32)
    print(metric)
    # filename = os.path.basename(args.datadir)[:-4]
    # print(f'saved at /orion/u/w4756677/diffusion/anchorDIff/weights/{filename}_param_{args.cat}.pkl')
    # pickle.dump(dict(ref=ref_params, pred=pred_params, metric=metric), open(f"/orion/u/w4756677/diffusion/anchorDIff/weights/{filename}_param_{args.cat}.pkl", "wb"))

tools/test_bbox_metric_robust_.py
diffFacto/diffFacto
import argparse 
import torch
import os
from anchor_diff.datasets.evaluation_utils import _pairwise_EMD_CD_ 
from anchor_diff.datasets.dataset_utils import load_ply_withNormals, shapenet_part_normal_cat_to_id
import numpy as np
from anchor_diff.utils.misc import fps
import pickle
import json
import os
from tqdm import tqdm
from anchor_diff.datasets.evaluation_utils import compute_snapping_metric 


seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43],
                'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46],
                'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27],
                'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40],
                'Chair': [12, 13, 14, 15], 'Knife': [22, 23]}

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--data_dir',
        type=str) 
    parser.add_argument(
        '--ref_dir',
        type=str) 
    parser.add_argument(
        '--save_dir',
        type=str
    )
    parser.add_argument(
        '--prefix',
        type=str
    )
    parser.add_argument(
        '--cat',
        type=str) 
    parser.add_argument(
        '--num_points_sample',
        type=int,
        default=2048) 
    parser.add_argument(
        '--num_points_compute',
        type=int,
        default=2048) 
    parser.add_argument(
        '--bs',
        type=int,
        default=32) 
    parser.add_argument(
        '--thresh',
        type=int,
        default=95) 
    args = parser.parse_args()
    thresh = args.thresh / 100
    data = pickle.load(open(args.data_dir, 'rb'))
    refs = pickle.load(open(args.ref_dir, 'rb'))
    cat_id = shapenet_part_normal_cat_to_id[args.cat]
    print(data.keys())
    assert 'pred' in data.keys() and ('pred_seg_mask' in data.keys() or 'seg_mask_ref' in data.keys())
    
    pred = torch.from_numpy(data['pred']).cuda()
    ref = torch.from_numpy(refs['ref']).cuda()
    # pred = torch.stack([-pred[..., 2], pred[..., 1], pred[..., 0]], dim=-1)
    pred_seg_mask = data.get('pred_seg_mask', None)
    ref_seg_mask = refs.get('ref_seg_mask', None)
    if pred_seg_mask is None:
        pred_seg_mask = data['seg_mask_ref']
    pred_seg_mask = torch.from_numpy(pred_seg_mask).cuda()
    ref_seg_mask = torch.from_numpy(ref_seg_mask).cuda()
    
    print("Normalizing---------")
    pred_max = pred.max(1)[0].reshape(-1, 1, 3) # (1, 3)
    pred_min = pred.min(1)[0].reshape(-1, 1, 3) # (1, 3)
    pred_shift = ((pred_min + pred_max) / 2).reshape(-1, 1, 3)
    pred_scale = (pred_max - pred_min).max(-1)[0].reshape(-1, 1, 1) / 2
    pred = (pred - pred_shift) / pred_scale
    ref_max = ref.max(1)[0].reshape(-1, 1, 3) # (1, 3)
    ref_min = ref.min(1)[0].reshape(-1, 1, 3) # (1, 3)
    ref_shift = ((ref_min + ref_max) / 2).reshape(-1, 1, 3)
    ref_scale = (ref_max - ref_min).max(-1)[0].reshape(-1, 1, 1) / 2
    ref = (ref - ref_shift) / ref_scale
    
    snapping_metrics = compute_snapping_metric(pred, pred_seg_mask, cls=args.cat)
    oracle_snapping_metrics = compute_snapping_metric(ref, ref_seg_mask, cls=args.cat)
    for k, v in snapping_metrics.items():
        print(f"{k}: {v}")
        print(f"oracle_{k}: {oracle_snapping_metrics[k]}")
    with open(os.path.join(args.save_dir, f"{args.prefix}_{args.cat}_snapping.txt"), 'w') as f:
        for k, v in snapping_metrics.items():
            f.write(f"{k}: {v}")
            f.write('\n')
            f.write(f"oracle_{k}: {oracle_snapping_metrics[k]}")
            f.write('\n')
    
tools/test_bbox_metric_snapping.py
diffFacto/diffFacto
import torch 
from tqdm import tqdm
from anchor_diff.utils.misc import fps
import numpy as np
from  anchor_diff.datasets.evaluation_utils import compute_all_metrics
import argparse
import pickle
import json
import os
from datetime import datetime
def evaluate(pred, ref, batch_size):
    preds = []
    refs = []
    print(pred.shape, ref.shape)
    for i in tqdm(range(pred.shape[0])):
        _pred = pred[i]
        if _pred.shape[0] > 2048:
            _pred = fps(_pred, 2048)
        pred_max = _pred.max(0).reshape(1, 3) # (1, 3)
        pred_min = _pred.min(0).reshape(1, 3) # (1, 3)
        pred_shift = ((pred_min + pred_max) / 2).reshape(1, 3)
        pred_scale = (pred_max - pred_min).max(-1)[0].reshape(1, 1) / 2
        _pred = (_pred - pred_shift) / pred_scale
        preds.append(_pred)
    for i in tqdm(range(ref.shape[0])):
        _ref = ref[i]
        if _ref.shape[0] > 2048:
            _ref = fps(_ref, 2048)
        ref_max = _ref.max(0).reshape(1, 3) # (1, 3)
        ref_min = _ref.min(0).reshape(1, 3) # (1, 3)
        ref_shift = ((ref_min + ref_max) / 2).reshape(1, 3)
        ref_scale = (ref_max - ref_min).max(-1)[0].reshape(1, 1) / 2
        _ref = (_ref - ref_shift) / ref_scale
        refs.append(_ref)

                     
    preds = torch.from_numpy(np.stack(preds, axis=0)).cuda()
    refs = torch.from_numpy(np.stack(refs, axis=0)).cuda()
    print(preds.shape, refs.shape)
    metrics = compute_all_metrics(preds, refs, batch_size)
    print(metrics)
    return metrics
    
    
def get_data(data, num_sample, num_part_sample):
    assert num_sample % num_part_sample == 0
    ref = data['input_ref']
    out = []
    preds = []
    for i in range(10):
        preds.append(data[f'pred_sample {i}'])
    for i in range(num_sample // num_part_sample):
        out.append(preds[i][:num_part_sample])
    return np.concatenate(out, axis=0), ref
def main():
    parser = argparse.ArgumentParser(description="Jittor Object Detection Training")



    parser.add_argument(
        "--data_dir",
        type=str,
    )
    parser.add_argument(
        "--save_dir",
        type=str,
    )
    parser.add_argument(
        "--tag",
        default='exp',
        type=str,
    )
    parser.add_argument(
        "--num_sample",
        default=600,
        type=int,
    )
    parser.add_argument(
        "--num_part_sample",
        default=200,
        type=int,
    )
    parser.add_argument(
        "--batch_size",
        default=32,
        type=int,
    )
    
    args = parser.parse_args()
    data = pickle.load(open(args.data_dir, "rb"))
    pred, ref = get_data(data, args.num_sample, args.num_part_sample)
    metrics = evaluate(pred, ref, args.batch_size)
    now = datetime.now()
    dt_string = now.strftime("%d_%m_%Y__%H:%M:%S")
    with open(os.path.join(args.save_dir, f"gen_score{dt_string}_{args.tag}.json"), "w") as outfile:
        json.dump({k: v.item() for k, v in metrics.items()}, outfile)
    


if __name__ == "__main__":
    main()
tools/test_gen.py
diffFacto/diffFacto