perf: cache CPU position embeddings across DiT runners (#2211)

This commit is contained in:
leejet authored and GitHub committed 2026-10-25 05:18:30 +08:00
1 parent a85db96aa5
commit ea209dbfcf
31 files changed
+855 -610

No files matched your search

+5 -5
View File
@@ -130,15 +130,15 @@ namespace sd::pipeline {
}
// TAE encode already returns the latent; KL/Flux VAE returns a mean/logvar pair.
sd::Tensor<float> ctrl_latent = sd->first_stage_model->get_desc() == "taesd"
? vae_output
: sd::ops::chunk(vae_output, 2, 2)[0];
ctrl_latent = sd->first_stage_model->vae_to_diffusion_latents(ctrl_latent);
? vae_output
: sd::ops::chunk(vae_output, 2, 2)[0];
ctrl_latent = sd->first_stage_model->vae_to_diffusion_latents(ctrl_latent);
if (ctrl_latent.empty() || ctrl_latent.dim() != 4 || ctrl_latent.shape()[2] > 33) {
LOG_ERROR("failed to encode control image");
return {};
}
std::vector<int64_t> padded_shape = ctrl_latent.shape();
padded_shape[2] = 33;
padded_shape[2] = 33;
sd::Tensor<float> padded(padded_shape);
sd::ops::slice_assign(&padded, 2, 0, ctrl_latent.shape()[2], ctrl_latent);
return padded;
@@ -1028,7 +1028,7 @@ namespace sd::pipeline {
std::vector<int64_t> ctrl_shape = latents.control_image.shape();
ctrl_shape[0] = hires_latent_w;
ctrl_shape[1] = hires_latent_h;
hires_control_image = sd::ops::interpolate(latents.control_image, ctrl_shape, sd::ops::InterpolateMode::Bicubic);
hires_control_image = sd::ops::interpolate(latents.control_image, ctrl_shape, sd::ops::InterpolateMode::Bicubic);
}
} else {
hires_control_image = latents.control_image;