Caching a stylegan2-ada model?

Hello,

Been poking around Streamlit for the first time with a GAN. I tried to use the Demo GAN app in the galley as a template. However, I am not experienced with caching, hashmaps, etc…

This is what I have tried. I assume I need to move the call of ‘tflib.init_tf()’ to this hash_funcs. But nothing has worked.

TL_GAN_HASH_FUNCS = {
tflib.init_tf():id
}
@st.cache(allow_output_mutation=True, hash_funcs=TL_GAN_HASH_FUNCS)
def load_model(network_pkl):
print(‘Loading networks from “%s”…’ % network_pkl)
with dnnlib.util.open_url(network_pkl) as fp:
_G, _D, Gs = pickle.load(fp)
noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith(‘noise’)] Gs_kwargs = dnnlib.EasyDict()
Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
Gs_kwargs.randomize_noise = False
return Gs, noise_vars, Gs_kwargs

Has anyone had luck caching their stylegan GANs?

Hello Frewil,

I’m facing the same issue. Were you able to find any solution?

No, unfortunately. Keep me posted if you get it :pensive:

I have spent good quality time on this but I wasn’t able to find a solution. So, I changed the codebase to PyTorch, and is working well.

For your issue, u can use
hash_funcs={’_thread.RLock’ : lambda _: None,
‘_thread.lock’ : lambda _: None,
‘builtins.PyCapsule’: lambda _: None,
‘_io.TextIOWrapper’ : lambda _: None,
‘builtins.weakref’: lambda _: None,
‘builtins.dict’ : lambda _:None,
}
as the hash function which can solve but I’m getting errors after that. Hope this fixes ur code.

1 Like

thanks for the tip!

I think I may be too lazy to convert to pytorch. I’ll wait it out for an update…