fasterr

app.py

@@ -12,41 +12,10 @@ import subprocess
 def check_gpu():
     return tf.test.is_gpu_available(cuda_only=False, min_cuda_compute_capability=None)
 
-model_path = 'best_net.pkl'
-#define load model functions
-_cached_networks = dict()
-def load_networks(path):
-    if path in _cached_networks:
-        return _cached_networks[path]
-    stream = open(path, 'rb')
-    tflib.init_tf()
-    with stream:
-        G, D, Gs = pickle.load(stream, encoding='latin1')
-    _cached_networks[path] = G, D, Gs
-    return G, D, Gs
-
-# Code to load the StyleGAN2 Model
-def load_model():
-    _G, _D, Gs = load_networks(model_path)
-    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
-#define helper functions
-def get_control_latent_vectors(path):
-    files = [x for x in Path(path).iterdir() if str(x).endswith('.npy')]
-    latent_vectors = {f.name[:-4]:np.load(f) for f in files}
-    return latent_vectors
-
-#load latent directions
-latent_controls = get_control_latent_vectors('trajectories/')
-
 def generate_image_from_projected_latents(latent_vector):
     images = Gs.components.synthesis.run(latent_vector, **Gs_kwargs)
     return images
 
-
 def frame_to_frame(latent_code):
     modified_latent_code = np.copy(latent_code)
     full_video = [generate_image_from_projected_latents(modified_latent_code)]
@@ -56,32 +25,50 @@ def frame_to_frame(latent_code):
         full_video.append(ims)
     return np.array(full_video).squeeze()
 
-#load the model
-Gs, noise_vars, Gs_kwargs = load_model()
-#select a random latent code
-rnd = np.random.RandomState(3)
-z = rnd.randn(1, *Gs.input_shape[1:])
-noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]
-tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars})
-random_img_latent_code = Gs.components.mapping.run(z,None)
+@st.cache(allow_output_mutation=True)  # Cache to avoid reloading the model every time
+def load_initial_setup():
+    stream = open('best_net.pkl', 'rb')
+    tflib.init_tf()
+    sess=tf.get_default_session()
+
+    with stream:
+        G, D, Gs = pickle.load(stream, encoding='latin1')
+    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
+    #load latent directions
+    files = [x for x in Path('trajectories/').iterdir() if str(x).endswith('.npy')]
+    latent_controls = {f.name[:-4]:np.load(f) for f in files}
+    #select a random latent code
+    noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]
+    rnd = np.random.RandomState()
 
-#make it be ED frame
-random_img_latent_code -= 0.7*latent_controls['time']
-vid = frame_to_frame(random_img_latent_code)
+    tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars})
 
+    return Gs, Gs_kwargs, latent_controls, sess
 
-temp_video_path="output.mp4"
-writer=imageio.get_writer(temp_video_path, fps=20)
-for i in range(vid.shape[0]):
-    frame = vid[i]
-    writer.append_data(frame)
+if __name__=="__main__":
+    rnd = np.random.RandomState()
+    Gs, Gs_kwargs, latent_controls, sess = load_initial_setup()
+    with sess.as_default():
+        z = rnd.randn(1, *Gs.input_shape[1:])
+        random_img_latent_code = Gs.components.mapping.run(z,None)
+        #make it be ED frame
+        random_img_latent_code -= 0.7*latent_controls['time']
+        vid = frame_to_frame(random_img_latent_code)
 
-writer.close()
-out_path = "fixed_out.mp4"
-command = ["ffmpeg", "-i", temp_video_path, "-vcodec", "libx264", out_path]
-subprocess.run(command)
+        temp_video_path="output.mp4"
+        writer=imageio.get_writer(temp_video_path, fps=20)
+        for i in range(vid.shape[0]):
+            frame = vid[i]
+            writer.append_data(frame)
 
-st.video(out_path)
+        writer.close()
+        out_path = "fixed_out.mp4"
+        command = ["ffmpeg", "-i", temp_video_path, "-vcodec", "libx264", out_path]
+        subprocess.run(command)
 
-os.remove(temp_video_path)
-os.remove(out_path)
+        st.video(out_path)
+        os.remove(temp_video_path)
+        os.remove(out_path)