batch inputs for compute_clip_text_embedding

[piercus]

Context

I need this for #165, to do GPU efficient evaluation (batch_size > 1).
This is the extension of #213
This is also a baby step in the context of #255 @limiteinductive

One question

In compute_clip_image_embedding there is a concat_batches bool, i'm shared if we need to put it in compute_clip_text_embedding or not.

2 contradictory POV

• I did not put it now cause the functionnal use case is not clear,
• For convergence of inputs formats (in the context of SD : classifier free guidance tensor shape #255 SD refactoring), we might want to tend to uniformity here.

[piercus]

@limiteinductive just added corresponding changes + unit tests for SDXL

[deltheil]

In compute_clip_image_embedding there is a concat_batches bool, i'm shared if we need to put it in compute_clip_text_embedding or not.

@piercus no, you can ignore it IMO. It has been added mainly for documentation purposes, in the context of IP-Adapter with multiple image prompts (see #218). And in this context, you want to create a longer sequence of tokens (like double the size for two images).

[limiteinductive]

@deltheil lgtm

[deltheil]

@deltheil lgtm

Thanks! @piercus will do the final round shortly, stay tuned

[deltheil]

Some comments/suggestions, please take a look

[piercus]

The batch-stability on GPU suffers from the below effect

Source : Torch Double Conv2d

In the above example, a pure pytorch double Conv2d leads to 2e-6 l1 norm diff

import torch
from torch.nn import Conv2d 
device = "cuda:0"

def distance (x: torch.Tensor, y: torch.Tensor) -> float:
    return torch.max((x - y).abs()).item()

with torch.no_grad():
    torch.cuda.manual_seed_all(0)
    x_b2 = torch.randn(2, 4, 32, 32).to(device)
    conv2d_1 = Conv2d(in_channels=4, out_channels=320, kernel_size=3, padding=1, device=device)
    conv2d_2 = Conv2d(in_channels=320, out_channels=640, kernel_size=3, padding=1, device=device)

    output_b2 = conv2d_2(conv2d_1(x_b2))
    output_b1 = conv2d_2(conv2d_1(x_b2[0:1]))

    print(distance(output_b2[0], output_b1[0]))

will output 2e-06

Amplification

This batch discrepancy is then amplified by the following layers

Level	Error max((x - y).abs())
Double Conv	2e-6
Conv + ResidualBlock	5e-5
DownBlocks	2e-3
Unet	1e-3
SD	2e-3

Details of the analysis can be found in https://gist.github.com/piercus/07d03f258907542d312c0c735445e793

Result

As a results the batch stability with torch.allclose is only calculated with tolerance of 5e-3

[deltheil]

See final nits. Please squash the extra commit afterwards. Thanks!

[piercus]

@deltheil please review