search

eladrich / pixel2style2pixel

Official Implementation for "Encoding in Style: a StyleGAN Encoder for Image-to-Image Translation" (CVPR 2021) presenting the pixel2style2pixel (pSp) framework
https://eladrich.github.io/pixel2style2pixel/
MIT License
3.19k stars 570 forks source link

Which part of the code is related to the generator? #171

Closed zbw0329 closed 3 years ago

zbw0329 commented 3 years ago

Hi!Thanks for your great work! I notice that the 'Generator' is defined in ‘stylegan2.model.py',and called by the 'psp.py'. But I don't see the relevant code used in the 'coach.train()' function. So I wonder which part of the codes calls generator to sample the pictures,when I was training the model? I want to change the code to use my own generator,which part of codes should I change?

yuval-alaluf commented 3 years ago

The Generator class is wrapped in the pSp class as the decoder. That is why you see no direct call to the class in coach.train(). If you wish to change the generator class you should change the following lines to call your class: For initializing the decoder: https://github.com/eladrich/pixel2style2pixel/blob/3e2ebd82bb3e3719507898f365abcec16597061a/models/psp.py#L31 And when calling the decoder's forward: https://github.com/eladrich/pixel2style2pixel/blob/3e2ebd82bb3e3719507898f365abcec16597061a/models/psp.py#L94-L97

Just note that the current decoder is designed to work with StyleGAN that operates on 18 inputs. If you change the decoder to something very different, you may need to change the encoder logic as well.

zbw0329 commented 3 years ago

OK, I see. Thanks for your reply