Skip to content

Commit e77e4c4

Browse files
authored
feat: adapt LCM for flow models (#1413)
1 parent 7d33d4b commit e77e4c4

1 file changed

Lines changed: 6 additions & 2 deletions

File tree

src/denoiser.hpp

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1137,7 +1137,8 @@ static sd::Tensor<float> sample_dpmpp_2m_v2(denoise_cb_t model,
11371137
static sd::Tensor<float> sample_lcm(denoise_cb_t model,
11381138
sd::Tensor<float> x,
11391139
const std::vector<float>& sigmas,
1140-
std::shared_ptr<RNG> rng) {
1140+
std::shared_ptr<RNG> rng,
1141+
bool is_flow_denoiser) {
11411142
int steps = static_cast<int>(sigmas.size()) - 1;
11421143
for (int i = 0; i < steps; i++) {
11431144
auto denoised_opt = model(x, sigmas[i], i + 1);
@@ -1146,6 +1147,9 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
11461147
}
11471148
x = std::move(denoised_opt);
11481149
if (sigmas[i + 1] > 0) {
1150+
if (is_flow_denoiser) {
1151+
x *= (1 - sigmas[i + 1]);
1152+
}
11491153
x += sd::Tensor<float>::randn_like(x, rng) * sigmas[i + 1];
11501154
}
11511155
}
@@ -1671,7 +1675,7 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
16711675
case DPMPP2Mv2_SAMPLE_METHOD:
16721676
return sample_dpmpp_2m_v2(model, std::move(x), sigmas);
16731677
case LCM_SAMPLE_METHOD:
1674-
return sample_lcm(model, std::move(x), sigmas, rng);
1678+
return sample_lcm(model, std::move(x), sigmas, rng, is_flow_denoiser);
16751679
case IPNDM_SAMPLE_METHOD:
16761680
return sample_ipndm(model, std::move(x), sigmas);
16771681
case IPNDM_V_SAMPLE_METHOD:

0 commit comments

Comments
 (0)