@@ -603,7 +603,43 @@ namespace sd::ggml_graph_cut {
603603 GGML_ASSERT (gf != nullptr );
604604 GGML_ASSERT (graph_ctx_out != nullptr );
605605
606- const size_t graph_size = segment.internal_node_indices .size () + segment.input_refs .size () + 8 ;
606+ // Collect leaf inputs and internal nodes, then any tensor they
607+ // reference that is not already represented, notably the view_src of a
608+ // view-typed input leaf. ggml_gallocr sizes its hash set from
609+ // n_nodes + n_leafs (plus a 25% margin that rounds down to zero for a
610+ // one-node segment), so every distinct tensor it will hash must be
611+ // counted here or a tiny segment overflows the hash set and aborts.
612+ std::vector<ggml_tensor*> leaves;
613+ std::unordered_set<ggml_tensor*> represented;
614+ for (const auto & input : segment.input_refs ) {
615+ ggml_tensor* current_input = input_tensor (gf, input);
616+ if (current_input == nullptr ) {
617+ continue ;
618+ }
619+ if (represented.insert (current_input).second ) {
620+ leaves.push_back (current_input);
621+ }
622+ }
623+ for (int node_idx : segment.internal_node_indices ) {
624+ represented.insert (ggml_graph_node (gf, node_idx));
625+ }
626+ auto add_reference = [&](ggml_tensor* tensor) {
627+ if (tensor != nullptr && represented.insert (tensor).second ) {
628+ leaves.push_back (tensor);
629+ }
630+ };
631+ for (int node_idx : segment.internal_node_indices ) {
632+ ggml_tensor* node = ggml_graph_node (gf, node_idx);
633+ for (int src_idx = 0 ; src_idx < GGML_MAX_SRC ; ++src_idx) {
634+ add_reference (node->src [src_idx]);
635+ }
636+ add_reference (node->view_src );
637+ }
638+ for (size_t i = 0 ; i < leaves.size (); ++i) {
639+ add_reference (leaves[i]->view_src );
640+ }
641+
642+ const size_t graph_size = segment.internal_node_indices .size () + leaves.size () + 8 ;
607643 ggml_init_params params = {
608644 /* .mem_size =*/ ggml_graph_overhead_custom (graph_size, false ) + 1024 ,
609645 /* .mem_buffer =*/ nullptr ,
@@ -614,13 +650,9 @@ namespace sd::ggml_graph_cut {
614650 ggml_cgraph* segment_graph = ggml_new_graph_custom (graph_ctx, graph_size, false );
615651 GGML_ASSERT (segment_graph != nullptr );
616652
617- for (const auto & input : segment.input_refs ) {
618- ggml_tensor* current_input = input_tensor (gf, input);
619- if (current_input == nullptr ) {
620- continue ;
621- }
653+ for (ggml_tensor* leaf : leaves) {
622654 GGML_ASSERT (segment_graph->n_leafs < segment_graph->size );
623- segment_graph->leafs [segment_graph->n_leafs ++] = current_input ;
655+ segment_graph->leafs [segment_graph->n_leafs ++] = leaf ;
624656 }
625657
626658 for (int output_node_index : segment.output_node_indices ) {
0 commit comments