| 142 | |
| 143 | |
| 144 | def initialize_handlers( |
| 145 | checkpoint_dir: Path, |
| 146 | config_path: str, |
| 147 | lm_model_path: str, |
| 148 | backend: str, |
| 149 | lm_backend: str, |
| 150 | lm_dtype: str, |
| 151 | init_llm: bool, |
| 152 | ) -> tuple[AceStepHandler, LLMHandler]: |
| 153 | os.environ["ACESTEP_CHECKPOINTS_DIR"] = str(checkpoint_dir) |
| 154 | dit_handler = AceStepHandler() |
| 155 | init_message, ok = dit_handler.initialize_service( |
| 156 | project_root=str(checkpoint_dir), |
| 157 | config_path=config_path, |
| 158 | device=backend, |
| 159 | force_dtype=torch.float32, |
| 160 | use_flash_attention=False, |
| 161 | compile_model=False, |
| 162 | offload_to_cpu=False, |
| 163 | offload_dit_to_cpu=False, |
| 164 | quantization=None, |
| 165 | ) |
| 166 | if not ok: |
| 167 | raise RuntimeError(f"ACE-Step DiT init failed: {init_message}") |
| 168 | # Keep the Python parity path in fp32 end to end so pre-DiT, DiT, and VAE |
| 169 | # all run with stable float32 math. The forced initialize_service dtype |
| 170 | # above avoids loading reduced-precision weights and then upcasting them |
| 171 | # later, which perturbs parity boundaries before and after DiT. |
| 172 | dit_handler.dtype = torch.float32 |
| 173 | if dit_handler.model is not None: |
| 174 | dit_handler.model = dit_handler.model.to(device=backend).to(dtype=torch.float32) |
| 175 | if dit_handler.text_encoder is not None: |
| 176 | dit_handler.text_encoder = dit_handler.text_encoder.to(device=backend).to(dtype=torch.float32) |
| 177 | if dit_handler.vae is not None: |
| 178 | dit_handler.vae = dit_handler.vae.to(device=backend).to(dtype=torch.float32) |
| 179 | if dit_handler.silence_latent is not None: |
| 180 | dit_handler.silence_latent = dit_handler.silence_latent.to(device=backend).to(dtype=torch.float32) |
| 181 | |
| 182 | llm_handler = LLMHandler() |
| 183 | if init_llm: |
| 184 | dtype = None |
| 185 | if lm_dtype != "auto": |
| 186 | dtype = { |
| 187 | "float32": torch.float32, |
| 188 | "float16": torch.float16, |
| 189 | "bfloat16": torch.bfloat16, |
| 190 | }[lm_dtype] |
| 191 | init_message, ok = llm_handler.initialize( |
| 192 | checkpoint_dir=str(checkpoint_dir), |
| 193 | lm_model_path=lm_model_path, |
| 194 | backend=lm_backend, |
| 195 | device=backend, |
| 196 | offload_to_cpu=False, |
| 197 | dtype=dtype, |
| 198 | ) |
| 199 | if not ok: |
| 200 | raise RuntimeError(f"ACE-Step LM init failed: {init_message}") |
| 201 | return dit_handler, llm_handler |