Skip to content

Commit 84765f5

Browse files
wbrunawbruna
authored
sd: sync to master-447-ccb6b0a (#1898)
* sd: sync to master-438-298b110 * sd: sync to master-440-3e81246 * sd: sync to master-444-a0adcfb * sd: sync to master-447-ccb6b0a
1 parent 9bb362c commit 84765f5

19 files changed

Lines changed: 2569 additions & 308 deletions

‎otherarch/sdcpp/cache_dit.hpp‎

Lines changed: 975 additions & 0 deletions
Large diffs are not rendered by default.

‎otherarch/sdcpp/common.hpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ class DownSampleBlock : public GGMLBlock {
2828
if (vae_downsample) {
2929
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
3030

31-
x = ggml_pad(ctx->ggml_ctx, x, 1, 1, 0, 0);
31+
x = ggml_ext_pad(ctx->ggml_ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
3232
x = conv->forward(ctx, x);
3333
} else {
3434
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["op"]);

‎otherarch/sdcpp/common/common.hpp‎

Lines changed: 221 additions & 79 deletions
Large diffs are not rendered by default.

‎otherarch/sdcpp/denoiser.hpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -366,18 +366,18 @@ struct KLOptimalScheduler : SigmaScheduler {
366366

367367
for (uint32_t i = 0; i < n; ++i) {
368368
// t goes from 0.0 to 1.0
369-
float t = static_cast<float>(i) / static_cast<float>(n-1);
369+
float t = static_cast<float>(i) / static_cast<float>(n - 1);
370370

371371
// Interpolate in the angle domain
372372
float angle = t * alpha_min + (1.0f - t) * alpha_max;
373373

374374
// Convert back to sigma
375375
sigmas.push_back(std::tan(angle));
376-
}
376+
}
377377

378378
// Append the final zero to sigma
379379
sigmas.push_back(0.0f);
380-
380+
381381
return sigmas;
382382
}
383383
};

‎otherarch/sdcpp/diffusion_model.hpp‎

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,9 @@ struct DiffusionModel {
3737
virtual void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) = 0;
3838
virtual size_t get_params_buffer_size() = 0;
3939
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter){};
40-
virtual int64_t get_adm_in_channels() = 0;
41-
virtual void set_flash_attn_enabled(bool enabled) = 0;
40+
virtual int64_t get_adm_in_channels() = 0;
41+
virtual void set_flash_attn_enabled(bool enabled) = 0;
42+
virtual void set_circular_axes(bool circular_x, bool circular_y) = 0;
4243
};
4344

