Added support for vision direct adapter for flux

2024-08-26 16:27:28 -06:00
parent e127c079da
commit 3843e0d148
3 changed files with 268 additions and 41 deletions
--- a/toolkit/custom_adapter.py
+++ b/toolkit/custom_adapter.py
@@ -852,6 +852,9 @@ class CustomAdapter(torch.nn.Module):
                    if self.adapter_type == 'te_augmenter':
                        clip_image_embeds = self.te_augmenter(clip_image_embeds)

+                    if self.adapter_type == 'vision_direct':
+                        clip_image_embeds = self.vd_adapter(clip_image_embeds)
+
                    # save them to the conditional and unconditional
                    try:
                        self.unconditional_embeds, self.conditional_embeds = clip_image_embeds.chunk(2, dim=0)