.gitignore

@@ -8,4 +8,5 @@ wandb
 exps*
 .vscode
 build
-lora_diffusion.egg-info
+lora_diffusion.egg-info
+training_batch_preview

lora_diffusion/cli_lora_add.py

@@ -1,19 +1,38 @@
-from typing import Literal, Union, Dict
+import sys
+if sys.version_info >= (3,8):
+    from typing import Literal
+else : 
+    from typing_extensions import Literal
+from typing import Union, Dict
 import os
 import shutil
 import fire
 from diffusers import StableDiffusionPipeline
 from safetensors.torch import safe_open, save_file
 
 import torch
-from .lora import (
-    tune_lora_scale,
-    patch_pipe,
-    collapse_lora,
-    monkeypatch_remove_lora,
-)
-from .lora_manager import lora_join
-from .to_ckpt_v2 import convert_to_ckpt
+
+try:
+    from .lora import (
+        tune_lora_scale,
+        patch_pipe,
+        collapse_lora,
+        monkeypatch_remove_lora,
+    )
+
+    from .lora_manager import lora_join
+    from .to_ckpt_v2 import convert_to_ckpt
+
+except:  # allows running the repo without installing it (can mess up existing dependencies)
+    from lora_diffusion import (
+        tune_lora_scale,
+        patch_pipe,
+        collapse_lora,
+        monkeypatch_remove_lora,
+    )
+
+    from lora_diffusion.lora_manager import lora_join
+    from lora_diffusion.to_ckpt_v2 import convert_to_ckpt
 
 
 def _text_lora_path(path: str) -> str:
@@ -185,3 +204,7 @@ def add(
 
 def main():
     fire.Fire(add)
+
+
+if __name__ == "__main__":
+    main()

lora_diffusion/cli_pt_to_safetensors.py

@@ -62,9 +62,11 @@ def convert(*paths, outpath, overwrite=False, **settings):
             }
 
             prefix = f"{name}."
-
-            arg_settings = { k[len(prefix) :]: v for k, v in settings.items() if k.startswith(prefix) }
-            model_settings = { **model_settings, **arg_settings }
+
+            arg_settings = {
+                k[len(prefix) :]: v for k, v in settings.items() if k.startswith(prefix)
+            }
+            model_settings = {**model_settings, **arg_settings}
 
             print(f"Loading Lora for {name} from {path} with settings {model_settings}")
 

lora_diffusion/dataset.py

@@ -1,7 +1,7 @@
 import random
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple, Union
-
+import numpy as np
 from PIL import Image
 from torch import zeros_like
 from torch.utils.data import Dataset
@@ -39,7 +39,38 @@
     "a photo of a small {}",
 ]
 
