Skip to content

How to Contribute Code

Robin Huang edited this page Jun 29, 2024 · 8 revisions

Adding support for new models

Here are a few tips to make reviewing your PR easier for us:

  • Have a minimal implementation of the model code that only depends on pytorch under a license compatible with the GPL license that ComfyUI uses.
  • Provide a reference image with sampling settings/seed/etc.. so that I can make sure the ComfyUI implementation matches the reference one.
  • Replace all attention functions with the comfyui "optimized_attention" attention function
  • If possible, please release your primary models in .safetensor file format.

Example of the sdpa implementation:

def optimized_attention(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False):
    if skip_reshape:
        b, _, _, dim_head = q.shape
    else:
        b, _, dim_head = q.shape
        dim_head //= heads
        q, k, v = map(
            lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2),
            (q, k, v),
        )

    out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
    out = (
        out.transpose(1, 2).reshape(b, -1, heads * dim_head)
    )
    return out

Audio

SDXL Cascade

mmdit

Clone this wiki locally