nodes/smthemex-ComfyUI_Pops
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source 
add
This commit is contained in:
smthemex
2024-06-12 19:46:59 +08:00
committed by GitHub
parent d7bb8fb101
commit 1f93a40b87
3 changed files with 1505 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/coach.py
+720
View File
@@ -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
scripts/dataset.py
+696
View File
@@ -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
scripts/train_config.py
+89
View File
@@ -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