-STYLE_TEMPLATE = [
+PERSON_TEMPLATE = [
+    "{}",
+    "{}",
+    "a picture of {}",
+    "a closeup of {}",
+    "a closeup photo of {}",
+    "a close-up picture of {}",
+    "a photo of {}",
+    "a photo of {}",
+    "the photo of {}",
+    "a cropped photo of {}",
+    "a funny photo of {}",
+    "a selfie of {}",
+    "a photo of the handsome {}",
+    "a photo of the beautiful {}",
+    "a selfie taken by the handsome {}",
+    "a selfie taken by {}",
+    "{} taking a selfie",
+    "{} is having fun, 4k photograph",
+    "{} wearing a plaidered shirt standing next to another person",
+    "smiling {} in a hoodie and sweater",
+    "a photo of the cool {}",
+    "a close-up photo of {}",
+    "a bright photo of {}",
+    "a cropped photo of {}",
+    "a brilliant HD photo of {}",
+    "a beautiful picture of {}",
+    "a photo showing {}",
+    "a great photo of {}",
+]
+
+STYLE_TEMPLATE_ORIG = [
     "a painting in the style of {}",
     "a rendering in the style of {}",
     "a cropped painting in the style of {}",
@@ -61,10 +92,28 @@
     "a large painting in the style of {}",
 ]
 
+STYLE_TEMPLATE = [
+    "a painting in the style of {}",
+    "a rendering in the style of {}",
+    "an artwork in the style of {}",
+    "a magnificent painting in the style of {}",
+    "a picture in the style of {}",
+    "a photograph, {} style",
+    "{} style painting",
+    "a {}-styled artwork",
+    "a nice painting in the style of {}",
+    "a goregous example of {} style",
+    "image in the style of {}",
+    "{}, painting",
+    "{} artwork",
+]
+
+
 NULL_TEMPLATE = ["{}"]
 
 TEMPLATE_MAP = {
     "object": OBJECT_TEMPLATE,
+    "person": PERSON_TEMPLATE,
     "style": STYLE_TEMPLATE,
     "null": NULL_TEMPLATE,
 }
@@ -116,6 +165,35 @@ def _generate_random_mask(image):
     return mask, masked_image
 
 
+def expand_rectangle(mask, f):
+    rows, cols = np.where(mask == 255)
+    top_row, bottom_row = np.min(rows), np.max(rows)
+    left_col, right_col = np.min(cols), np.max(cols)
+
+    rect_height, rect_width = bottom_row - top_row + 1, right_col - left_col + 1
+    new_height, new_width = np.round(rect_height * f), np.round(rect_width * f)
+
+    center_row, center_col = top_row + rect_height // 2, left_col + rect_width // 2
+    top_row, bottom_row = np.round(center_row - new_height / 2), np.round(
+        center_row + new_height / 2
+    )
+    left_col, right_col = np.round(center_col - new_width / 2), np.round(
+        center_col + new_width / 2
+    )
+
+    top_row, bottom_row = int(np.clip(top_row, 0, mask.shape[0] - 1)), int(
+        np.clip(bottom_row, 0, mask.shape[0] - 1)
+    )
+    left_col, right_col = int(np.clip(left_col, 0, mask.shape[1] - 1)), int(
+        np.clip(right_col, 0, mask.shape[1] - 1)
+    )
+
+    expanded_mask = np.ones_like(mask)
+    expanded_mask[top_row : bottom_row + 1, left_col : right_col + 1] = 255
+
+    return expanded_mask
+
+
 class PivotalTuningDatasetCapation(Dataset):
     """
     A dataset to prepare the instance and class images with the prompts for fine-tuning the model.
@@ -141,6 +219,8 @@ def __init__(
         self.tokenizer = tokenizer
         self.resize = resize
         self.train_inpainting = train_inpainting
+        self.h_flip_prob = 0.5
+        self.final_flip_prob = 0.33 if use_template == "person" else 0.5
 
         instance_data_root = Path(instance_data_root)
         if not instance_data_root.exists():
@@ -156,6 +236,10 @@ def __init__(
         # Prepare the instance images
         if use_mask_captioned_data:
             src_imgs = glob.glob(str(instance_data_root) + "/*src.jpg")
+            src_imgs = sorted(
+                src_imgs, key=lambda x: int(str(Path(x).stem).split(".")[0])
+            )
+
             for f in src_imgs:
                 idx = int(str(Path(f).stem).split(".")[0])
                 mask_path = f"{instance_data_root}/{idx}.mask.png"
@@ -218,6 +302,18 @@ def __init__(
                         ]
                     )
                     for idx, mask in enumerate(masks):
+                        avg_pixel_value = np.array(mask.getdata()).mean()
+                        if avg_pixel_value == 1.0:
+                            print(f"No mask detected for {idx}..")
+                        else:
+                            if 1:
+                                # convert to numpy array:
+                                mask = np.array(mask)
+                                # Make the rectangular mask region bigger:
+                                mask = expand_rectangle(mask, 1.25)
+                                # convert back to PIL image:
+                                mask = Image.fromarray(mask).convert("L")
+
                         mask.save(f"{instance_data_root}/{idx}.mask.png")
 
                     break
@@ -237,12 +333,13 @@ def __init__(
         self.h_flip = h_flip
         self.image_transforms = transforms.Compose(
             [
+                transforms.RandomAffine(degrees=0, translate=(0, 0), scale=(1.0, 1.2)),
                 transforms.Resize(
                     size, interpolation=transforms.InterpolationMode.BILINEAR
                 )
                 if resize
                 else transforms.Lambda(lambda x: x),
-                transforms.ColorJitter(0.1, 0.1)
+                transforms.ColorJitter(0.1, 0.1, 0.02, 0.02)
                 if color_jitter
                 else transforms.Lambda(lambda x: x),
                 transforms.CenterCrop(size),
@@ -253,6 +350,15 @@ def __init__(
 
         self.blur_amount = blur_amount
 
+        print("Captions:")
+        print(self.captions)
+
+    def tune_h_flip_prob(self, training_progress):
+        if self.h_flip:
+            # Tune the h_flip probability to be 0.5 training_progress is 0 and end_prob when training_progress is 1
+            self.h_flip_prob = 0.5 + (self.final_flip_prob - 0.5) * training_progress
+            print(f"h_flip_prob: {self.h_flip_prob:.3f}")
+
     def __len__(self):
         return self._length
 
@@ -283,18 +389,14 @@ def __getitem__(self, index):
                 for token, value in self.token_map.items():
                     text = text.replace(token, value)
 
-        print(text)
+        if random.random() < 0.1:
+            print(text)
 
         if self.use_mask:
-            example["mask"] = (
-                self.image_transforms(
-                    Image.open(self.mask_path[index % self.num_instance_images])
-                )
-                * 0.5
-                + 1.0
-            )
+            img_mask = Image.open(self.mask_path[index % self.num_instance_images])
+            example["mask"] = self.image_transforms(img_mask) * 0.5 + 1.0
 
-        if self.h_flip and random.random() > 0.5:
+        if self.h_flip and random.random() < self.h_flip_prob:
             hflip = transforms.RandomHorizontalFlip(p=1)
 
             example["instance_images"] = hflip(example["instance_images"])

lora_diffusion/lora.py

@@ -1,7 +1,12 @@
 import json
 import math
 from itertools import groupby
-from typing import Callable, Dict, List, Optional, Set, Tuple, Type, Union
+import sys
+if sys.version_info >= (3,9):
+    from typing import Type
+else : 
+    from typing_extensions import Type
+from typing import Callable, Dict, List, Optional, Set, Tuple, Union
 
 import numpy as np
 import PIL
@@ -914,7 +919,7 @@ def apply_learned_embed_in_clip(
         trained_tokens = list(learned_embeds.keys())
 
     for token in trained_tokens:
-        print(token)
+        print("Adding new token: ", token)
         embeds = learned_embeds[token]
 
         # cast to dtype of text_encoder

lora_diffusion/preprocess_files.py

@@ -2,7 +2,12 @@
 # Have BLIP auto caption
 # Have CLIPSeg auto mask concept
 
-from typing import List, Literal, Union, Optional, Tuple
+import sys
+if sys.version_info >= (3,8):
+    from typing import Literal
+else : 
+    from typing_extensions import Literal
+from typing import List, Union, Optional, Tuple
 import os
 from PIL import Image, ImageFilter
 import torch
@@ -244,7 +249,7 @@ def _center_of_mass(mask: Image.Image):
 def load_and_save_masks_and_captions(
     files: Union[str, List[str]],
     output_dir: str,
-    caption_text: Optional[str] = None,
+    captions_text: Optional[Union[List[str], str]] = None,
     target_prompts: Optional[Union[List[str], str]] = None,
     target_size: int = 512,
     crop_based_on_salience: bool = True,
@@ -263,8 +268,10 @@ def load_and_save_masks_and_captions(
         # check if it is a directory
         if os.path.isdir(files):
             # get all the .png .jpg in the directory
-            files = glob.glob(os.path.join(files, "*.png")) + glob.glob(
-                os.path.join(files, "*.jpg")
+            files = (
+                glob.glob(os.path.join(files, "*.png"))
+                + glob.glob(os.path.join(files, "*.jpg"))
+                + glob.glob(os.path.join(files, "*.jpeg"))
             )
 
         if len(files) == 0:
@@ -278,8 +285,10 @@ def load_and_save_masks_and_captions(
     images = [Image.open(file) for file in files]
 
     # captions
-    print(f"Generating {len(images)} captions...")
-    captions = blip_captioning_dataset(images, text=caption_text)
+    captions = caption_text
+    if not isinstance(caption_text, List):
+        print(f"Generating {len(images)} captions...")
+        captions = blip_captioning_dataset(images, text=caption_text)
 
     if target_prompts is None:
         target_prompts = captions
@@ -325,3 +334,7 @@ def load_and_save_masks_and_captions(
 
 def main():
     fire.Fire(load_and_save_masks_and_captions)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -8,3 +8,4 @@ safetensors
 opencv-python
 torchvision
 mediapipe
+typing_extensions; python_version < '3.9'

setup.py

@@ -6,7 +6,7 @@
 setup(
     name="lora_diffusion",
     py_modules=["lora_diffusion"],
-    version="0.1.7",
+    version="0.1.8",
     description="Low Rank Adaptation for Diffusion Models. Works with Stable Diffusion out-of-the-box.",
     author="Simo Ryu",
     packages=find_packages(),