ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-29 10:41:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

Moved SD batch processing to a shared method and added it for use in slider training. Still testing if it affects quality over sampling

Browse Source
This commit is contained in:
Jaret Burkett
2023-08-26 08:55:00 -06:00
parent aeaca13d69
commit 3367ab6b2c
6 changed files with 178 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
toolkit/llvae.py
View File

@@ -101,14 +101,16 @@ if __name__ == '__main__':
from PIL import Image
import torchvision.transforms as transforms
user_path = os.path.expanduser('~')
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
dtype = torch.float32
input_path = os.path.join(user_path, "Pictures/sample_2_512.png")
output_path = os.path.join(user_path, "Pictures/sample_2_512_llvae.png")
img = Image.open(input_path)
img_tensor = transforms.ToTensor()(img)
img_tensor = img_tensor.unsqueeze(0)
img_tensor = img_tensor.unsqueeze(0).to(device=device, dtype=dtype)
print("input_shape: ", list(img_tensor.shape))
vae = LosslessLatentVAE(in_channels=3, latent_depth=8)
vae = LosslessLatentVAE(in_channels=3, latent_depth=8, dtype=dtype).to(device=device, dtype=dtype)
latent = vae.encode(img_tensor)
print("latent_shape: ", list(latent.shape))
out_tensor = vae.decode(latent)