blocks["layers." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new MLPEmbedder(hidden_dim, hidden_dim));

blocks["norms." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new RMSNorm(hidden_dim));

}

}

struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* timestep) {

// Implement forward pass based on the pseudo-code in the plan (Phase 2.2)

return norm->forward(ctx, x);

}

};

blocks["norm"] = std::shared_ptr<GGMLBlock>(new QKNorm(head_dim));

blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim));

}

std::vector<struct ggml_tensor*> pre_attention(struct ggml_context* ctx, struct ggml_tensor* x) {

auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);

auto norm = std::dynamic_pointer_cast<QKNorm>(blocks["norm"]);

auto v = ggml_view_3d(ctx, qkv_split, dim, L, N, qkv_split->nb[1], qkv_split->nb[2], offset * 2); // [dim, L, N]

// Reshape q, k, v for QKNorm and ggml_nn_attention_ext

// Let's reshape q, k, v to [d_head, L, N*n_head] and [d_head, L, n_head, N] respectively

auto q_reshaped = ggml_reshape_3d(ctx, q, head_dim, L, N * num_heads); // [dim, L, N] -> [d_head, L, N*n_head]

// Modulation block is created and called in the forward pass based on the plan

}

struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* pe, struct ModulationOut& mod, struct ggml_tensor* attn_mask = NULL) {

// x: [N, L, hidden_size]

// pe: Positional embeddings (shape TBD)

blocks["txt_mlp.2"] = std::shared_ptr<GGMLBlock>(new Linear(mlp_hidden_dim, hidden_size));

}

std::pair<struct ggml_tensor*, struct ggml_tensor*> forward(struct ggml_context* ctx, struct ggml_tensor* img, struct ggml_tensor* txt, struct ggml_tensor* pe, const std::vector<ModulationOut>& img_mods, const std::vector<ModulationOut>& txt_mods, struct ggml_tensor* attn_mask = NULL) {

auto img_norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm1"]);

auto img_attn = std::dynamic_pointer_cast<SelfAttention>(blocks["img_attn"]);

}

struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* shift, struct ggml_tensor* scale) {

auto norm_final = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_final"]);

: hidden_size(hidden_size), num_heads(num_heads), mlp_ratio(mlp_ratio),

depth(depth), single_depth(single_depth), in_channels(in_channels),

out_channels(out_channels), flash_attn(flash_attn) {

blocks["img_in"] = std::shared_ptr<GGMLBlock>(new Linear(in_channels, hidden_size, true));

blocks["txt_in"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, hidden_size, true)); // T5 embeddings are already hidden_size

blocks["final_layer"] = std::shared_ptr<GGMLBlock>(new LastLayer_ggml(hidden_size, 1, out_channels)); // patch_size is 1 for Chroma

}

struct ggml_tensor* forward(struct ggml_context* ctx,

struct ggml_tensor* img_latent,

struct ggml_tensor* pe,

struct ggml_tensor* t5_padding_mask,

std::vector<int> skip_layers = std::vector<int>()) {

auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);

auto txt_in = std::dynamic_pointer_cast<Linear>(blocks["txt_in"]);

auto final_layer = std::dynamic_pointer_cast<LastLayer_ggml>(blocks["final_layer"]);

// This means approx_output needs to be split into one set of (scale, shift, gate).

// This is inconsistent with the DoubleStreamBlock.

// Re-reading the plan:

// Approximator output: `conditioning signal (Tensor of shape [batch_size, 1, out_dim])`, likely split into scale, shift, gate.

// SingleStreamBlock: `mod = vec`

// Let's assume `approx_output` is the `vec` and it contains all necessary modulation parameters.

// And the `Modulation` class will be used to extract them.

// Let's assume `approx_output` is `hidden_size * 6` (for 2 sets of modulations).

// And `SingleStreamBlock` will use the first set.

ChromaRunner(

ggml_backend_t backend,

std::map<std::string, enum ggml_type>& tensor_types,

) :

GGMLRunner(backend),

chroma_params({

chroma_params.in_channels, chroma_params.out_channels, chroma_params.flash_attn

)

{

}

std::string get_desc() override {

std::map<std::string, enum ggml_type>& tensor_types,

SDVersion version = VERSION_CHROMA,

bool flash_attn = false)

}

void alloc_params_buffer() {

struct ggml_tensor* dummy = ggml_get_tensor(ctx_meta_, name.c_str());

size_t offset = data_offset + gguf_get_tensor_offset(ctx_gguf_, i);

// LOG_DEBUG("%s", name.c_str());

TensorStorage tensor_storage(prefix + name, dummy->type, dummy->ne, ggml_n_dims(dummy), file_index, offset);

#define found_family (is_xl || is_flux)

for (auto& tensor_storage : tensor_storages) {

if (!found_family) {

return VERSION_CHROMA;

}

if (tensor_storage.name.find("model.diffusion_model.double_blocks.") != std::string::npos) {