Bug encountered when running generate.py

[gene-rative]

Hi,

I encounter the following bug when running generate.py (command included in the snapshot). I guess the reason is that the current code is inconsistent with the persistent class stored in the pickle.

(base) user@xxx:/data/user/Diffusion-GAN/diffusion-stylegan2$ python generate.py --outdir=out --seeds=1-100 \
>     --network=https://tsciencescu.blob.core.windows.net/projectshzheng/DiffusionGAN/diffusion-stylegan2-ffhq.pkl
Loading networks from "https://tsciencescu.blob.core.windows.net/projectshzheng/DiffusionGAN/diffusion-stylegan2-ffhq.pkl"...
Downloading https://tsciencescu.blob.core.windows.net/projectshzheng/DiffusionGAN/diffusion-stylegan2-ffhq.pkl ... done
Traceback (most recent call last):
  File "/data/user/Diffusion-GAN/diffusion-stylegan2/generate.py", line 127, in <module>
    generate_images() # pylint: disable=no-value-for-parameter
  File "/home/user/anaconda3/lib/python3.9/site-packages/click/core.py", line 1128, in __call__
    return self.main(*args, **kwargs)
  File "/home/user/anaconda3/lib/python3.9/site-packages/click/core.py", line 1053, in main
    rv = self.invoke(ctx)
  File "/home/user/anaconda3/lib/python3.9/site-packages/click/core.py", line 1395, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/home/user/anaconda3/lib/python3.9/site-packages/click/core.py", line 754, in invoke
    return __callback(*args, **kwargs)
  File "/home/user/anaconda3/lib/python3.9/site-packages/click/decorators.py", line 26, in new_func
    return f(get_current_context(), *args, **kwargs)
  File "/data/user/Diffusion-GAN/diffusion-stylegan2/generate.py", line 84, in generate_images
    G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore
  File "/data/user/Diffusion-GAN/diffusion-stylegan2/legacy.py", line 21, in load_network_pkl
    data = _LegacyUnpickler(f).load()
  File "/data/user/Diffusion-GAN/diffusion-stylegan2/legacy.py", line 71, in find_class
    return super().find_class(module, name)
ModuleNotFoundError: No module named 'training.augment'

[Zhendong-Wang]

Hi,

Thanks for pointing out the issue. You are correct! I changed the file name and didn't notice that the saved pkl is based on the old filename. I have uploaded the old file training/augment.py for loading pkl use, as a currently quick fix. Will try to update the pkl as well.

The problem is solved. :)

[Zhendong-Wang]

Hi,

I fixed the issue in another way by ignoring the load of augpipe class, which is not used during generaton. You could simply git pull and the problem is fixed.