app.py

import streamlit as st
import tensorflow as tf
import pickle
import numpy as np
from pathlib import Path
import dnnlib
from dnnlib import tflib
import imageio
import os
import subprocess

def check_gpu():
    return tf.test.is_gpu_available(cuda_only=False, min_cuda_compute_capability=None)

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)]
    for i in range(49):
        modified_latent_code = modified_latent_code +  latent_controls[f'{i}{i+1}']
        ims = generate_image_from_projected_latents(modified_latent_code)
        full_video.append(ims)
    return np.array(full_video).squeeze()

@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()

    tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars})

    return Gs, Gs_kwargs, latent_controls, sess

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)

        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)

        writer.close()
        out_path = "fixed_out.mp4"
        command = ["ffmpeg", "-i", temp_video_path, "-vcodec", "libx264", out_path]
        subprocess.run(command)

        st.video(out_path)
        os.remove(temp_video_path)
        os.remove(out_path)