Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,568 Bytes
f2f17f4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import argparse
import os
import torch
from diffusers import AutoencoderKL
from torch import nn
from torch.optim import Adam
from utils import load_image, save_image
def main(args):
os.makedirs(args.out_dir, exist_ok=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
vae = AutoencoderKL.from_pretrained(args.vae_model_path).to(
device, dtype=torch.float32
)
vae.requires_grad_(False)
image = load_image(args.image_path, size=(512, 512)).to(device, dtype=torch.float32)
image = image * 2 - 1
save_image(image / 2 + 0.5, f"{args.out_dir}/ori_image.png")
latents = vae.encode(image)["latent_dist"].mean
save_image(latents, f"{args.out_dir}/latents.png")
rec_image = vae.decode(latents, return_dict=False)[0]
save_image(rec_image / 2 + 0.5, f"{args.out_dir}/rec_image.png")
for param in vae.decoder.parameters():
param.requires_grad = True
loss_fn = nn.L1Loss()
optimizer = Adam(vae.decoder.parameters(), lr=args.learning_rate)
# Training loop
for epoch in range(args.num_epochs):
reconstructed = vae.decode(latents, return_dict=False)[0]
loss = loss_fn(reconstructed, image)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1}/{args.num_epochs}, Loss: {loss.item()}")
rec_image = vae.decode(latents, return_dict=False)[0]
save_image(rec_image / 2 + 0.5, f"{args.out_dir}/trained_rec_image.png")
vae.save_pretrained(
f"{args.out_dir}/trained_vae_{os.path.basename(args.image_path)}"
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Train a VAE with given image and settings."
)
# Add arguments
parser.add_argument(
"--out_dir",
type=str,
default="./trained_vae/",
help="Output directory to save results",
)
parser.add_argument(
"--vae_model_path",
type=str,
required=True,
help="Path to the pretrained VAE model",
)
parser.add_argument(
"--image_path", type=str, required=True, help="Path to the input image"
)
parser.add_argument(
"--learning_rate",
type=float,
default=1e-4,
help="Learning rate for the optimizer",
)
parser.add_argument(
"--num_epochs", type=int, default=75, help="Number of training epochs"
)
args = parser.parse_args()
main(args)
|