danielroich / PTI

Official Implementation for "Pivotal Tuning for Latent-based editing of Real Images" (ACM TOG 2022) https://arxiv.org/abs/2106.05744
MIT License
897 stars 112 forks source link

Unexpected size of W features in generated_images = self.G.synthesis(w, noise_mode='const', force_fp32=True) #40

Closed usmancheema89 closed 1 year ago

usmancheema89 commented 2 years ago

I am using a self-trained model trained using the StyleGan-ada pytorch repository. While using use_multi_id_training=True I get a size mismatch error in the forward call of G.synthesis.

The full trace is shown:

Traceback (most recent call last): File "test.py", line 6, in run_PTI(use_multi_id_training=True) File "/home/usman/Documents/Work/PTI/scripts/run_pti.py", line 44, in run_PTI coach.train() File "/home/usman/Documents/Work/PTI/training/coaches/multi_id_coach.py", line 58, in train generated_images = self.forward(w_pivot) File "/home/usman/Documents/Work/PTI/training/coaches/base_coach.py", line 130, in forward generated_images = self.G.synthesis(w, noise_mode='const', force_fp32=True) File "/home/usman/anaconda3/envs/eg3d/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl return forward_call(*input, **kwargs) File "", line 460, in forward File "/home/usman/Documents/Work/PTI/torch_utils/misc.py", line 93, in assert_shape raise AssertionError(f'Wrong size for dimension {idx}: got {size}, expected {ref_size}') AssertionError: Wrong size for dimension 1: got 18, expected 16

usmancheema89 commented 2 years ago

Think I figured out the issue. In w_projector.py change:

del G return w_opt.repeat([1, 18, 1])

to

G_map_num_ws = G.mapping.num_ws del G return w_opt.repeat([1, G_map_num_ws, 1])

danielroich commented 1 year ago

Great! The code correction you suggested is actually the right way to do this. It makes the code agnostic of the generator number of layers.