22#define __DENOISER_HPP__
33
44#include < cmath>
5+ #include < string>
56#include < utility>
67
78#include " ggml_extend.hpp"
@@ -1148,7 +1149,80 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
11481149 sd::Tensor<float > x,
11491150 const std::vector<float >& sigmas,
11501151 std::shared_ptr<RNG> rng,
1151- bool is_flow_denoiser) {
1152+ bool is_flow_denoiser,
1153+ const char * extra_sample_args = nullptr ) {
1154+ struct LCMSampleArgs {
1155+ float noise_clip_std = 0 .0f ;
1156+ float noise_scale_start = 1 .0f ;
1157+ float noise_scale_end = 1 .0f ;
1158+ };
1159+
1160+ auto trim = [](std::string value) -> std::string {
1161+ const char * whitespace = " \t\r\n "
1162+ size_t begin = value.find_first_not_of (whitespace);
1163+ if (begin == std::string::npos) {
1164+ return " "
1165+ }
1166+ size_t end = value.find_last_not_of (whitespace);
1167+ return value.substr (begin, end - begin + 1 );
1168+ };
1169+
1170+ LCMSampleArgs args;
1171+ if (extra_sample_args != nullptr && extra_sample_args[0 ] != ' \0 '
1172+ std::string raw (extra_sample_args);
1173+ size_t start = 0 ;
1174+ bool noise_scale_end_was_set = false ;
1175+ bool noise_scale_start_was_set = false ;
1176+ auto parse_arg = [&](const std::string& item) {
1177+ std::string token = trim (item);
1178+ if (token.empty ()) {
1179+ return ;
1180+ }
1181+ size_t eq = token.find (' ='
1182+ if (eq == std::string::npos) {
1183+ LOG_WARN (" ignoring invalid lcm extra sample arg '%s'" c_str ());
1184+ return ;
1185+ }
1186+
1187+ std::string key = trim (token.substr (0 , eq));
1188+ std::string value = trim (token.substr (eq + 1 ));
1189+ float parsed = 0 .0f ;
1190+ try {
1191+ size_t consumed = 0 ;
1192+ parsed = std::stof (value, &consumed);
1193+ if (trim (value.substr (consumed)).size () != 0 ) {
1194+ LOG_WARN (" ignoring invalid lcm extra sample arg '%s'" c_str ());
1195+ return ;
1196+ }
1197+ } catch (const std::exception&) {
1198+ LOG_WARN (" ignoring invalid lcm extra sample arg '%s'" c_str ());
1199+ return ;
1200+ }
1201+
1202+ if (key == " noise_clip_std"
1203+ args.noise_clip_std = parsed;
1204+ } else if (key == " noise_scale_start"
1205+ args.noise_scale_start = parsed;
1206+ noise_scale_start_was_set = true ;
1207+ } else if (key == " noise_scale_end"
1208+ args.noise_scale_end = parsed;
1209+ noise_scale_end_was_set = true ;
1210+ } else {
1211+ LOG_WARN (" ignoring unknown lcm extra sample arg '%s'" c_str ());
1212+ }
1213+ };
1214+
1215+ for (size_t pos = 0 ; pos <= raw.size (); ++pos) {
1216+ if (pos == raw.size () || raw[pos] == ' ,' ' ;'
1217+ parse_arg (raw.substr (start, pos - start));
1218+ start = pos + 1 ;
1219+ }
1220+ }
1221+ if (noise_scale_start_was_set && !noise_scale_end_was_set) {
1222+ args.noise_scale_end = args.noise_scale_start ;
1223+ }
1224+ }
1225+
11521226 int steps = static_cast <int >(sigmas.size ()) - 1 ;
11531227 for (int i = 0 ; i < steps; i++) {
11541228 auto denoised_opt = model (x, sigmas[i], i + 1 , nullptr );
@@ -1160,7 +1234,27 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
11601234 if (is_flow_denoiser) {
11611235 x *= (1 - sigmas[i + 1 ]);
11621236 }
1163- x += sd::Tensor<float >::randn_like (x, rng) * sigmas[i + 1 ];
1237+ auto noise = sd::Tensor<float >::randn_like (x, rng);
1238+ if (args.noise_clip_std > 0 .0f && noise.numel () > 0 ) {
1239+ double mean = 0.0 ;
1240+ for (int64_t j = 0 ; j < noise.numel (); ++j) {
1241+ mean += static_cast <double >(noise[j]);
1242+ }
1243+ mean /= static_cast <double >(noise.numel ());
1244+
1245+ double variance = 0.0 ;
1246+ for (int64_t j = 0 ; j < noise.numel (); ++j) {
1247+ double centered = static_cast <double >(noise[j]) - mean;
1248+ variance += centered * centered;
1249+ }
1250+ variance /= static_cast <double >(noise.numel ());
1251+
1252+ float clip_val = args.noise_clip_std * static_cast <float >(std::sqrt (variance));
1253+ noise = sd::ops::clamp (noise, -clip_val, clip_val);
1254+ }
1255+ float t = steps > 1 ? static_cast <float >(i) / static_cast <float >(steps - 1 ) : 0 .0f ;
1256+ float noise_scale = args.noise_scale_start + (args.noise_scale_end - args.noise_scale_start ) * t;
1257+ x += noise * (sigmas[i + 1 ] * noise_scale);
11641258 }
11651259 }
11661260 return x;
@@ -1656,15 +1750,15 @@ static sd::Tensor<float> sample_euler_cfg_pp(denoise_cb_t model,
16561750 for (int i = 0 ; i < steps; i++) {
16571751 float sigma = sigmas[i];
16581752 sd::Tensor<float > uncond_denoised;
1659-
1753+
16601754 auto denoised_opt = model (x, sigma, i + 1 , &uncond_denoised);
16611755 if (denoised_opt.empty () || uncond_denoised.empty ()) {
16621756 return {};
16631757 }
1664-
1758+
16651759 sd::Tensor<float > denoised = std::move (denoised_opt);
1666- sd::Tensor<float > d = (x - uncond_denoised) / sigma;
1667-
1760+ sd::Tensor<float > d = (x - uncond_denoised) / sigma;
1761+
16681762 x = denoised + d * sigmas[i + 1 ];
16691763 }
16701764 return x;
@@ -1679,19 +1773,19 @@ static sd::Tensor<float> sample_euler_ancestral_cfg_pp(denoise_cb_t model,
16791773 for (int i = 0 ; i < steps; i++) {
16801774 float sigma = sigmas[i];
16811775 sd::Tensor<float > uncond_denoised;
1682-
1776+
16831777 auto denoised_opt = model (x, sigma, i + 1 , &uncond_denoised);
16841778 if (denoised_opt.empty () || uncond_denoised.empty ()) {
16851779 return {};
16861780 }
1687-
1781+
16881782 sd::Tensor<float > denoised = std::move (denoised_opt);
1689- sd::Tensor<float > d = (x - uncond_denoised) / sigma;
1690-
1783+ sd::Tensor<float > d = (x - uncond_denoised) / sigma;
1784+
16911785 auto [sigma_down, sigma_up] = get_ancestral_step (sigmas[i], sigmas[i + 1 ], eta);
1692-
1786+
16931787 x = denoised + d * sigma_down;
1694-
1788+
16951789 if (sigmas[i + 1 ] > 0 ) {
16961790 x += sd::Tensor<float >::randn_like (x, rng) * sigma_up;
16971791 }
@@ -1706,7 +1800,8 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
17061800 std::vector<float > sigmas,
17071801 std::shared_ptr<RNG> rng,
17081802 float eta,
1709- bool is_flow_denoiser) {
1803+ bool is_flow_denoiser,
1804+ const char * extra_sample_args) {
17101805 switch (method) {
17111806 case EULER_A_SAMPLE_METHOD:
17121807 if (is_flow_denoiser)
@@ -1729,7 +1824,7 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
17291824 case DPMPP2Mv2_SAMPLE_METHOD:
17301825 return sample_dpmpp_2m_v2 (model, std::move (x), sigmas);
17311826 case LCM_SAMPLE_METHOD:
1732- return sample_lcm (model, std::move (x), sigmas, rng, is_flow_denoiser);
1827+ return sample_lcm (model, std::move (x), sigmas, rng, is_flow_denoiser, extra_sample_args );
17331828 case IPNDM_SAMPLE_METHOD:
17341829 return sample_ipndm (model, std::move (x), sigmas);
17351830 case IPNDM_V_SAMPLE_METHOD:
0 commit comments