| 12 | use_clamp = False |
| 13 | |
| 14 | class CompressionSimulation: |
| 15 | """ |
| 16 | """ |
| 17 | def __init__(self, entropy_model_enable: bool = False, |
| 18 | entropy_model_type: Literal["factorized_model", "gaussian_model"] = "factorized_model", |
| 19 | entropy_steps: Dict[str, int] = None, |
| 20 | device: device = None, |
| 21 | ada_mask_opt: bool = False, |
| 22 | ada_mask_step: int = 10_000, |
| 23 | ada_mask_strategy: Optional[str] = "learnable", |
| 24 | **kwargs) -> None: |
| 25 | self.entropy_model_enable = entropy_model_enable |
| 26 | self.entropy_model_type = entropy_model_type |
| 27 | self.entropy_steps = entropy_steps |
| 28 | self.device = device |
| 29 | |
| 30 | self.simulation_option = { |
| 31 | "means": False, |
| 32 | "scales": True, |
| 33 | "quats": True, |
| 34 | "opacities": True, |
| 35 | "sh0": True, |
| 36 | "shN": True |
| 37 | } |
| 38 | self.shN_qat = False |
| 39 | self.shN_ada_mask_opt = ada_mask_opt |
| 40 | self.shN_ada_mask_step = ada_mask_step |
| 41 | self.shN_ada_mask_strategy = ada_mask_strategy |
| 42 | |
| 43 | self.q_bitwidth = { |
| 44 | "means": None, |
| 45 | "scales": 8, |
| 46 | "quats": 8, |
| 47 | "opacities": 8, |
| 48 | "sh0": 8, |
| 49 | "shN": None |
| 50 | } |
| 51 | |
| 52 | self.bds = { |
| 53 | "means": None, |
| 54 | "scales": [-10, 2], |
| 55 | "quats": [-1, 1], |
| 56 | "opacities": [-15, 15], |
| 57 | "sh0": [-2, 4], |
| 58 | "shN": None |
| 59 | } |
| 60 | |
| 61 | self.entropy_model_option = { |
| 62 | "means": False, |
| 63 | "scales": True, |
| 64 | "quats": True, |
| 65 | "opacities": False, |
| 66 | "sh0": True, |
| 67 | "shN": False |
| 68 | } |
| 69 | |
| 70 | # turn off if entropy step < 0 |
| 71 | for name, flag in self.entropy_model_option.items(): |