Fix Value dimension in ImageCrossAttention

[hugojarkoff]

This is a minor issue, but while reading the codebase I noticed that the ImageCrossAttention uses the wrong in_features dimension in the second fl.Linear layer:

class ImageCrossAttention(fl.Chain):
    def __init__(self, text_cross_attention: fl.Attention, scale: float = 1.0) -> None:
            ...      
                    # This first Linear corresponds to the keys K' 
                    fl.Linear(
                      ...
                    ),
                    ...
                    # This second Linear corresponds to the values V'
                    fl.Linear(
                        in_features=text_cross_attention.key_embedding_dim,
                        out_features=text_cross_attention.inner_dim,
                        bias=text_cross_attention.use_bias,
                        device=text_cross_attention.device,
                        dtype=text_cross_attention.dtype,
                    ),
                ),
            ...

Should (IIUC) be changed to :

class ImageCrossAttention(fl.Chain):
    def __init__(self, text_cross_attention: fl.Attention, scale: float = 1.0) -> None:
            ...      
                    # This first Linear corresponds to the keys K' 
                    fl.Linear(
                      ...
                    ),
                    ...
                    # This second Linear corresponds to the values V'
                    fl.Linear(
                        in_features=text_cross_attention.value_embedding_dim,
                        out_features=text_cross_attention.inner_dim,
                        bias=text_cross_attention.use_bias,
                        device=text_cross_attention.device,
                        dtype=text_cross_attention.dtype,
                    ),
                ),
            ...

In practice, and IINM, this shouldn't change anything in the context of Image Cross-Attention (since both key and query dim are the same).