format : black

Author: cloneofsimo

lora_diffusion/utils.py

@@ -51,16 +51,15 @@
 ]
 
 
-def image_grid(_imgs, rows = None, cols = None):
-    
+def image_grid(_imgs, rows=None, cols=None):
+
     if rows is None and cols is None:
         rows = cols = math.ceil(len(_imgs) ** 0.5)
-    
+
     if rows is None:
         rows = math.ceil(len(_imgs) / cols)
     if cols is None:
         cols = math.ceil(len(_imgs) / rows)
-    
 
     w, h = _imgs[0].size
     grid = Image.new("RGB", size=(cols * w, rows * h))
@@ -176,25 +175,23 @@ def visualize_progress(
     text_sclae=1.0,
     num_inference_steps=50,
     guidance_scale=5.0,
-    offset : int = 0,
-    limit : int = 10,
-    seed : int = 0
+    offset: int = 0,
+    limit: int = 10,
+    seed: int = 0,
 ):
 
-    
     imgs = []
     if isinstance(path_alls, str):
         alls = list(set(glob.glob(path_alls)))
-        
+
         alls.sort(key=os.path.getmtime)
     else:
         alls = path_alls
-    
+
     pipe = StableDiffusionPipeline.from_pretrained(
         model_id, torch_dtype=torch.float16
     ).to(device)
 
-
     print(f"Found {len(alls)} checkpoints")
     for path in alls[offset:limit]:
         print(path)
@@ -207,8 +204,11 @@ def visualize_progress(
         tune_lora_scale(pipe.text_encoder, text_sclae)
 
         torch.manual_seed(seed)
-        image = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
+        image = pipe(
+            prompt,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+        ).images[0]
         imgs.append(image)
 
     return imgs
-    

training_scripts/train_lora_dreambooth.py

@@ -90,20 +90,25 @@ def __init__(
             self.class_prompt = class_prompt
         else:
             self.class_data_root = None
-         
-        img_transforms = []
 
+        img_transforms = []
 
         if resize:
-            img_transforms.append(transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR))
+            img_transforms.append(
+                transforms.Resize(
+                    size, interpolation=transforms.InterpolationMode.BILINEAR
+                )
+            )
         if center_crop:
             img_transforms.append(transforms.CenterCrop(size))
         if color_jitter:
             img_transforms.append(transforms.ColorJitter(0.2, 0.1))
         if h_flip:
             img_transforms.append(transforms.RandomHorizontalFlip())
 
-        self.image_transforms = transforms.Compose([*img_transforms, transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
+        self.image_transforms = transforms.Compose(
+            [*img_transforms, transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]
+        )
 
     def __len__(self):
         return self._length