diff --git a/scripts/coach.py b/scripts/coach.py
new file mode 100644
index 0000000..4bf426b
--- /dev/null
+++ b/scripts/coach.py
@@ -0,0 +1,720 @@
+import os
+import random
+import sys
+from pathlib import Path
+
+import accelerate
+import diffusers
+import numpy as np
+import pyrallis
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+import transformers
+from accelerate import Accelerator
+# matplotlib.use('Agg')  # Set the backend to non-interactive (Agg)
+from accelerate.logging import get_logger
+from accelerate.state import AcceleratorState
+from accelerate.utils import ProjectConfiguration, set_seed
+from diffusers import PriorTransformer, DDPMScheduler
+from diffusers import UNet2DConditionModel, KandinskyV22Pipeline
+from diffusers.loaders import AttnProcsLayers
+from diffusers.models.attention_processor import LoRAAttnProcessor
+from diffusers.training_utils import EMAModel
+from packaging import version
+from torchvision import transforms
+from tqdm import tqdm
+from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer, CLIPVisionModelWithProjection
+from transformers.utils import ContextManagers
+
+from model.pipeline_pops import pOpsPipeline
+from model import pops_utils
+from training.dataset import TexturedDataset, InstructDataset, PairsDataset, \
+    SceneDataset, UnionDataset, ClothesDataset
+from training.train_config import TrainConfig
+from utils import vis_utils
+
+logger = get_logger(__name__, log_level="INFO")
+
+
+class Coach:
+    def __init__(self, config: TrainConfig):
+        self.cfg = config
+        self.cfg.output_dir.mkdir(exist_ok=True, parents=True)
+        (self.cfg.output_dir / 'cfg.yaml').write_text(pyrallis.dump(self.cfg))
+        (self.cfg.output_dir / 'run.sh').write_text(f'python {Path(__file__).name} {" ".join(sys.argv)}')
+
+        self.logging_dir = self.cfg.output_dir / 'logs'
+        accelerator_project_config = ProjectConfiguration(
+            total_limit=2, project_dir=self.cfg.output_dir, logging_dir=self.logging_dir
+        )
+        self.accelerator = Accelerator(
+            gradient_accumulation_steps=self.cfg.gradient_accumulation_steps,
+            mixed_precision=self.cfg.mixed_precision,
+            log_with=self.cfg.report_to,
+            project_config=accelerator_project_config,
+        )
+        logger.info(self.accelerator.state, main_process_only=False)
+        if self.accelerator.is_local_main_process:
+            transformers.utils.logging.set_verbosity_warning()
+            diffusers.utils.logging.set_verbosity_info()
+        else:
+            transformers.utils.logging.set_verbosity_error()
+            diffusers.utils.logging.set_verbosity_error()
+
+        if self.cfg.seed is not None:
+            set_seed(self.cfg.seed)
+
+        if self.accelerator.is_main_process:
+            self.logging_dir.mkdir(exist_ok=True, parents=True)
+
+        self.noise_scheduler = DDPMScheduler(beta_schedule='squaredcos_cap_v2', prediction_type='sample')
+        self.image_processor = CLIPImageProcessor.from_pretrained(self.cfg.image_processor_path,
+                                                                  subfolder='image_processor')
+        self.tokenizer = CLIPTokenizer.from_pretrained(self.cfg.tokenizer_path, subfolder='tokenizer')
+
+        def deepspeed_zero_init_disabled_context_manager():
+            """
+            returns either a context list that includes one that will disable zero.Init or an empty context list
+            """
+            deepspeed_plugin = AcceleratorState().deepspeed_plugin if accelerate.state.is_initialized() else None
+            if deepspeed_plugin is None:
+                return []
+
+            return [deepspeed_plugin.zero3_init_context_manager(enable=False)]
+
+        self.weight_dtype = torch.float32
+        if self.accelerator.mixed_precision == "fp16":
+            self.weight_dtype = torch.float16
+        elif self.accelerator.mixed_precision == "bf16":
+            self.weight_dtype = torch.bfloat16
+        with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
+            self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(self.cfg.pretrained_image_encoder,
+                                                                               subfolder='image_encoder',
+                                                                               torch_dtype=self.weight_dtype).eval()
+            self.text_encoder = CLIPTextModelWithProjection.from_pretrained(self.cfg.text_encoder_path,
+                                                                            subfolder='text_encoder',
+                                                                            torch_dtype=self.weight_dtype).eval()
+        print('args.pretrained_prior_path =', self.cfg.pretrained_prior_path)
+        self.prior = PriorTransformer.from_pretrained(
+            self.cfg.pretrained_prior_path, subfolder="prior"
+        )
+
+        self.clip_mean = self.prior.clip_mean.clone()
+        self.clip_std = self.prior.clip_std.clone()
+
+        # Freeze text_encoder and image_encoder
+        self.text_encoder.requires_grad_(False)
+        self.image_encoder.requires_grad_(False)
+
+        # Load full model for vis
+        self.unet = UNet2DConditionModel.from_pretrained('kandinsky-community/kandinsky-2-2-decoder',
+                                                         subfolder='unet').to(torch.float16).to(self.accelerator.device)
+        self.prior_pipeline = pOpsPipeline.from_pretrained(
+            'kandinsky-community/kandinsky-2-2-prior',
+            prior=self.prior,
+            image_encoder=self.image_encoder,
+            torch_dtype=torch.float16)
+        self.prior_pipeline = self.prior_pipeline.to(self.accelerator.device)
+        self.decoder = KandinskyV22Pipeline.from_pretrained('kandinsky-community/kandinsky-2-2-decoder', unet=self.unet,
+                                                            torch_dtype=torch.float16)
+        self.decoder = self.decoder.to(self.accelerator.device)
+
+        if self.cfg.lora_rank is not None:
+            self.prior.requires_grad_(False)
+            self.prior.to(self.accelerator.device, dtype=self.weight_dtype)
+            lora_attn_procs = {}
+            for name in self.prior.attn_processors.keys():
+                lora_attn_procs[name] = LoRAAttnProcessor(hidden_size=2048, rank=self.cfg.lora_rank)
+
+            self.prior.set_attn_processor(lora_attn_procs)
+
+            self.lora_layers = AttnProcsLayers(self.prior.attn_processors)
+
+        # Create EMA for the prior.
+        if self.cfg.use_ema:
+            ema_prior = PriorTransformer.from_pretrained(
+                self.cfg.pretrained_prior_path, subfolder="prior"
+            )
+            ema_prior = EMAModel(ema_prior.parameters(), model_cls=PriorTransformer, model_config=ema_prior.config)
+            ema_prior.to(self.accelerator.device)
+
+        # `accelerate` 0.16.0 will have better support for customized saving
+        if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
+            # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
+            def save_model_hook(models, weights, output_dir):
+                if self.cfg.use_ema:
+                    ema_prior.save_pretrained(os.path.join(output_dir, "prior_ema"))
+
+                for i, model in enumerate(models):
+                    model.save_pretrained(os.path.join(output_dir, "prior"))
+
+                    # make sure to pop weight so that corresponding model is not saved again
+                    weights.pop()
+
+            def load_model_hook(models, input_dir):
+                if self.cfg.use_ema:
+                    load_model = EMAModel.from_pretrained(os.path.join(input_dir, "prior_ema"), PriorTransformer)
+                    ema_prior.load_state_dict(load_model.state_dict())
+                    ema_prior.to(self.accelerator.device)
+                    del load_model
+
+                for i in range(len(models)):
+                    # pop models so that they are not loaded again
+                    model = models.pop()
+
+                    # load diffusers style into model
+                    load_model = PriorTransformer.from_pretrained(input_dir, subfolder="prior")
+                    model.register_to_config(**load_model.config)
+
+                    model.load_state_dict(load_model.state_dict())
+                    del load_model
+
+            self.accelerator.register_save_state_pre_hook(save_model_hook)
+            self.accelerator.register_load_state_pre_hook(load_model_hook)
+
+        if self.cfg.allow_tf32:
+            torch.backends.cuda.matmul.allow_tf32 = True
+
+        if self.cfg.use_8bit_adam:
+            try:
+                import bitsandbytes as bnb
+            except ImportError:
+                raise ImportError(
+                    "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
+                )
+
+            optimizer_cls = bnb.optim.AdamW8bit
+        else:
+            optimizer_cls = torch.optim.AdamW
+
+        if self.cfg.lora_rank is not None:
+            params_to_optimize = list(self.lora_layers.parameters())
+        else:
+            if self.cfg.training_mode == 'full':
+                params_to_optimize = list(self.prior.parameters())
+            elif '_layers' in self.cfg.training_mode:
+                layers_to_train = int(self.cfg.training_mode.split('_')[0])
+                params_to_optimize = []
+                self.prior.norm_out.requires_grad_(False)
+                self.prior.proj_to_clip_embeddings.requires_grad_(False)
+                self.prior.time_proj.requires_grad_(False)
+                self.prior.time_embedding.requires_grad_(False)
+                self.prior.transformer_blocks[layers_to_train:].requires_grad_(False)
+                self.prior.clip_mean.requires_grad_(False)
+                self.prior.clip_std.requires_grad_(False)
+                # Populate params_to_optimize with all params that require grad and print its name
+                for name, param in self.prior.named_parameters():
+                    if param.requires_grad:
+                        print(name)
+                        params_to_optimize.append(param)
+            else:
+                raise ValueError(f"Invalid training mode: {self.cfg.training_mode}")
+                # Count the number of parameters
+        num_params = sum(p.numel() for p in params_to_optimize)
+        # Pretty print the number of parameters
+        print(f"Number of parameters: {num_params}")
+
+        self.optimizer = optimizer_cls(
+            params_to_optimize,
+            lr=self.cfg.lr,
+            betas=(self.cfg.adam_beta1, self.cfg.adam_beta2),
+            weight_decay=self.cfg.adam_weight_decay,
+            eps=self.cfg.adam_epsilon,
+        )
+
+        self.train_dataloader, self.validation_dataloader = self.get_dataloaders()
+
+        if self.cfg.lora_rank is not None:
+            self.lora_layers, self.optimizer, self.train_dataloader = self.accelerator.prepare(
+                self.lora_layers, self.optimizer, self.train_dataloader
+            )
+        else:
+            self.prior, self.optimizer, self.train_dataloader = self.accelerator.prepare(
+                self.prior, self.optimizer, self.train_dataloader
+            )
+
+        self.image_encoder.to(self.accelerator.device, dtype=self.weight_dtype)
+        self.text_encoder.to(self.accelerator.device, dtype=self.weight_dtype)
+        self.clip_mean = self.clip_mean.to(self.weight_dtype).to(self.accelerator.device)
+        self.clip_std = self.clip_std.to(self.weight_dtype).to(self.accelerator.device)
+
+        self.train_step = 0 if self.cfg.resume_from_step is None else self.cfg.resume_from_step
+        print(self.train_step)
+        self.device = "cuda"
+
+        if self.cfg.resume_from_path is not None:
+            prior_state_dict = torch.load(self.cfg.resume_from_path, map_location=self.device)
+            msg = self.prior.load_state_dict(prior_state_dict, strict=False)
+            print(msg)
+
+    def save_model(self, save_path):
+        save_path.parent.mkdir(exist_ok=True, parents=True)
+        prior_state_dict = self.prior.state_dict()
+        del prior_state_dict["clip_mean"]
+        del prior_state_dict["clip_std"]
+        torch.save(prior_state_dict, save_path)
+
+        if self.cfg.lora_rank is not None:
+            lora_state_dict = self.lora_layers.state_dict()
+            torch.save(lora_state_dict, save_path.parent / "lora.ckpt")
+
+
+    def unnormalize_and_pil(self, tensor):
+        unnormed = tensor * torch.tensor(self.image_processor.image_std).view(3, 1, 1).to(
+            tensor.device) + torch.tensor(self.image_processor.image_mean).view(3, 1, 1).to(tensor.device)
+        return transforms.ToPILImage()(unnormed)
+
+    def save_images(self, image, conds, input_embeds, hidden_states, target_embeds, label='', save_path=''):
+        self.prior.eval()
+        input_images = []
+        captions = []
+        for i in range(len(conds)):
+            pil_image = self.unnormalize_and_pil(conds[i]).resize((self.cfg.img_size, self.cfg.img_size))
+            # Check if all zeros
+            if self.cfg.mode == 'clothes' and np.all(np.array(pil_image) > 250):
+                continue
+            input_images.append(pil_image)
+            captions.append("Condition")
+        if image is not None:
+            input_images.append(self.unnormalize_and_pil(image).resize((self.cfg.img_size, self.cfg.img_size)))
+            captions.append(f"Target {label}")
+
+        seeds = range(2)
+        output_images = []
+        embebds_to_vis = []
+        embeds_captions = []
+        if self.cfg.mode not in ['hybrid']:
+            embebds_to_vis += [target_embeds]
+            embeds_captions += ["Target Reconstruct" if image is not None else "Source Reconstruct"]
+        if self.cfg.mode not in ['instruct', 'colorization', 'clothes']:  # ,'histogram']:
+            for alpha in [0.5]:
+                embebds_to_vis.append(alpha * hidden_states[:, 0] + (1 - alpha) * hidden_states[:, 1])
+                embeds_captions.append(f"{alpha:.2f}A + {1 - alpha:.2f}B")
+        for embs in embebds_to_vis:
+            embs_unnormed = (embs * self.clip_std) + self.clip_mean
+            zero_embeds = self.prior_pipeline.get_zero_embed(embs.shape[0], device=embs.device)
+            direct_from_emb = self.decoder(image_embeds=embs_unnormed, negative_image_embeds=zero_embeds,
+                                           num_inference_steps=50, height=512,
+                                           width=512, guidance_scale=4).images
+            output_images = output_images + direct_from_emb
+        captions += embeds_captions
+
+        for seed in seeds:
+            negative_input_embeds = torch.zeros_like(input_embeds)
+            negative_hidden_states = torch.zeros_like(hidden_states)
+            for scale in [1, 4]:
+                img_emb = self.prior_pipeline(input_embeds=input_embeds, input_hidden_states=hidden_states,
+                                              negative_input_embeds=negative_input_embeds,
+                                              negative_input_hidden_states=negative_hidden_states,
+                                              num_inference_steps=25,
+                                              num_images_per_prompt=1,
+                                              guidance_scale=scale,
+                                              generator=torch.Generator(device="cuda").manual_seed(seed))
+
+                negative_emb = img_emb.negative_image_embeds
+                for seed_2 in range(1):
+                    images = self.decoder(image_embeds=img_emb.image_embeds, negative_image_embeds=negative_emb,
+                                          num_inference_steps=50, height=512,
+                                          width=512, guidance_scale=4,
+                                          generator=torch.Generator(device="cuda").manual_seed(seed_2)).images
+                    # torch.Generator(device="cuda").manual_seed(s) for s in range(4)
+                    output_images += images
+                    captions.append(f"prior_s {seed}, cfg {scale}, unet_s {seed_2}")
+
+        all_images = input_images + output_images
+        gen_images = vis_utils.create_table_plot(images=all_images, captions=captions)
+        gen_images.save(save_path)
+        self.prior.train()
+
+    def get_dataloaders(self) -> torch.utils.data.DataLoader:
+        sampler_train = None
+        if self.cfg.mode == 'texture':
+            dataset_path = self.cfg.dataset_path
+            if not isinstance(self.cfg.dataset_path, list):
+                dataset_path = [self.cfg.dataset_path]
+            datasets = []
+            for path in dataset_path:
+                datasets.append(TexturedDataset(
+                    dataset_dir=path,
+                    randoms_dir=self.cfg.randoms_dir,
+                    image_processor=self.image_processor,
+                    tokenizer=self.tokenizer
+                ))
+            dataset = torch.utils.data.ConcatDataset(datasets)
+            dataset_weights = []
+            for single_dataset in datasets:
+                dataset_weights.extend([len(dataset) / len(single_dataset)] * len(single_dataset))
+            sampler_train = torch.utils.data.WeightedRandomSampler(weights=dataset_weights,
+                                                                   num_samples=len(dataset_weights))
+
+            validation_dataset = TexturedDataset(
+                dataset_dir=self.cfg.val_dataset_path,
+                multi_dirs=False,
+                textures_dir=self.cfg.textures_dir,
+                image_processor=self.image_processor,
+                tokenizer=self.tokenizer
+            )
+        elif self.cfg.mode == 'instruct':
+            dataset = InstructDataset(
+                dataset_dir=self.cfg.dataset_path,
+                image_processor=self.image_processor,
+                tokenizer=self.tokenizer,
+            )
+            validation_dataset = InstructDataset(dataset_dir=self.cfg.val_dataset_path,
+                                                 image_processor=self.image_processor, tokenizer=self.tokenizer)
+        elif self.cfg.mode == 'clothes':
+            dataset = ClothesDataset(
+                image_processor=self.image_processor,
+                tokenizer=self.tokenizer,
+            )
+            validation_dataset = ClothesDataset(
+                image_processor=self.image_processor,
+                tokenizer=self.tokenizer,
+                val_mode=True
+            )
+        elif self.cfg.mode == 'union':
+            dataset_path = self.cfg.dataset_path
+            if not isinstance(self.cfg.dataset_path, list):
+                dataset_path = [self.cfg.dataset_path]
+            datasets = []
+            for path in dataset_path:
+                datasets.append(UnionDataset(dataset_dir=path,
+                                             random_dir=self.cfg.randoms_dir,
+                                             image_processor=self.image_processor,
+                                             tokenizer=self.tokenizer))
+            dataset = torch.utils.data.ConcatDataset(datasets)
+            dataset_weights = []
+            for single_dataset in datasets:
+                dataset_weights.extend([len(dataset) / len(single_dataset)] * len(single_dataset))
+            sampler_train = torch.utils.data.WeightedRandomSampler(weights=dataset_weights,
+                                                                   num_samples=len(dataset_weights))
+
+            validation_dataset = PairsDataset(dataset_dir=self.cfg.val_dataset_path,
+                                              image_processor=self.image_processor, tokenizer=self.tokenizer)
+        elif self.cfg.mode == 'scene':
+            dataset_path = self.cfg.dataset_path
+            if not isinstance(self.cfg.dataset_path, list):
+                dataset_path = [self.cfg.dataset_path]
+            datasets = []
+            for path in dataset_path:
+                datasets.append(SceneDataset(
+                    dataset_dir=path,
+                    background_dir=self.cfg.backgrounds_dir,
+                    random_dir=self.cfg.randoms_dir,
+                    image_processor=self.image_processor,
+                    tokenizer=self.tokenizer
+                ))
+            dataset = torch.utils.data.ConcatDataset(datasets)
+            dataset_weights = []
+            for single_dataset in datasets:
+                dataset_weights.extend([len(dataset) / len(single_dataset)] * len(single_dataset))
+            sampler_train = torch.utils.data.WeightedRandomSampler(weights=dataset_weights,
+                                                                   num_samples=len(dataset_weights))
+            validation_dataset = TexturedDataset(
+                dataset_dir=self.cfg.val_dataset_path,
+                multi_dirs=False,
+                textures_dir=self.cfg.textures_dir,
+                image_processor=self.image_processor,
+                tokenizer=self.tokenizer
+            )
+        else:
+            raise ValueError(f"Invalid mode: {self.cfg.mode}")
+        train_dataloader = torch.utils.data.DataLoader(
+            dataset, batch_size=self.cfg.train_batch_size, shuffle=sampler_train is None,
+            num_workers=self.cfg.num_workers,
+            sampler=sampler_train
+        )
+
+        validation_dataloader = torch.utils.data.DataLoader(
+            validation_dataset, batch_size=1, shuffle=True, num_workers=self.cfg.num_workers
+        )
+        return train_dataloader, validation_dataloader
+
+    def train(self):
+        pbar = tqdm(range(self.train_step, self.cfg.max_train_steps + 1))
+        while self.train_step < self.cfg.max_train_steps:
+            train_loss = 0.0
+            self.prior.train()
+
+            for sample_idx, batch in enumerate(self.train_dataloader):
+                # self.optimizer.zero_grad()
+                with self.accelerator.accumulate(self.prior):
+                    image, cond = batch
+
+                    image = image.to(self.weight_dtype).to(self.accelerator.device)
+                    if 'source_a' in cond:
+                        cond['source_a'] = cond['source_a'].to(self.weight_dtype).to(self.accelerator.device)
+                    if 'source_b' in cond:
+                        cond['source_b'] = cond['source_b'].to(self.weight_dtype).to(self.accelerator.device)
+                    if 'crops' in cond:
+                        for crop_ind in range(len(cond['crops'])):
+                            cond['crops'][crop_ind] = cond['crops'][crop_ind].to(self.weight_dtype).to(
+                                self.accelerator.device)
+                    for key in cond.keys():
+                        if isinstance(cond[key], torch.Tensor):
+                            cond[key] = cond[key].to(self.accelerator.device)
+
+                    with torch.no_grad():
+                        text_mask = cond['mask']
+                        text_encoder_output = self.text_encoder(cond['input_ids'])
+                        text_encoder_hidden_states = text_encoder_output.last_hidden_state
+                        # Take only the unmasked tokens
+                        text_encoder_concat = None
+
+                        if self.cfg.mode == 'instruct':
+                            text_target = self.text_encoder(cond['condition_ids']).text_embeds
+                            text_encoder_concat = text_encoder_hidden_states[:, :text_mask.sum().item()]
+
+                        image_embeds = self.image_encoder(image).image_embeds
+
+                        # Sample noise that we'll add to the image_embeds
+                        noise = torch.randn_like(image_embeds)
+                        bsz = image_embeds.shape[0]
+                        # Sample a random timestep for each image
+                        timesteps = torch.randint(0, self.noise_scheduler.config.num_train_timesteps, (bsz,),
+                                                  device=image_embeds.device)
+                        timesteps = timesteps.long()
+                        image_embeds = (image_embeds - self.clip_mean) / self.clip_std
+                        noisy_latents = self.noise_scheduler.add_noise(image_embeds, noise, timesteps)
+
+                        target = image_embeds
+
+                        input_image_embeds, input_hidden_state = self.preprocess_inputs(cond=cond,
+                                                                                        text_encoder_concat=text_encoder_concat)
+
+                # Create the input sequence for the prior
+                loss = 0
+                image_feat_seq = torch.zeros_like(text_encoder_hidden_states)
+                image_feat_seq[:, :input_hidden_state.shape[1]] = input_hidden_state
+                image_txt_mask = torch.zeros_like(text_mask)
+                image_txt_mask[:, :input_hidden_state.shape[1]] = 1
+
+                # Predict the noise residual and compute loss
+                model_pred = self.prior(
+                    noisy_latents,
+                    timestep=timesteps,
+                    proj_embedding=input_image_embeds,
+                    encoder_hidden_states=image_feat_seq,
+                    attention_mask=image_txt_mask,
+                ).predicted_image_embedding
+
+                # Calculate the loss
+                if self.cfg.mode == 'instruct':
+                    prior_loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+                    normed_pred = model_pred.float() * self.clip_std.detach() + self.clip_mean.detach()
+                    normed_pred = F.normalize(normed_pred, p=2, dim=-1)
+                    normed_text_target = F.normalize(text_target.float(), p=2, dim=-1)
+                    clip_loss = 1 - (normed_pred * normed_text_target).sum(dim=-1)
+                    prior_loss = prior_loss + self.cfg.clip_strength * clip_loss
+                else:
+                    prior_loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+
+                loss += prior_loss
+                # Gather the losses across all processes for logging (if we use distributed training).
+                avg_loss = self.accelerator.gather(loss.repeat(self.cfg.train_batch_size)).mean()
+                train_loss += avg_loss.item() / self.cfg.gradient_accumulation_steps
+
+                # Backprop
+                self.accelerator.backward(loss)
+                if self.accelerator.sync_gradients:
+                    self.accelerator.clip_grad_norm_(self.prior.parameters(), self.cfg.max_grad_norm)
+                self.optimizer.step()
+                self.optimizer.zero_grad()
+
+                # Checks if the accelerator has performed an optimization step behind the scenes
+                if self.accelerator.sync_gradients:
+                    if self.cfg.use_ema:
+                        self.ema_prior.step(self.prior.parameters())
+                    pbar.update(1)
+                    self.train_step += 1
+                    train_loss = 0.0
+
+                    if self.train_step % self.cfg.checkpointing_steps == 1:
+                        if self.accelerator.is_main_process:
+                            save_path = self.cfg.output_dir / f"learned_prior.pth"
+                            self.save_model(save_path)
+                            logger.info(f"Saved state to {save_path}")
+                    pbar.set_postfix(**{"loss": loss.cpu().detach().item()})
+
+                    if self.cfg.log_image_frequency > 0 and (
+                            self.train_step % self.cfg.log_image_frequency == 1) or self.train_step < 10:
+                        image_save_path = self.cfg.output_dir / 'images' / f"{self.train_step}_step_images.jpg"
+                        image_save_path.parent.mkdir(exist_ok=True, parents=True)
+                        # Apply the full diffusion process
+                        conds_list = []
+                        if 'source_a' in cond:
+                            conds_list.append(cond['source_a'][0])
+                        if 'source_b' in cond:
+                            conds_list.append(cond['source_b'][0])
+                        if 'crops' in cond:
+                            for crop_ind in range(len(cond['crops'])):
+                                conds_list.append(cond['crops'][crop_ind][0])
+
+                        label = ''
+                        if 'text_condition' in cond:
+                            label += cond['text_condition'][0]
+                        self.save_images(image=image[0], conds=conds_list,
+                                         label=label,
+                                         input_embeds=input_image_embeds[0].unsqueeze(0),
+                                         hidden_states=input_hidden_state[0].unsqueeze(0),
+                                         target_embeds=image_embeds[0].unsqueeze(0),
+                                         save_path=image_save_path)
+
+                    if self.cfg.log_validation > 0 and (self.train_step % self.cfg.log_validation == 10):
+                        # Run validation
+                        self.log_validation()
+
+                    if self.train_step >= self.cfg.max_train_steps:
+                        break
+
+            self.train_dataloader, self.validation_dataloader = self.get_dataloaders()
+        pbar.close()
+
+    def preprocess_inputs(self, cond, text_encoder_concat=None):
+
+        # Specific pre-processing for each mode
+        if self.cfg.mode == 'instruct':
+            input_image_embeds, input_hidden_state = pops_utils.preprocess(cond['source_a'], None,
+                                                                    self.image_encoder,
+                                                                    self.clip_mean.detach(),
+                                                                    self.clip_std.detach(),
+                                                                    concat_hidden_states=text_encoder_concat)
+            # For CFG
+            if random.random() < 0.2:
+                input_hidden_state = torch.zeros_like(input_hidden_state)
+                for b_ind in range(cond['source_a'].shape[0]):
+                    cond['source_a'][b_ind] = cond['source_a'][b_ind] * 0
+        elif self.cfg.mode == 'clothes':
+            input_image_embeds, input_hidden_state = pops_utils.preprocess(None, None,
+                                                                    self.image_encoder,
+                                                                    self.clip_mean.detach(),
+                                                                    self.clip_std.detach(),
+                                                                    image_list=cond['crops'])
+            if random.random() < 0.2:
+                input_hidden_state = torch.zeros_like(input_hidden_state)
+                # For vis purposes zero all images
+                cond['crops'] = [torch.zeros_like(cond['crops'][0])]
+        else:
+            if self.cfg.mode == 'texture':
+                should_drop_cond = []
+                for b_ind in range(cond['source_a'].shape[0]):
+                    # For CFG and random data
+                    if 'is_random' in cond and cond['is_random'][b_ind]:
+                        should_drop_a = random.random() < 0.5
+                        should_drop_b = True
+                    else:
+                        # Randomly dropping one of the conditions to make the model use both?
+                        should_drop_a = random.random() < 0.25
+                        should_drop_b = random.random() < 0.25
+
+                    if should_drop_a:
+                        cond['source_a'][b_ind] = cond['source_a'][b_ind] * 0
+                    if should_drop_b:
+                        cond['source_b'][b_ind] = cond['source_b'][b_ind] * 0
+                    should_drop_cond.append((should_drop_a, should_drop_b))
+            elif self.cfg.mode == 'union':
+                should_drop_cond = []
+                for b_ind in range(cond['source_a'].shape[0]):
+                    should_drop_a, should_drop_b = False, False
+                    if ('is_random' in cond and cond['is_random'][b_ind]) or random.random() < 0.2:
+                        should_drop_a = True
+                        should_drop_b = True
+
+                    if should_drop_a:
+                        cond['source_a'][b_ind] = cond['source_a'][b_ind] * 0
+                    if should_drop_b:
+                        cond['source_b'][b_ind] = cond['source_b'][b_ind] * 0
+                    should_drop_cond.append((should_drop_a, should_drop_b))
+            elif self.cfg.mode == 'scene':
+                should_drop_cond = []
+                for b_ind in range(cond['source_a'].shape[0]):
+                    should_drop_a, should_drop_b = False, False
+                    if ('is_random' in cond and cond['is_random'][b_ind]) or (random.random() < 0.1):
+                        should_drop_a = True
+                        should_drop_b = True
+
+                    if should_drop_a:
+                        cond['source_a'][b_ind] = cond['source_a'][b_ind] * 0
+                    if should_drop_b:
+                        cond['source_b'][b_ind] = cond['source_b'][b_ind] * 0
+                    should_drop_cond.append((should_drop_a, should_drop_b))
+            else:
+                raise ValueError(f"Invalid mode: {self.cfg.mode}")
+            input_image_embeds, input_hidden_state = pops_utils.preprocess(cond['source_a'], cond['source_b'],
+                                                                    self.image_encoder,
+                                                                    self.clip_mean.detach(),
+                                                                    self.clip_std.detach(),
+                                                                    should_drop_cond=should_drop_cond)
+
+        return input_image_embeds, input_hidden_state
+
+    def log_validation(self):
+        for sample_idx, batch in tqdm(enumerate(self.validation_dataloader)):
+            image, cond = batch
+            image = image.to(self.weight_dtype).to(self.accelerator.device)
+            if 'source_a' in cond:
+                cond['source_a'] = cond['source_a'].to(self.weight_dtype).to(self.accelerator.device)
+            if 'source_b' in cond:
+                cond['source_b'] = cond['source_b'].to(self.weight_dtype).to(self.accelerator.device)
+            if 'crops' in cond:
+                for crop_ind in range(len(cond['crops'])):
+                    cond['crops'][crop_ind] = cond['crops'][crop_ind].to(self.weight_dtype).to(
+                        self.accelerator.device)
+            for key in cond.keys():
+                if isinstance(cond[key], torch.Tensor):
+                    cond[key] = cond[key].to(self.accelerator.device)
+
+            text_mask = cond['mask']
+
+            with torch.no_grad():
+                target_embeds = self.image_encoder(image).image_embeds
+                target_embeds = (target_embeds - self.clip_mean) / self.clip_std
+
+                text_encoder_output = self.text_encoder(cond['input_ids'])
+                text_encoder_hidden_states = text_encoder_output.last_hidden_state
+                # Take only the unmasked tokens
+                if self.cfg.mode == 'instruct':
+                    text_encoder_concat = text_encoder_hidden_states[:, :text_mask.sum().item()]
+                    input_image_embeds, input_hidden_state = pops_utils.preprocess(cond['source_a'],
+                                                                            None,
+                                                                            self.image_encoder,
+                                                                            self.clip_mean.detach(),
+                                                                            self.clip_std.detach(),
+                                                                            concat_hidden_states=text_encoder_concat)
+                elif self.cfg.mode == 'clothes':
+                    input_image_embeds, input_hidden_state = pops_utils.preprocess(None, None,
+                                                                            self.image_encoder,
+                                                                            self.clip_mean.detach(),
+                                                                            self.clip_std.detach(),
+                                                                            image_list=cond['crops'])
+
+                else:
+                    input_image_embeds, input_hidden_state = pops_utils.preprocess(cond['source_a'],
+                                                                            cond['source_b'],
+                                                                            self.image_encoder,
+                                                                            self.clip_mean.detach(),
+                                                                            self.clip_std.detach())
+
+            image_save_path = self.cfg.output_dir / 'val_images' / f"{self.train_step}_step_{sample_idx}_images.jpg"
+            image_save_path.parent.mkdir(exist_ok=True, parents=True)
+
+            save_target_image = image[0] if self.cfg.textures_dir is None else None
+            conds_list = []
+            if 'source_a' in cond:
+                conds_list.append(cond['source_a'][0])
+            if 'source_b' in cond:
+                conds_list.append(cond['source_b'][0])
+            if 'crops' in cond:
+                for crop_ind in range(len(cond['crops'])):
+                    conds_list.append(cond['crops'][crop_ind][0])
+            # Apply the full diffusion process
+            self.save_images(image=save_target_image, conds=conds_list,
+                             label=cond['text_condition'][0] if 'text_condition' in cond else '',
+                             input_embeds=input_image_embeds[0].unsqueeze(0),
+                             hidden_states=input_hidden_state[0].unsqueeze(0),
+                             target_embeds=target_embeds[0].unsqueeze(0),
+                             save_path=image_save_path)
+
+            if sample_idx == self.cfg.n_val_images:
+                break
diff --git a/scripts/dataset.py b/scripts/dataset.py
new file mode 100644
index 0000000..46410d1
--- /dev/null
+++ b/scripts/dataset.py
@@ -0,0 +1,696 @@
+import json
+import random
+import traceback
+from pathlib import Path
+from typing import Optional
+
+import cv2
+import numpy as np
+from PIL import Image
+from torch.utils.data import Dataset
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+
+from data_generation import words_bank
+
+try:
+    from torchvision.transforms import InterpolationMode
+
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+
+def _transform(n_px):
+    return Compose(
+        [
+            Resize(n_px, interpolation=BICUBIC),
+            CenterCrop(n_px),
+            _convert_image_to_rgb,
+            ToTensor(),
+            Normalize(
+                (0.48145466, 0.4578275, 0.40821073),
+                (0.26862954, 0.26130258, 0.27577711),
+            ),
+        ]
+    )
+
+
+class TexturedDataset(Dataset):
+    def __init__(
+            self,
+            dataset_dir: Path,
+            multi_dirs: bool = True,
+            textures_dir: Optional[Path] = None,
+            randoms_dir: Optional[Path] = None,
+            clip_image_size: int = 224,
+            image_processor=None,
+            tokenizer=None
+    ):
+        """ Texturing dataset
+        Args:
+            dataset_dir (Path): Path to the directory containing images.
+            multi_dirs (bool): Whether to expect multiple directories in dataset_dir.
+            textures_dir (Optional[Path]): Path to the directory containing textures, used for validation
+            randoms_dir (Optional[Path]): Path to the directory containing random images, to further diversify the dataset.
+            clip_image_size (int): Size of the image to be fed into the model.
+            image_processor: ImageProcessor: Image processor to be used.
+            tokenizer: Tokenizer: Tokenizer to be used.
+        """
+        if multi_dirs:
+            self.paths = [list(d.glob('*.jpg')) for d in dataset_dir.iterdir() if d.is_dir()]
+        else:
+            # Read each directory in dataset_dir
+            self.paths = [list(dataset_dir.glob('*.jpg'))]
+        self.flat_inds = []
+        for i, path_list in enumerate(self.paths):
+            self.flat_inds.extend([i] * len(path_list))
+
+        print(f'For dataset {dataset_dir}')
+        print(f'Loaded {len(self.flat_inds)} images.')
+        if textures_dir:
+            self.textures = list(textures_dir.glob('*.jpg'))
+            print(f'Loaded {len(self.textures)} textures.')
+        else:
+            self.textures = None
+
+        if randoms_dir:
+            self.randoms = list(randoms_dir.glob('*.jpg'))
+            print(f'Loaded {len(self.randoms)} randoms.')
+        else:
+            self.randoms = None
+        self.tokenizer = tokenizer
+        self.transform1 = _transform(clip_image_size)
+
+        self.clip_image_size = clip_image_size
+
+        self.image_processor = image_processor
+
+
+    def __len__(self):
+        return len(self.flat_inds)
+
+    def crop_patch(self, image, box, mask):
+        """ Logic for cropping an exemplar patch from an image. """
+
+        x_min = box[0]
+        x_max = box[2]
+        y_min = box[1]
+        y_max = box[3]
+        max_rect_size = min(x_max - x_min, y_max - y_min) // 2
+        min_rect_size = min(96, max_rect_size - 1)
+        rect_size = random.randint(min_rect_size, max_rect_size)
+
+        base_crop = image.crop(box)
+        # Erode mask so we don't choose center from boundary
+        mask = np.array(mask)
+        mask[mask < 250] = 0
+        mask[mask >= 250] = 255
+
+        if np.count_nonzero(mask) == 0:
+            x_center = (x_min + x_max) // 2
+            y_center = (y_min + y_max) // 2
+            out_crop = image.crop((x_center - rect_size // 2, y_center - rect_size // 2, x_center + rect_size // 2,
+                                   y_center + rect_size // 2))
+        else:
+            eroded_mask = cv2.erode(mask, np.ones((min_rect_size + 1, min_rect_size + 1), np.uint8), iterations=1,
+                                    borderType=cv2.BORDER_CONSTANT, borderValue=0)
+            foreground_y, foreground_x = np.where(eroded_mask == 255)
+            foreground_pixels = list(zip(foreground_y, foreground_x))
+            if len(foreground_pixels) == 0:
+                last_non_empty_mask = mask.copy()
+                curr_mask = mask.copy()
+                kernel = np.ones((11, 11), np.uint8)
+                while np.count_nonzero(curr_mask) > 0:
+                    last_non_empty_mask = curr_mask.copy()
+                    curr_mask = cv2.erode(curr_mask, kernel, iterations=1, borderType=cv2.BORDER_CONSTANT,
+                                          borderValue=0)
+                foreground_y, foreground_x = np.where(last_non_empty_mask == 255)
+                foreground_pixels = list(zip(foreground_y, foreground_x))
+
+            y_center, x_center = random.choice(foreground_pixels)
+            # Check if rect_size is all in mask, if not shorten
+            curr_size = rect_size
+            # Bound to size of image based on y_center and x_center
+            curr_size = min(curr_size, min(x_center, y_center, image.width - x_center, image.height - y_center) * 2)
+
+            while curr_size > min_rect_size - 1:
+                if np.all(mask[y_center - curr_size // 2:y_center + curr_size // 2,
+                          x_center - curr_size // 2:x_center + curr_size // 2] == 255):
+                    break
+                curr_size -= 1
+
+            out_crop = base_crop.crop((x_center - curr_size // 2, y_center - curr_size // 2, x_center + curr_size // 2,
+                                       y_center + curr_size // 2))
+
+        return out_crop
+
+    def __getitem__(self, i: int):
+        is_random = False
+        try:
+            # If a random directory is given load a random image with some probability
+            if self.randoms is not None and random.random() < 0.2:
+                random_path = random.choice(self.randoms)
+                random_image = Image.open(random_path).convert("RGB")
+                source_image = random_image
+
+                target_image = random_image
+                # Calculate random crop size of scale 0.1 to 0.5
+                rand_scale = random.uniform(0.2, 0.4)
+                crop_size = min(target_image.width, target_image.height) * rand_scale
+                # Randomly choose crop location around image center, shift by 0.1width max
+                crop_center_x = random.randint(int(target_image.width * 0.45), int(target_image.width * 0.55))
+                crop_center_y = random.randint(int(target_image.height * 0.45), int(target_image.height * 0.55))
+                target_patch = target_image.crop((crop_center_x - crop_size // 2, crop_center_y - crop_size // 2,
+                                                  crop_center_x + crop_size // 2, crop_center_y + crop_size // 2))
+                is_random = True
+            else:
+                relevant_paths = self.paths[self.flat_inds[i]]
+
+                target_path = random.choice(relevant_paths)
+                if random.random() < 0.2:
+                    source_path = target_path
+                else:
+                    source_path = random.choice(relevant_paths)
+
+                target_image = Image.open(target_path).convert("RGB")
+                source_image = Image.open(source_path).convert("RGB")
+
+                # For validation where textures are taken from a different directory
+                if self.textures is not None:
+                    texture_path = random.choice(self.textures)
+                    target_patch = Image.open(texture_path).convert("RGB")
+                    target_image = source_image
+                else:
+                    box_path = target_path.parent / 'box.json'
+                    if box_path.exists():
+                        with open(box_path, 'r') as f:
+                            box = json.load(f)
+                            box = (box['x1'], box['y1'], box['x2'], box['y2'])
+                    else:
+                        box = (0, 0, target_image.width, target_image.height)
+                    mask_path = target_path.parent / 'mask.png'
+                    if mask_path.exists():
+                        mask = Image.open(mask_path).convert("L")
+                    else:
+                        mask = Image.new('L', (target_image.width, target_image.height), 0)
+                    target_patch = self.crop_patch(target_image, box, mask)
+
+            target_image = self.image_processor(target_image)['pixel_values'][0]
+            source_image = self.image_processor(source_image)['pixel_values'][0]
+            target_patch = self.image_processor(target_patch)['pixel_values'][0]
+        except Exception as e:
+            print(f'Error processing image: {e}')
+            traceback.print_exc()
+            empty_image = Image.new('RGB', (self.clip_image_size, self.clip_image_size), (255, 255, 255))
+            target_image = self.image_processor(empty_image)['pixel_values'][0]
+            source_image = self.image_processor(empty_image)['pixel_values'][0]
+            target_patch = self.image_processor(empty_image)['pixel_values'][0]
+
+        out_dict = {}
+
+        text = ""
+
+        text_inputs = self.tokenizer(
+            text,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        out_dict["input_ids"] = text_inputs.input_ids[0]
+        out_dict["mask"] = text_inputs.attention_mask.bool()[0]
+        out_dict["text"] = text
+
+        out_dict["source_a"] = source_image
+        out_dict["source_b"] = target_patch
+        out_dict["is_random"] = is_random
+
+        return target_image, out_dict
+
+
+class InstructDataset(Dataset):
+    def __init__(
+            self,
+            dataset_dir: Path,
+            clip_image_size: int = 224,
+            image_processor=None,
+            tokenizer=None,
+
+    ):
+        self.paths = list(dataset_dir.glob('*.jpg'))
+        self.tokenizer = tokenizer
+        self.transform1 = _transform(clip_image_size)
+
+        self.clip_image_size = clip_image_size
+
+        self.image_processor = image_processor
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, i: int):
+        path = self.paths[i]
+        prompt = ' '.join(path.stem.split('_')[:-1])
+        image = Image.open(path).convert("RGB")
+        clip_image = self.image_processor(image)['pixel_values'][0]
+
+        out_dict = {}
+
+        if random.random() < 0.5:
+            adjective = random.choice(words_bank.adjectives)
+            text = adjective
+            text_condition = f'A {adjective} {prompt}'
+        else:
+            art_type = random.choice(words_bank.art_types)
+            text = art_type
+            text_condition = f'A {art_type} of a {prompt}'
+
+        text_inputs = self.tokenizer(
+            text,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        out_dict["input_ids"] = text_inputs.input_ids[0]
+
+        out_dict["mask"] = text_inputs.attention_mask.bool()[0]
+        out_dict["text"] = text
+
+        text_condition_inputs = self.tokenizer(
+            text_condition,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        out_dict["condition_ids"] = text_condition_inputs.input_ids[0]
+
+        out_dict["text_condition"] = text_condition
+
+        out_dict["source_a"] = clip_image
+
+        return clip_image, out_dict
+
+
+class PairsDataset(Dataset):
+    """ Used for validation, simply creates pairs of images from the same directory."""
+    def __init__(
+            self,
+            dataset_dir: Path,
+            clip_image_size: int = 224,
+            image_processor=None,
+            tokenizer=None,
+
+    ):
+        self.paths = list(dataset_dir.glob('*.jpg'))
+        self.tokenizer = tokenizer
+        self.transform1 = _transform(clip_image_size)
+
+        self.clip_image_size = clip_image_size
+
+        self.image_processor = image_processor
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, i: int):
+        path = self.paths[i]
+        path_b = random.choice(self.paths)
+        image = Image.open(path).convert("RGB")
+        image_b = Image.open(path_b).convert("RGB")
+        clip_image = self.image_processor(image)['pixel_values'][0]
+
+        clip_image_a = self.image_processor(image)['pixel_values'][0]
+        clip_image_b = self.image_processor(image_b)['pixel_values'][0]
+
+        out_dict = {}
+
+        text = ''
+
+        text_inputs = self.tokenizer(
+            text,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        out_dict["input_ids"] = text_inputs.input_ids[0]
+
+        out_dict["mask"] = text_inputs.attention_mask.bool()[0]
+        out_dict["text"] = text
+
+        out_dict["source_a"] = clip_image_a
+        out_dict["source_b"] = clip_image_b
+
+        return clip_image, out_dict
+
+
+class SceneDataset(Dataset):
+    def __init__(
+            self,
+            dataset_dir: Path,
+            background_dir: Path,
+            random_dir: Path,
+            clip_image_size: int = 224,
+            image_processor=None,
+            tokenizer=None,
+
+    ):
+        self.paths = [p for p in list(dataset_dir.glob('*.jpg')) if not p.name.endswith('_inpainted.jpg')]
+        self.backgrounds = list(background_dir.glob('*.jpg'))
+        self.randoms = list(random_dir.glob('*.jpg')) if random_dir is not None else None
+        self.tokenizer = tokenizer
+        self.transform1 = _transform(clip_image_size)
+
+        self.clip_image_size = clip_image_size
+
+        self.image_processor = image_processor
+
+    def __len__(self):
+        return len(self.paths)
+
+    def crop_patch(self, image, mask):
+
+        # if mask is empty return zeros
+        if np.count_nonzero(mask) == 0:
+            return Image.new('RGB', (self.clip_image_size, self.clip_image_size), (255, 255, 255))
+        last_non_empty_mask = mask.copy()
+        curr_mask = mask.copy()
+        kernel = np.ones((11, 11), np.uint8)
+        rect_size = 11
+        while np.count_nonzero(curr_mask) > 1000:
+            rect_size += 11
+            last_non_empty_mask = curr_mask.copy()
+            curr_mask = cv2.erode(curr_mask, kernel, iterations=1, borderType=cv2.BORDER_CONSTANT,
+                                  borderValue=0)
+        foreground_y, foreground_x = np.where(last_non_empty_mask == 1)
+        foreground_pixels = list(zip(foreground_y, foreground_x))
+        y_center, x_center = random.choice(foreground_pixels)
+        # Check if rect_size is all in mask, if not shorten
+        curr_size = rect_size
+        # Bound to size of image based on y_center and x_center
+        curr_size = min(curr_size, min(x_center, y_center, image.width - x_center, image.height - y_center) * 2)
+        out_crop = image.crop((x_center - curr_size // 2, y_center - curr_size // 2, x_center + curr_size // 2,
+                               y_center + curr_size // 2))
+
+        return out_crop
+
+    def __getitem__(self, i: int):
+        is_random = False
+        try:
+            if self.randoms is not None and random.random() < 0.1:
+                path = random.choice(self.randoms)
+                image = Image.open(path).convert("RGB")
+                object_with_background = Image.open(path).convert("RGBA")
+                background_patch = Image.open(path).convert("RGBA")
+                is_random = True
+            else:
+
+                path = self.paths[i]
+                image = Image.open(path).convert("RGB")
+                if path.with_suffix('.npy').exists():
+                    mask = np.load(path.with_suffix('.npy'))  # [:,:,3]
+                else:
+                    mask = np.load(path.with_name(path.stem + '_inpainted.npy'))
+                if random.random() < 0.5:
+                    # White background
+                    background = Image.new("RGB", image.size, (255, 255, 255))
+                else:
+                    background_path = random.choice(self.backgrounds)
+                    background = Image.open(background_path).convert("RGB")
+
+
+                background.paste(image, (0, 0), Image.fromarray(mask))
+                object_with_background = background
+
+                # Generate a patch from the original background
+                # Check if has _inpainted file
+                inpainted_path = path.with_name(path.stem + '_inpainted.jpg')
+                if inpainted_path.exists():
+                    background_patch = Image.open(inpainted_path).convert("RGB")
+                else:
+                    mask[mask < 10] = 0
+                    mask[mask >= 10] = 1
+                    mask = 1 - mask
+                    background_patch = self.crop_patch(image, mask)
+                    # Check size of patch
+                    if background_patch.size[0] < 50 or background_patch.size[1] < 50:
+                        background_patch = Image.new('RGB', (self.clip_image_size, self.clip_image_size),
+                                                     (255, 255, 255))
+        except Exception as e:
+            print(f'Error processing image: {e}')
+            traceback.print_exc()
+            empty_image = Image.new('RGB', (self.clip_image_size, self.clip_image_size), (255, 255, 255))
+            image = empty_image
+            object_with_background = empty_image
+            background_patch = empty_image
+
+        clip_image = self.image_processor(image)['pixel_values'][0]
+        clip_image_a = self.image_processor(object_with_background.convert('RGB'))['pixel_values'][0]
+        try:
+            clip_image_b = self.image_processor(background_patch.convert('RGB'))['pixel_values'][0]
+        except Exception as e:
+            print(f'Error processing image: {e}')
+            traceback.print_exc()
+            clip_image_b = \
+                self.image_processor(Image.new('RGB', (self.clip_image_size, self.clip_image_size), (255, 255, 255)))[
+                    'pixel_values'][0]
+
+        out_dict = {}
+
+        text = ''
+
+        text_inputs = self.tokenizer(
+            text,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        out_dict["input_ids"] = text_inputs.input_ids[0]
+
+        out_dict["mask"] = text_inputs.attention_mask.bool()[0]
+        out_dict["text"] = text
+
+        out_dict["source_a"] = clip_image_a
+        out_dict["source_b"] = clip_image_b
+        out_dict["is_random"] = is_random
+
+        return clip_image, out_dict
+
+
+class UnionDataset(Dataset):
+    def __init__(
+            self,
+            dataset_dir: Path,
+            random_dir: Path,
+            clip_image_size: int = 224,
+            image_processor=None,
+            tokenizer=None,
+
+    ):
+        self.paths = list(dataset_dir.glob('*.jpg'))
+        self.randoms = list(random_dir.glob('*.jpg')) if random_dir is not None else None
+        self.tokenizer = tokenizer
+        self.transform1 = _transform(clip_image_size)
+
+        self.clip_image_size = clip_image_size
+
+        self.image_processor = image_processor
+
+    def __len__(self):
+        return len(self.paths)
+
+    def paste_on_background(self, image, background, min_scale=0.4, max_scale=0.8):
+        # Calculate aspect ratio and determine resizing based on the smaller dimension of the background
+        aspect_ratio = image.width / image.height
+        scale = random.uniform(min_scale, max_scale)
+        new_width = int(min(background.width, background.height * aspect_ratio) * scale)
+        new_height = int(new_width / aspect_ratio)
+
+        # Resize image and calculate position
+        image = image.resize((new_width, new_height), resample=Image.LANCZOS)
+        pos_x = random.randint(0, background.width - new_width)
+        pos_y = random.randint(0, background.height - new_height)
+
+        # Paste the image using its alpha channel as mask if present
+        background.paste(image, (pos_x, pos_y), image if 'A' in image.mode else None)
+        return background
+
+    def __getitem__(self, i: int):
+        is_random = False
+        try:
+            if self.randoms is not None and random.random() < 0.1:
+                path = random.choice(self.randoms)
+                image = Image.open(path).convert("RGB")
+                object_a = Image.open(path).convert("RGBA")
+                object_b = Image.open(path).convert("RGBA")
+                is_random = True
+            else:
+                path = self.paths[i]
+                image = Image.open(path).convert("RGB")
+                object_a = Image.open(str(path).replace('.jpg', f'_OBJ_0.png')).convert("RGBA")
+                object_b = Image.open(str(path).replace('.jpg', f'_OBJ_1.png')).convert("RGBA")
+
+                # Use white background
+                background_a = Image.new("RGB", image.size, (255, 255, 255))
+                background_b = Image.new("RGB", image.size, (255, 255, 255))
+
+                # Convert to RGB with white background
+                self.paste_on_background(object_a, background_a, min_scale=0.8, max_scale=0.95)
+                self.paste_on_background(object_b, background_b, min_scale=0.8, max_scale=0.95)
+                object_a, object_b = background_a, background_b
+
+                # Random switch between a,b
+                if random.random() < 0.5:
+                    object_a, object_b = object_b, object_a
+        except Exception as e:
+            print(f'Error processing image: {e}')
+            traceback.print_exc()
+            empty_image = Image.new('RGB', (self.clip_image_size, self.clip_image_size), (255, 255, 255))
+            image = empty_image
+            object_a = empty_image
+            object_b = empty_image
+
+        clip_image = self.image_processor(image)['pixel_values'][0]
+        clip_image_a = self.image_processor(object_a.convert('RGB'))['pixel_values'][0]
+        clip_image_b = self.image_processor(object_b.convert('RGB'))['pixel_values'][0]
+
+        # clip_image = self.image_processor(image)['pixel_values'][0]
+
+        out_dict = {}
+
+        text = ''
+
+        text_inputs = self.tokenizer(
+            text,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        out_dict["input_ids"] = text_inputs.input_ids[0]
+
+        out_dict["mask"] = text_inputs.attention_mask.bool()[0]
+        out_dict["text"] = text
+
+        out_dict["source_a"] = clip_image_a
+        out_dict["source_b"] = clip_image_b
+        out_dict["is_random"] = is_random
+
+        return clip_image, out_dict
+
+
+class ClothesDataset(Dataset):
+    def __init__(
+            self,
+            clip_image_size: int = 224,
+            image_processor=None,
+            val_mode=False,
+            tokenizer=None,
+
+    ):
+        from datasets import load_dataset
+        self.dataset = load_dataset("mattmdjaga/human_parsing_dataset")['train']
+        self.tokenizer = tokenizer
+        self.transform1 = _transform(clip_image_size)
+
+        self.relevant_inds = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 16, 17]
+
+        self.clip_image_size = clip_image_size
+
+        self.image_processor = image_processor
+        self.val_mode = val_mode
+        self.val_size = 100
+
+    def __len__(self):
+        if self.val_mode:
+            return self.val_size
+        else:
+            return len(self.dataset) - self.val_size
+
+    def crop_by_mask(self, image, mask):
+        image = np.array(image)
+        # Apply mask on image
+        image = image * mask[:, :, None] + (1 - mask[:, :, None]) * 255
+        image = image.astype(np.uint8)
+        rows = np.any(mask, axis=1)
+        cols = np.any(mask, axis=0)
+        ymin, ymax = np.where(rows)[0][[0, -1]]
+        xmin, xmax = np.where(cols)[0][[0, -1]]
+
+        # Crop the image
+        cropped_image = Image.fromarray(image[ymin:ymax + 1, xmin:xmax + 1])
+        cropped_mask = mask[ymin:ymax + 1, xmin:xmax + 1]
+        return cropped_image, cropped_mask
+
+    def paste_on_background(self, image, background, min_scale=0.4, max_scale=0.8):
+        # Calculate aspect ratio and determine resizing based on the smaller dimension of the background
+        aspect_ratio = image.width / image.height
+        scale = random.uniform(min_scale, max_scale)
+        new_width = int(min(background.width, background.height * aspect_ratio) * scale)
+        new_height = int(new_width / aspect_ratio)
+
+        # Resize image and calculate position
+        image = image.resize((new_width, new_height), resample=Image.LANCZOS)
+        pos_x = random.randint(0, background.width - new_width)
+        pos_y = random.randint(0, background.height - new_height)
+
+        # Paste the image using its alpha channel as mask if present
+        background.paste(image, (pos_x, pos_y), image if 'A' in image.mode else None)
+        return background
+
+    def __getitem__(self, i: int):
+        if self.val_mode:
+            i += len(self.dataset) - self.val_size
+        image = self.dataset[i]['image']
+        mask = self.dataset[i]['mask']
+
+        # Pad image to rect with white background
+
+        crops = []
+        for ind in self.relevant_inds:
+            curr_mask = np.array(mask) == ind
+            object_with_background = Image.new("RGB", (self.clip_image_size, self.clip_image_size), (255, 255, 255))
+            if np.any(curr_mask):
+                cropped_object, cropped_mask = self.crop_by_mask(image, curr_mask)
+                # Make sure crop is above certain size
+                if min(cropped_object.size) < 60:  # or max(cropped_object.size) < 80:
+                    continue
+
+                if np.count_nonzero(cropped_mask) < 1800:
+                    continue
+                self.paste_on_background(cropped_object, object_with_background, min_scale=0.8, max_scale=0.95)
+
+            crops.append(self.image_processor(object_with_background)['pixel_values'][0])
+        large_axis = max(image.size)
+        rect_image = Image.new('RGB', (large_axis, large_axis), (255, 255, 255))
+        # Paste in the middle
+        rect_image.paste(image, (large_axis // 2 - self.dataset[i]['image'].size[0] // 2,
+                                 large_axis // 2 - self.dataset[i]['image'].size[1] // 2))
+
+        target_clip_image = self.image_processor(rect_image)['pixel_values'][0]
+
+        out_dict = {}
+        text = ''
+
+        text_inputs = self.tokenizer(
+            text,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        out_dict["input_ids"] = text_inputs.input_ids[0]
+
+        out_dict["mask"] = text_inputs.attention_mask.bool()[0]
+        out_dict["text"] = text
+
+        out_dict["crops"] = crops
+
+        return target_clip_image, out_dict
diff --git a/scripts/train_config.py b/scripts/train_config.py
new file mode 100644
index 0000000..ee37858
--- /dev/null
+++ b/scripts/train_config.py
@@ -0,0 +1,89 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import List, Optional, Union
+
+
+@dataclass
+class TrainConfig:
+    # Dataset path
+    dataset_path: Union[Path, List[Path]] = Path('datasets/generated/generated_things')
+    # Validation dataset path
+    val_dataset_path: Path = Path('datasets/generated/generated_things_val')
+    # Path to pretrained model WITHOUT 2_1 folder
+    cache_root: Path = Path('/tmp/kandinsky2')
+    # The output directory where the model predictions and checkpoints will be written.
+    output_dir: Path = Path('results/my_pops_model')
+    # GPU device
+    device: str = 'cuda:0'
+    # The resolution for input images, all the images will be resized to this size
+    img_size: int = 512
+    # Batch size (per device) for the training dataloader
+    train_batch_size: int = 1
+    # Initial learning rate (after the potential warmup period) to use
+    lr: float = 1e-5
+    # Dataloader num workers.
+    num_workers: int = 8
+    # The beta1 parameter for the Adam optimizer.
+    adam_beta1: float = 0.9
+    # The beta2 parameter for the Adam optimizer
+    adam_beta2: float = 0.999
+    # Weight decay to use
+    adam_weight_decay: float = 0.0  # 1e-2
+    # Epsilon value for the Adam optimizer
+    adam_epsilon: float = 1e-08
+    # How often save images. Values less zero - disable saving
+    log_image_frequency: int = 500
+    # How often to run validation
+    log_validation: int = 5000
+    # The number of images to save during each validation
+    n_val_images: int = 10
+    # A seed for reproducible training
+    seed: Optional[int] = None
+    # The number of accumulation steps to use
+    gradient_accumulation_steps: int = 1
+    # Whether to use mixed precision training
+    mixed_precision: Optional[str] = 'fp16'
+    # Log to wandb
+    report_to: str = 'wandb'
+    # Path to pretrained prior model or model identifier from huggingface.co/models.
+    pretrained_prior_path: str = 'kandinsky-community/kandinsky-2-2-prior'
+    # Path to pretrained image encoder.
+    pretrained_image_encoder: str = 'kandinsky-community/kandinsky-2-2-prior'
+    # Path to scheduler.
+    scheduler_path: str = 'kandinsky-community/kandinsky-2-2-prior'
+    # Path to image_processor.
+    image_processor_path: str = 'kandinsky-community/kandinsky-2-2-prior'
+    # Path to text_encoder.
+    text_encoder_path: str = 'kandinsky-community/kandinsky-2-2-prior'
+    # Path to tokenizer.
+    tokenizer_path: str = 'kandinsky-community/kandinsky-2-2-prior'
+    use_ema: bool = False
+    allow_tf32: bool = False
+    use_8bit_adam: bool = False
+    lr_scheduler: str = 'constant'
+    # The number of training steps to run
+    max_train_steps: int = 1000000
+    # Max grad for clipping
+    max_grad_norm: float = 1.0
+    # How often to save checkpoints
+    checkpointing_steps: int = 5000
+    # The path to resume from
+    resume_from_path: Optional[Path] = None
+    # The step to resume from, mainly for logging
+    resume_from_step: Optional[int] = None
+    # Lora mode, untested
+    lora_rank: Optional[int] = None
+    # Which operator to train
+    mode: str = 'texture'
+    # The path to the textures dataset if used
+    textures_dir: Optional[Path] = None
+    # The path to the backgrounds dataset if used
+    backgrounds_dir: Optional[Path] = None
+    # optional directory of plain images to use for unconditional denoising
+    randoms_dir: Optional[Path] = None
+    # Whether full model is trained or only some layers, x_layers is the format for training only x layers
+    training_mode: str = 'full'
+    # Whether to use clip loss
+    use_clip_loss: bool = False
+    # Clip lambda
+    clip_strength: float = 10.0