storytold/FineTrainers-Conditioning
2
0
Fork 0
mirror of https://github.com/storytold/FineTrainers-Conditioning.git synced 2026-10-24 08:25:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

other condtions -> text conditions

Browse Source
This commit is contained in:
Aryan
2024-12-23 13:48:58 +01:00
parent 363841ae3f
commit 86e1d1888a
2 changed files with 24 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
finetrainers/dataset.py
+4 -4
View File
@@ -271,9 +271,9 @@ class PrecomputedDataset(Dataset):
self.conditions_path = self.data_root / PRECOMPUTED_DIR_NAME / PRECOMPUTED_CONDITIONS_DIR_NAME self.conditions_path = self.data_root / PRECOMPUTED_DIR_NAME / PRECOMPUTED_CONDITIONS_DIR_NAME
self.latent_conditions = sorted(os.listdir(self.latents_path)) self.latent_conditions = sorted(os.listdir(self.latents_path))
self.other_conditions = sorted(os.listdir(self.conditions_path)) self.text_conditions = sorted(os.listdir(self.conditions_path))
assert len(self.latent_conditions) == len(self.other_conditions), "Number of captions and videos do not match" assert len(self.latent_conditions) == len(self.text_conditions), "Number of captions and videos do not match"
def __len__(self) -> int: def __len__(self) -> int:
return len(self.latent_conditions) return len(self.latent_conditions)
@@ -281,9 +281,9 @@ class PrecomputedDataset(Dataset):
def __getitem__(self, index: int) -> Dict[str, Any]: def __getitem__(self, index: int) -> Dict[str, Any]:
conditions = {} conditions = {}
latent_path = self.latents_path / self.latent_conditions[index] latent_path = self.latents_path / self.latent_conditions[index]
condition_path = self.conditions_path / self.other_conditions[index] condition_path = self.conditions_path / self.text_conditions[index]
conditions["latent_conditions"] = torch.load(latent_path, map_location="cpu", weights_only=True) conditions["latent_conditions"] = torch.load(latent_path, map_location="cpu", weights_only=True)
conditions["other_conditions"] = torch.load(condition_path, map_location="cpu", weights_only=True) conditions["text_conditions"] = torch.load(condition_path, map_location="cpu", weights_only=True)
return conditions return conditions
finetrainers/trainer.py
+20 -21
View File
@@ -187,25 +187,25 @@ class Trainer:
def collate_fn(batch): def collate_fn(batch):
latent_conditions = [x["latent_conditions"] for x in batch] latent_conditions = [x["latent_conditions"] for x in batch]
other_conditions = [x["other_conditions"] for x in batch] text_conditions = [x["text_conditions"] for x in batch]
batched_latent_conditions = {} batched_latent_conditions = {}
batched_other_conditions = {} batched_text_conditions = {}
for key in list(latent_conditions[0].keys()): for key in list(latent_conditions[0].keys()):
if torch.is_tensor(latent_conditions[0][key]): if torch.is_tensor(latent_conditions[0][key]):
batched_latent_conditions[key] = torch.cat([x[key] for x in latent_conditions], dim=0) batched_latent_conditions[key] = torch.cat([x[key] for x in latent_conditions], dim=0)
else: else:
# TODO(aryan): implement batch sampler for precomputed latents # TODO(aryan): implement batch sampler for precomputed latents
batched_latent_conditions[key] = [x[key] for x in latent_conditions][0] batched_latent_conditions[key] = [x[key] for x in latent_conditions][0]
for key in list(other_conditions[0].keys()): for key in list(text_conditions[0].keys()):
if torch.is_tensor(other_conditions[0][key]): if torch.is_tensor(text_conditions[0][key]):
batched_other_conditions[key] = torch.cat([x[key] for x in other_conditions], dim=0) batched_text_conditions[key] = torch.cat([x[key] for x in text_conditions], dim=0)
else: else:
# TODO(aryan): implement batch sampler for precomputed latents # TODO(aryan): implement batch sampler for precomputed latents
batched_other_conditions[key] = [x[key] for x in other_conditions][0] batched_text_conditions[key] = [x[key] for x in text_conditions][0]
return {"latent_conditions": batched_latent_conditions, "other_conditions": batched_other_conditions} return {"latent_conditions": batched_latent_conditions, "text_conditions": batched_text_conditions}
should_recompute = should_perform_precomputation(self.args.data_root) should_precompute = should_perform_precomputation(self.args.data_root)
if not should_recompute: if not should_precompute:
logger.info("Precomputed conditions and latents found. Loading precomputed data.") logger.info("Precomputed conditions and latents found. Loading precomputed data.")
self.dataloader = torch.utils.data.DataLoader( self.dataloader = torch.utils.data.DataLoader(
PrecomputedDataset(self.args.data_root), PrecomputedDataset(self.args.data_root),
@@ -259,7 +259,7 @@ class Trainer:
f"Precomputing conditions and latents for batch {i + 1}/{len(self.dataset)} on process {self.state.accelerator.process_index}" f"Precomputing conditions and latents for batch {i + 1}/{len(self.dataset)} on process {self.state.accelerator.process_index}"
) )
other_conditions = self.model_config["prepare_conditions"]( text_conditions = self.model_config["prepare_conditions"](
tokenizer=self.tokenizer, tokenizer=self.tokenizer,
tokenizer_2=self.tokenizer_2, tokenizer_2=self.tokenizer_2,
tokenizer_3=self.tokenizer_3, tokenizer_3=self.tokenizer_3,
@@ -271,7 +271,7 @@ class Trainer:
dtype=self.state.weight_dtype, dtype=self.state.weight_dtype,
) )
filename = conditions_dir / f"conditions-{i}-{index}.pt" filename = conditions_dir / f"conditions-{i}-{index}.pt"
torch.save(other_conditions, filename.as_posix()) torch.save(text_conditions, filename.as_posix())
index += 1 index += 1
progress_bar.update(1) progress_bar.update(1)
self._delete_components() self._delete_components()
@@ -596,7 +596,7 @@ class Trainer:
dtype=weight_dtype, dtype=weight_dtype,
generator=generator, generator=generator,
) )
other_conditions = self.model_config["prepare_conditions"]( text_conditions = self.model_config["prepare_conditions"](
tokenizer=self.tokenizer, tokenizer=self.tokenizer,
text_encoder=self.text_encoder, text_encoder=self.text_encoder,
tokenizer_2=self.tokenizer_2, tokenizer_2=self.tokenizer_2,
@@ -607,27 +607,27 @@ class Trainer:
) )
else: else:
latent_conditions = batch["latent_conditions"] latent_conditions = batch["latent_conditions"]
other_conditions = batch["other_conditions"] text_conditions = batch["text_conditions"]
latent_conditions["latents"] = DiagonalGaussianDistribution( latent_conditions["latents"] = DiagonalGaussianDistribution(
latent_conditions["latents"] latent_conditions["latents"]
).sample(generator) ).sample(generator)
if "post_latent_preparation" in self.model_config.keys(): if "post_latent_preparation" in self.model_config.keys():
latent_conditions = self.model_config["post_latent_preparation"](**latent_conditions) latent_conditions = self.model_config["post_latent_preparation"](**latent_conditions)
align_device_and_dtype(latent_conditions, accelerator.device, weight_dtype) align_device_and_dtype(latent_conditions, accelerator.device, weight_dtype)
align_device_and_dtype(other_conditions, accelerator.device, weight_dtype) align_device_and_dtype(text_conditions, accelerator.device, weight_dtype)
batch_size = latent_conditions["latents"].shape[0] batch_size = latent_conditions["latents"].shape[0]
latent_conditions = make_contiguous(latent_conditions) latent_conditions = make_contiguous(latent_conditions)
other_conditions = make_contiguous(other_conditions) text_conditions = make_contiguous(text_conditions)
if self.args.caption_dropout_technique == "zero": if self.args.caption_dropout_technique == "zero":
if random.random() < self.args.caption_dropout_p: if random.random() < self.args.caption_dropout_p:
other_conditions["prompt_embeds"].fill_(0) text_conditions["prompt_embeds"].fill_(0)
other_conditions["prompt_attention_mask"].fill_(False) text_conditions["prompt_attention_mask"].fill_(False)
# TODO(aryan): refactor later # TODO(aryan): refactor later
if "pooled_prompt_embeds" in other_conditions: if "pooled_prompt_embeds" in text_conditions:
other_conditions["pooled_prompt_embeds"].fill_(0) text_conditions["pooled_prompt_embeds"].fill_(0)
# These weighting schemes use a uniform timestep sampling and instead post-weight the loss # These weighting schemes use a uniform timestep sampling and instead post-weight the loss
weights = compute_density_for_timestep_sampling( weights = compute_density_for_timestep_sampling(
@@ -650,14 +650,13 @@ class Trainer:
noisy_latents = (1.0 - sigmas) * latent_conditions["latents"] + sigmas * noise noisy_latents = (1.0 - sigmas) * latent_conditions["latents"] + sigmas * noise
latent_conditions.update({"noisy_latents": noisy_latents}) latent_conditions.update({"noisy_latents": noisy_latents})
other_conditions.update({"timesteps": timesteps})
# These weighting schemes use a uniform timestep sampling and instead post-weight the loss # These weighting schemes use a uniform timestep sampling and instead post-weight the loss
weights = compute_loss_weighting_for_sd3( weights = compute_loss_weighting_for_sd3(
weighting_scheme=self.args.flow_weighting_scheme, sigmas=sigmas weighting_scheme=self.args.flow_weighting_scheme, sigmas=sigmas
) )
pred = self.model_config["forward_pass"]( pred = self.model_config["forward_pass"](
transformer=self.transformer, **latent_conditions, **other_conditions transformer=self.transformer, timesteps=timesteps, **latent_conditions, **text_conditions
) )
target = noise - latent_conditions["latents"] target = noise - latent_conditions["latents"]