Added functional TI strategy

1 files changed, 118 insertions, 0 deletions

diff --git a/training/functional.py b/training/functional.py
index 1f2ca6d..e54c9c8 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -2,6 +2,8 @@ import math
 from contextlib import _GeneratorContextManager, nullcontext
 from typing import Callable, Any, Tuple, Union, Optional
 from functools import partial
+from pathlib import Path
+import itertools
 
 import torch
 import torch.nn.functional as F
@@ -26,6 +28,14 @@ def const(result=None):
     return fn
 
 
+def make_grid(images, rows, cols):
+    w, h = images[0].size
+    grid = Image.new('RGB', size=(cols*w, rows*h))
+    for i, image in enumerate(images):
+        grid.paste(image, box=(i % cols*w, i//cols*h))
+    return grid
+
+
 def get_models(pretrained_model_name_or_path: str):
     tokenizer = MultiCLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder='tokenizer')
     text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder')
@@ -40,6 +50,107 @@ def get_models(pretrained_model_name_or_path: str):
     return tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler, embeddings
 
 
+def save_samples(
+    accelerator: Accelerator,
+    unet: UNet2DConditionModel,
+    text_encoder: CLIPTextModel,
+    tokenizer: MultiCLIPTokenizer,
+    vae: AutoencoderKL,
+    sample_scheduler: DPMSolverMultistepScheduler,
+    train_dataloader: DataLoader,
+    val_dataloader: DataLoader,
+    dtype: torch.dtype,
+    output_dir: Path,
+    seed: int,
+    step: int,
+    batch_size: int = 1,
+    num_batches: int = 1,
+    num_steps: int = 20,
+    guidance_scale: float = 7.5,
+    image_size: Optional[int] = None,
+):
+    print(f"Saving samples for step {step}...")
+
+    samples_path = output_dir.joinpath("samples")
+
+    grid_cols = min(batch_size, 4)
+    grid_rows = (num_batches * batch_size) // grid_cols
+
+    unet = accelerator.unwrap_model(unet)
+    text_encoder = accelerator.unwrap_model(text_encoder)
+
+    orig_unet_dtype = unet.dtype
+    orig_text_encoder_dtype = text_encoder.dtype
+
+    unet.to(dtype=dtype)
+    text_encoder.to(dtype=dtype)
+
+    pipeline = VlpnStableDiffusion(
+        text_encoder=text_encoder,
+        vae=vae,
+        unet=unet,
+        tokenizer=tokenizer,
+        scheduler=sample_scheduler,
+    ).to(accelerator.device)
+    pipeline.set_progress_bar_config(dynamic_ncols=True)
+
+    generator = torch.Generator(device=accelerator.device).manual_seed(seed)
+
+    for pool, data, gen in [
+        ("stable", val_dataloader, generator),
+        ("val", val_dataloader, None),
+        ("train", train_dataloader, None)
+    ]:
+        all_samples = []
+        file_path = samples_path.joinpath(pool, f"step_{step}.jpg")
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        batches = list(itertools.islice(itertools.cycle(data), batch_size * num_batches))
+        prompt_ids = [
+            prompt
+            for batch in batches
+            for prompt in batch["prompt_ids"]
+        ]
+        nprompt_ids = [
+            prompt
+            for batch in batches
+            for prompt in batch["nprompt_ids"]
+        ]
+
+        for i in range(num_batches):
+            start = i * batch_size
+            end = (i + 1) * batch_size
+            prompt = prompt_ids[start:end]
+            nprompt = nprompt_ids[start:end]
+
+            samples = pipeline(
+                prompt=prompt,
+                negative_prompt=nprompt,
+                height=image_size,
+                width=image_size,
+                generator=gen,
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_steps,
+                output_type='pil'
+            ).images
+
+            all_samples += samples
+
+        image_grid = make_grid(all_samples, grid_rows, grid_cols)
+        image_grid.save(file_path, quality=85)
+
+    unet.to(dtype=orig_unet_dtype)
+    text_encoder.to(dtype=orig_text_encoder_dtype)
+
+    del unet
+    del text_encoder
+    del generator
+    del pipeline
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
 def generate_class_images(
     accelerator: Accelerator,
     text_encoder: CLIPTextModel,
@@ -109,6 +220,10 @@ def get_models(pretrained_model_name_or_path: str):
 
     embeddings = patch_managed_embeddings(text_encoder)
 
+    vae.requires_grad_(False)
+    unet.requires_grad_(False)
+    text_encoder.requires_grad_(False)
+
     return tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler, embeddings
 
 
@@ -427,6 +542,9 @@ def train(
         seed,
     )
 
+    if accelerator.is_main_process:
+        accelerator.init_trackers("textual_inversion")
+
     train_loop(
         accelerator=accelerator,
         optimizer=optimizer,