sd: delay VAE dtype archive until after override (#12388)

VAEs have host specific dtype logic that should override the dynamic
_model_dtype. Defer the archiving of model dtypes until after.

Author: rattus128

comfy/sd.py

@@ -793,8 +793,6 @@ def estimate_memory(shape, dtype, num_layers = 16, kv_cache_multiplier = 2):
             self.first_stage_model = AutoencoderKL(**(config['params']))
         self.first_stage_model = self.first_stage_model.eval()
 
-        model_management.archive_model_dtypes(self.first_stage_model)
-
         if device is None:
             device = model_management.vae_device()
         self.device = device
@@ -803,6 +801,7 @@ def estimate_memory(shape, dtype, num_layers = 16, kv_cache_multiplier = 2):
             dtype = model_management.vae_dtype(self.device, self.working_dtypes)
         self.vae_dtype = dtype
         self.first_stage_model.to(self.vae_dtype)
+        model_management.archive_model_dtypes(self.first_stage_model)
         self.output_device = model_management.intermediate_device()
 
         mp = comfy.model_patcher.CoreModelPatcher