Fixes stylegan import

stylegan3/__init__.py

@@ -0,0 +1,9 @@
+import sys
+
+sys.path.append("../stylegan3")
+
+from dnnlib.util import open_url
+from legacy import load_network_pkl
+
+sys.path.remove("../stylegan3")
+

test.py

@@ -1,5 +0,0 @@
-from stylegan3 import dnnlib, legacy
-
-network = 'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl'
-f = dnnlib.util.open_url(network)
-G = legacy.load_network_pkl(f)

visualize.py

@@ -4,8 +4,7 @@ import moviepy.editor as mpy
 import random
 import torch
 from tqdm import tqdm
-import stylegan3.dnnlib
-import stylegan3.legacy
+import stylegan3
 
 
 def visualize(audio_file, network, truncation, batch_size, *args, **kwargs):
@@ -34,8 +33,8 @@ def visualize(audio_file, network, truncation, batch_size, *args, **kwargs):
 
     # Load pre-trained model
     device = torch.device('cuda')
-    with stylegan3.dnnlib.open_url(network) as f:
-        G = stylegan3.legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore
+    with stylegan3.open_url(network) as f:
+        G = stylegan3.load_network_pkl(f)['G_ema'].to(device) # type: ignore
         G.eval()
 
     with torch.no_grad():