Add first successful sampling implementation

This commit is contained in:
Gadersd
2023-08-04 17:01:44 -04:00
parent 3bf5b8c477
commit 77f30aefa7
9 changed files with 42 additions and 34 deletions

View File

@@ -1,3 +1,4 @@
import pathlib
from autoencoder import save_autoencoder
from unet import save_unet_model
from clip import save_clip_text_transformer
@@ -5,8 +6,9 @@ from clip import save_clip_text_transformer
from save import save_scalar, save_tensor
def save_stable_diffusion(stable_diffusion, path):
pathlib.Path(path).mkdir(parents=True, exist_ok=True)
save_scalar(stable_diffusion.alphas_cumprod.shape[0], "n_steps", path)
save_tensor(stable_diffusion.alphas_cumprod, 'alphas_cumprod', path)
save_autoencoder(stable_diffusion.autoencoder, 'autoencoder', path)
save_unet_model(stable_diffusion.diffusion, 'unet', path)
save_clip_text_transformer(stable_diffusion.clip, 'clip', path)
save_autoencoder(stable_diffusion.first_stage_model, pathlib.Path(path, 'autoencoder'))
save_unet_model(stable_diffusion.model.diffusion_model, pathlib.Path(path, 'unet'))
save_clip_text_transformer(stable_diffusion.cond_stage_model.transformer.text_model, pathlib.Path(path, 'clip'))