Mirrors
/
waifu-diffusion
mirror of https://github.com/harubaru/waifu-diffusion.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Access Underlying Model

This commit is contained in:
cafeai 2022-12-03 19:57:00 +09:00
parent 3cefb57fc6
commit 34715bcc97
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
trainer/diffusers_trainer.py
View File

@ -500,6 +500,9 @@ class AspectDataset(torch.utils.data.Dataset):
self.device = device self.device = device
self.ucg = ucg self.ucg = ucg
if type(self.text_encoder) is torch.nn.parallel.DistributedDataParallel:
self.text_encoder = self.text_encoder.module
self.transforms = torchvision.transforms.Compose([ self.transforms = torchvision.transforms.Compose([
torchvision.transforms.RandomHorizontalFlip(p=0.5), torchvision.transforms.RandomHorizontalFlip(p=0.5),
torchvision.transforms.ToTensor(), torchvision.transforms.ToTensor(),