Big refactor of SD runner and added image generator

This commit is contained in:
Jaret Burkett
2023-08-03 14:51:25 -06:00
parent 75ec5d9292
commit 66c6f0f6f7
16 changed files with 923 additions and 430 deletions

View File

@@ -202,9 +202,11 @@ class TrainSDRescaleProcess(BaseSDTrainProcess):
)
# get noise
noise = self.get_latent_noise(
noise = self.sd.get_latent_noise(
pixel_height=self.rescale_config.from_resolution,
pixel_width=self.rescale_config.from_resolution,
batch_size=self.train_config.batch_size,
noise_offset=self.train_config.noise_offset,
).to(self.device_torch, dtype=dtype)
torch.set_default_device(self.device_torch)
@@ -238,7 +240,7 @@ class TrainSDRescaleProcess(BaseSDTrainProcess):
)
with torch.no_grad():
noise_pred_target = self.predict_noise(
noise_pred_target = self.sd.predict_noise(
latents,
text_embeddings=text_embeddings,
timestep=timestep,
@@ -256,7 +258,7 @@ class TrainSDRescaleProcess(BaseSDTrainProcess):
with self.network:
assert self.network.is_active
self.network.multiplier = 1.0
noise_pred_train = self.predict_noise(
noise_pred_train = self.sd.predict_noise(
reduced_latents,
text_embeddings=text_embeddings,
timestep=timestep,