4445
struct UNetModel : public DiffusionModel {
@@ -87,6 +88,10 @@ struct UNetModel : public DiffusionModel {
8788
unet.set_flash_attention_enabled(enabled);
8889
}
8990

91+
void set_circular_axes(bool circular_x, bool circular_y) override {
92+
unet.set_circular_axes(circular_x, circular_y);
93+
}
94+
9095
bool compute(int n_threads,
9196
DiffusionParams diffusion_params,
9297
struct ggml_tensor** output = nullptr,
@@ -148,6 +153,10 @@ struct MMDiTModel : public DiffusionModel {
148153
mmdit.set_flash_attention_enabled(enabled);
149154
}
150155

156+
void set_circular_axes(bool circular_x, bool circular_y) override {
157+
mmdit.set_circular_axes(circular_x, circular_y);
158+
}
159+
151160
bool compute(int n_threads,
152161
DiffusionParams diffusion_params,
153162
struct ggml_tensor** output = nullptr,
@@ -210,6 +219,10 @@ struct FluxModel : public DiffusionModel {
210219
flux.set_flash_attention_enabled(enabled);
211220
}
212221

222+
void set_circular_axes(bool circular_x, bool circular_y) override {
223+
flux.set_circular_axes(circular_x, circular_y);
224+
}
225+
213226
bool compute(int n_threads,
214227
DiffusionParams diffusion_params,
215228
struct ggml_tensor** output = nullptr,
@@ -277,6 +290,10 @@ struct WanModel : public DiffusionModel {
277290
wan.set_flash_attention_enabled(enabled);
278291
}
279292

293+
void set_circular_axes(bool circular_x, bool circular_y) override {
294+
wan.set_circular_axes(circular_x, circular_y);
295+
}
296+
280297
bool compute(int n_threads,
281298
DiffusionParams diffusion_params,
282299
struct ggml_tensor** output = nullptr,
@@ -303,8 +320,9 @@ struct QwenImageModel : public DiffusionModel {
303320
bool offload_params_to_cpu,
304321
const String2TensorStorage& tensor_storage_map = {},
305322
const std::string prefix = "model.diffusion_model",
306-
SDVersion version = VERSION_QWEN_IMAGE)
307-
: prefix(prefix), qwen_image(backend, offload_params_to_cpu, tensor_storage_map, prefix, version) {
323+
SDVersion version = VERSION_QWEN_IMAGE,
324+
bool zero_cond_t = false)
325+
: prefix(prefix), qwen_image(backend, offload_params_to_cpu, tensor_storage_map, prefix, version, zero_cond_t) {
308326
}
309327

310328
std::string get_desc() override {
@@ -343,6 +361,10 @@ struct QwenImageModel : public DiffusionModel {
343361
qwen_image.set_flash_attention_enabled(enabled);
344362
}
345363

364+
void set_circular_axes(bool circular_x, bool circular_y) override {
365+
qwen_image.set_circular_axes(circular_x, circular_y);
366+
}
367+
346368
bool compute(int n_threads,
347369
DiffusionParams diffusion_params,
348370
struct ggml_tensor** output = nullptr,
@@ -406,6 +428,10 @@ struct ZImageModel : public DiffusionModel {
406428
z_image.set_flash_attention_enabled(enabled);
407429
}
408430

431+
void set_circular_axes(bool circular_x, bool circular_y) override {
432+
z_image.set_circular_axes(circular_x, circular_y);
433+
}
434+
409435
bool compute(int n_threads,
410436
DiffusionParams diffusion_params,
411437
struct ggml_tensor** output = nullptr,

‎otherarch/sdcpp/flux.hpp‎

Lines changed: 23 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -233,14 +233,17 @@ namespace Flux {
233233
__STATIC_INLINE__ struct ggml_tensor* modulate(struct ggml_context* ctx,
234234
struct ggml_tensor* x,
235235
struct ggml_tensor* shift,
236-
struct ggml_tensor* scale) {
236+
struct ggml_tensor* scale,
237+
bool skip_reshape = false) {
237238
// x: [N, L, C]
238239
// scale: [N, C]
239240
// shift: [N, C]
240-
scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
241-
shift = ggml_reshape_3d(ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C]
242-
x = ggml_add(ctx, x, ggml_mul(ctx, x, scale));
243-
x = ggml_add(ctx, x, shift);
241+
if (!skip_reshape) {
242+
scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
243+
shift = ggml_reshape_3d(ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C]
244+
}
245+
x = ggml_add(ctx, x, ggml_mul(ctx, x, scale));
246+
x = ggml_add(ctx, x, shift);
244247
return x;
245248
}
246249

@@ -860,14 +863,14 @@ namespace Flux {
860863
}
861864
}
862865

863-
struct ggml_tensor* pad_to_patch_size(struct ggml_context* ctx,
866+
struct ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx,
864867
struct ggml_tensor* x) {
865868
int64_t W = x->ne[0];
866869
int64_t H = x->ne[1];
867870

868871
int pad_h = (params.patch_size - H % params.patch_size) % params.patch_size;
869872
int pad_w = (params.patch_size - W % params.patch_size) % params.patch_size;
870-
x = ggml_pad(ctx, x, pad_w, pad_h, 0, 0); // [N, C, H + pad_h, W + pad_w]
873+
x = ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
871874
return x;
872875
}
873876

@@ -893,11 +896,11 @@ namespace Flux {
893896
return x;
894897
}
895898

896-
struct ggml_tensor* process_img(struct ggml_context* ctx,
899+
struct ggml_tensor* process_img(GGMLRunnerContext* ctx,
897900
struct ggml_tensor* x) {
898901
// img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size)
899902
x = pad_to_patch_size(ctx, x);
900-
x = patchify(ctx, x);
903+
x = patchify(ctx->ggml_ctx, x);
901904
return x;
902905
}
903906

@@ -1076,7 +1079,7 @@ namespace Flux {
10761079
int pad_h = (patch_size - H % patch_size) % patch_size;
10771080
int pad_w = (patch_size - W % patch_size) % patch_size;
10781081

1079-
auto img = pad_to_patch_size(ctx->ggml_ctx, x);
1082+
auto img = pad_to_patch_size(ctx, x);
10801083
auto orig_img = img;
10811084

10821085
if (params.chroma_radiance_params.use_patch_size_32) {
@@ -1150,16 +1153,16 @@ namespace Flux {
11501153
int pad_h = (patch_size - H % patch_size) % patch_size;
11511154
int pad_w = (patch_size - W % patch_size) % patch_size;
11521155

1153-
auto img = process_img(ctx->ggml_ctx, x);
1156+
auto img = process_img(ctx, x);
11541157
uint64_t img_tokens = img->ne[1];
11551158

11561159
if (params.version == VERSION_FLUX_FILL) {
11571160
GGML_ASSERT(c_concat != nullptr);
11581161
ggml_tensor* masked = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], C, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], 0);
11591162
ggml_tensor* mask = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], 8 * 8, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * C);
11601163

1161-
masked = process_img(ctx->ggml_ctx, masked);
1162-
mask = process_img(ctx->ggml_ctx, mask);
1164+
masked = process_img(ctx, masked);
1165+
mask = process_img(ctx, mask);
11631166

11641167
img = ggml_concat(ctx->ggml_ctx, img, ggml_concat(ctx->ggml_ctx, masked, mask, 0), 0);
11651168
} else if (params.version == VERSION_FLEX_2) {
@@ -1168,21 +1171,21 @@ namespace Flux {
11681171
ggml_tensor* mask = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], 1, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * C);
11691172
ggml_tensor* control = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], C, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * (C + 1));
11701173

1171-
masked = process_img(ctx->ggml_ctx, masked);
1172-
mask = process_img(ctx->ggml_ctx, mask);
1173-
control = process_img(ctx->ggml_ctx, control);
1174+
masked = process_img(ctx, masked);
1175+
mask = process_img(ctx, mask);
1176+
control = process_img(ctx, control);
11741177

11751178
img = ggml_concat(ctx->ggml_ctx, img, ggml_concat(ctx->ggml_ctx, ggml_concat(ctx->ggml_ctx, masked, mask, 0), control, 0), 0);
11761179
} else if (params.version == VERSION_FLUX_CONTROLS) {
11771180
GGML_ASSERT(c_concat != nullptr);
11781181

1179-
auto control = process_img(ctx->ggml_ctx, c_concat);
1182+
auto control = process_img(ctx, c_concat);
11801183
img = ggml_concat(ctx->ggml_ctx, img, control, 0);
11811184
}
11821185

11831186
if (ref_latents.size() > 0) {
11841187
for (ggml_tensor* ref : ref_latents) {
1185-
ref = process_img(ctx->ggml_ctx, ref);
1188+
ref = process_img(ctx, ref);
11861189
img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
11871190
}
11881191
}
@@ -1472,6 +1475,8 @@ namespace Flux {
14721475
increase_ref_index,
14731476
flux_params.ref_index_scale,
14741477
flux_params.theta,
1478+
circular_y_enabled,
1479+
circular_x_enabled,
14751480
flux_params.axes_dim);
14761481
int pos_len = pe_vec.size() / flux_params.axes_dim_sum / 2;
14771482
// LOG_DEBUG("pos_len %d", pos_len);

0 commit comments

Comments
 (0)