Switched ip adapter dataloader to clip_image paths so the control paths can be used for training assistant adapters while training ip adapters

toolkit/config_modules.py

@@ -421,6 +421,9 @@ class DatasetConfig:
         self.caption_type = self.caption_ext
         self.guidance_type: GuidanceType = kwargs.get('guidance_type', 'targeted')
 
+        # ip adapter / reference dataset
+        self.clip_image_path: str = kwargs.get('clip_image_path', None)  # depth maps, etc
+
 
 def preprocess_dataset_raw_config(raw_config: List[dict]) -> List[dict]:
     """

toolkit/data_transfer_object/data_loader.py

@@ -8,7 +8,7 @@ from PIL.ImageOps import exif_transpose
 from toolkit import image_utils
 from toolkit.dataloader_mixins import CaptionProcessingDTOMixin, ImageProcessingDTOMixin, LatentCachingFileItemDTOMixin, \
     ControlFileItemDTOMixin, ArgBreakMixin, PoiFileItemDTOMixin, MaskFileItemDTOMixin, AugmentationFileItemDTOMixin, \
-    UnconditionalFileItemDTOMixin
+    UnconditionalFileItemDTOMixin, ClipImageFileItemDTOMixin
 
 if TYPE_CHECKING:
     from toolkit.config_modules import DatasetConfig
@@ -28,6 +28,7 @@ class FileItemDTO(
     CaptionProcessingDTOMixin,
     ImageProcessingDTOMixin,
     ControlFileItemDTOMixin,
+    ClipImageFileItemDTOMixin,
     MaskFileItemDTOMixin,
     AugmentationFileItemDTOMixin,
     UnconditionalFileItemDTOMixin,
@@ -71,6 +72,7 @@ class FileItemDTO(
         self.tensor = None
         self.cleanup_latent()
         self.cleanup_control()
+        self.cleanup_clip_image()
         self.cleanup_mask()
         self.cleanup_unconditional()
 
@@ -83,6 +85,7 @@ class DataLoaderBatchDTO:
             self.tensor: Union[torch.Tensor, None] = None
             self.latents: Union[torch.Tensor, None] = None
             self.control_tensor: Union[torch.Tensor, None] = None
+            self.clip_image_tensor: Union[torch.Tensor, None] = None
             self.mask_tensor: Union[torch.Tensor, None] = None
             self.unaugmented_tensor: Union[torch.Tensor, None] = None
             self.unconditional_tensor: Union[torch.Tensor, None] = None
@@ -113,6 +116,21 @@ class DataLoaderBatchDTO:
                         control_tensors.append(x.control_tensor)
                 self.control_tensor = torch.cat([x.unsqueeze(0) for x in control_tensors])
 
+            if any([x.clip_image_tensor is not None for x in self.file_items]):
+                # find one to use as a base
+                base_clip_image_tensor = None
+                for x in self.file_items:
+                    if x.clip_image_tensor is not None:
+                        base_clip_image_tensor = x.clip_image_tensor
+                        break
+                clip_image_tensors = []
+                for x in self.file_items:
+                    if x.clip_image_tensor is None:
+                        clip_image_tensors.append(torch.zeros_like(base_clip_image_tensor))
+                    else:
+                        clip_image_tensors.append(x.clip_image_tensor)
+                self.clip_image_tensor = torch.cat([x.unsqueeze(0) for x in clip_image_tensors])
+
             if any([x.mask_tensor is not None for x in self.file_items]):
                 # find one to use as a base
                 base_mask_tensor = None

toolkit/dataloader_mixins.py

@@ -350,6 +350,8 @@ class ImageProcessingDTOMixin:
             self.get_latent()
             if self.has_control_image:
                 self.load_control_image()
+            if self.has_clip_image:
+                self.load_clip_image()
             if self.has_mask_image:
                 self.load_mask_image()
             if self.has_unconditional:
@@ -443,6 +445,8 @@ class ImageProcessingDTOMixin:
         if not only_load_latents:
             if self.has_control_image:
                 self.load_control_image()
+            if self.has_clip_image:
+                self.load_clip_image()
             if self.has_mask_image:
                 self.load_mask_image()
             if self.has_unconditional:
@@ -523,6 +527,46 @@ class ControlFileItemDTOMixin:
         self.control_tensor = None
 
 
+class ClipImageFileItemDTOMixin:
+    def __init__(self: 'FileItemDTO', *args, **kwargs):
+        if hasattr(super(), '__init__'):
+            super().__init__(*args, **kwargs)
+        self.has_clip_image = False
+        self.clip_image_path: Union[str, None] = None
+        self.clip_image_tensor: Union[torch.Tensor, None] = None
+        dataset_config: 'DatasetConfig' = kwargs.get('dataset_config', None)
+        if dataset_config.clip_image_path is not None:
+            # find the control image path
+            clip_image_path = dataset_config.clip_image_path
+            # we are using control images
+            img_path = kwargs.get('path', None)
+            img_ext_list = ['.jpg', '.jpeg', '.png', '.webp']
+            file_name_no_ext = os.path.splitext(os.path.basename(img_path))[0]
+            for ext in img_ext_list:
+                if os.path.exists(os.path.join(clip_image_path, file_name_no_ext + ext)):
+                    self.clip_image_path = os.path.join(clip_image_path, file_name_no_ext + ext)
+                    self.has_clip_image = True
+                    break
+
+    def load_clip_image(self: 'FileItemDTO'):
+        img = Image.open(self.clip_image_path).convert('RGB')
+        try:
+            img = exif_transpose(img)
+        except Exception as e:
+            print(f"Error: {e}")
+            print(f"Error loading image: {self.clip_image_path}")
+
+        # we just scale them to 512x512:
+        img = img.resize((512, 512), Image.BICUBIC)
+
+        self.clip_image_tensor = transforms.ToTensor()(img)
+
+    def cleanup_clip_image(self: 'FileItemDTO'):
+        self.clip_image_tensor = None
+
+
+
+
 class AugmentationFileItemDTOMixin:
     def __init__(self: 'FileItemDTO', *args, **kwargs):
         if hasattr(super(), '__init__'):

toolkit/ip_adapter.py

@@ -153,7 +153,10 @@ class IPAdapter(torch.nn.Module):
         super().__init__()
         self.config = adapter_config
         self.sd_ref: weakref.ref = weakref.ref(sd)
-        self.clip_image_processor = CLIPImageProcessor.from_pretrained(adapter_config.image_encoder_path)
+        try:
+            self.clip_image_processor = CLIPImageProcessor.from_pretrained(adapter_config.image_encoder_path)
+        except EnvironmentError:
+            self.clip_image_processor = CLIPImageProcessor()
         self.device = self.sd_ref().unet.device
         self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(adapter_config.image_encoder_path,
                                                                            ignore_mismatched_sizes=True)