Stylegan3 .pkl warning and conversion

Models
1

I guys, I trained my Stylegan3 model and I uploaded my .pkl file with boring pickle warning.

Any experience on stylegen upload or conversion here on HF?

Even converting in .pt I’ve the same warning :frowning:

import torch
import pickle
from safetensors.torch import save_file, load_file
from train.networks_stylegan3 import Generator
from torch_utils import gen_utils

model = ‘mymodel.pkl’
device = torch.device(‘cuda’)

Load the network

G = gen_utils.load_network(‘G_ema’, model, None, device)

Save the state dict of the model

state_dict = G.state_dict()

Save the state dict in SafeTensors format

save_file(state_dict, ‘model.safetensors’)

Load the state dict from SafeTensors format

state_dict = load_file(‘mode.safetensors’)

Initialize a new model and load the state dict

G_kwargs = {
“z_dim”: 512,
“c_dim”: 0,
“w_dim”: 512,
“img_resolution”: 1024,
“img_channels”: 3,
“mapping_kwargs”: {
“num_layers”: 2,
“freeze_layers”: 0,
“freeze_embed”: False
},
“synthesis_kwargs”: {
“channel_base”: 32768,
“channel_max”: 512,
“magnitude_ema_beta”: 0.9988915792636801
}
}

G = Generator(
z_dim=G_kwargs[‘z_dim’],
c_dim=G_kwargs[‘c_dim’],
w_dim=G_kwargs[‘w_dim’],
img_resolution=G_kwargs[‘img_resolution’],
img_channels=G_kwargs[‘img_channels’],
mapping_kwargs=G_kwargs[‘mapping_kwargs’],
**G_kwargs[‘synthesis_kwargs’]
).to(device)

G.load_state_dict(state_dict)

torch.save(G, ‘mynewmodel.pt’)