app.py

import streamlit as st
import torch
import imageio
import numpy as np
import time
from skimage.transform import resize
import warnings
import cv2
import subprocess
import os
from demo import load_checkpoints
from demo import make_animation
from skimage import img_as_ubyte
import shutil


def save_image_from_upload(uploaded_file):
    # Specify the path to save
    save_path = './uploaded_images'
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    # Open the file in the desired location with write-binary ('wb') mode
    with open(os.path.join(save_path, uploaded_file.name), "wb") as f:
        f.write(uploaded_file.getbuffer())  # Write the file to the specified location
        #st.success(f'Saved file {uploaded_file.name} in {save_path}')
    try:
        shutil.copy2(os.path.join(save_path, uploaded_file.name),
                     st.session_state['source_image_path'])
        print("File copied successfully.")
    except FileNotFoundError:
        print("The source file was not found.")
    except PermissionError:
        print("Permission denied.")
    except Exception as e:
        print(f"Error occurred: {e}")
def save_video_from_upload(uploaded_file):
    # Specify the path to save
    save_path = './uploaded_videos'
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    # Open the file in the desired location with write-binary ('wb') mode
    with open(os.path.join(save_path, uploaded_file.name), "wb") as f:
        f.write(uploaded_file.getbuffer())  # Write the file to the specified location
        # st.success(f'Saved file {uploaded_file.name} in {save_path}')
    try:
        shutil.copy2(os.path.join(save_path, uploaded_file.name),
                     st.session_state['driving_video_path'])
        print("File copied successfully.")
    except FileNotFoundError:
        print("The source file was not found.")
    except PermissionError:
        print("Permission denied.")
    except Exception as e:
        print(f"Error occurred: {e}")

def create_image_video_side_by_side(source, driving, generated=None, output_file='assets/output_video.mp4', fps=20, progress_bar=None):
    #st.image(source,caption='create_image_video_side_by_side src')
    total_driving_frames = len(driving)
    progress_bar.progress(0)
    #images = l
    images = []
    print("going through video")

    for i in range(len(driving)):

        cols = [source]

        cols.append(driving[i])
        if generated is not None:
            #print("generated data length:"+str(len(generated)))
            #print("cols type"+str(type(cols[i])))
            #print("generated[i] shape" + str(generated[i].shape))
            cols.append(generated[i])
            #print("len(cols) afer append "+str(len(cols)))
        # else:
        #     print("generated is None!!!")

        # Concatenate the images horizontally
        full_image = np.concatenate(cols, axis=1)

        # Convert the image array to an RGB image
        full_image_rgb = np.clip(full_image * 255, 0, 255).astype(
            np.uint8) if full_image.max() <= 1 else full_image.astype(np.uint8)
        #print("full_image_rgb shape" + str(full_image_rgb.shape))
        # Append to the list of images
        # if i == 0:
        #     source_rgb = np.clip(source * 255, 0, 255).astype(
        #         np.uint8) if source.max() <= 1 else source.astype(np.uint8)
        #     st.image(source_rgb, caption="source_rgb")
        #     driving_0_rgb = np.clip(driving[i] * 255, 0, 255).astype(
        #         np.uint8) if driving[i].max() <= 1 else driving[i].astype(np.uint8)
        #     st.image(driving_0_rgb, caption="driving_0_rgb")
        #     cols_temp = [source_rgb]
        #     cols_temp.append(driving[i])
        #     full_image_temp = np.concatenate(cols_temp, axis=1)
        #     st.image(full_image_temp, caption="full_image_temp")


        images.append(full_image_rgb)
        progress_percentage = (i + 1) / total_driving_frames
        progress_bar.progress(progress_percentage)
    print("going through video done")
    # Determine the size of the frames
    height, width, layers = images[0].shape
    print("images[0].shape"+str(images[0].shape))
    size = (width, height)

    # Define the codec and create VideoWriter object
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 'mp4v' or 'XVID'
    delete_file_if_exists("temp_gen_video.mp4")
    out = cv2.VideoWriter("temp_gen_video.mp4", fourcc, fps, size)

    print("writing video start")
    for image in images:
        out.write(cv2.cvtColor(image, cv2.COLOR_RGB2BGR))


    out.release()  # Release the video writer
    print("writing video end")
    print("converting video to H264")
    delete_file_if_exists(output_file)
    command = ['ffmpeg', '-i', 'temp_gen_video.mp4', '-c:v', 'libx264', output_file]
    try:
        result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True)
        print("FFMPEG Output:\n", result.stdout)
        print("Video converted successfully.")
    except subprocess.CalledProcessError as e:
        # Output the error in case of failure
        print("Error during conversion:\n", e.stderr)
    print("converting video to H264 done")
    return output_file
def delete_file_if_exists(file_path):
    try:
        if os.path.exists(file_path):
            print(f"File {file_path} exists and will be deleted.")
            os.remove(file_path)
            print(f"File {file_path} has been deleted.")
        else:
            print(f"No file found at {file_path}, nothing to delete.")
    except OSError as e:
        print(f"Error deleting file {file_path}: {e}")


def read_src_image():
    print("reading src imaage")
    st.session_state['source_image'] = imageio.imread(st.session_state['source_image_path'])
    st.session_state['source_image'] = resize(st.session_state['source_image'],
                                              (st.session_state['pixel'], st.session_state['pixel']))[..., :3]

    # st.session_state['source_image'] = np.clip(st.session_state['source_image'] * 255, 0, 255).astype(
    #     np.uint8) if st.session_state['source_image'].max() <= 1 else st.session_state['source_image'].astype(np.uint8)
def read_driving_video(progress_bar=None):

    reader = imageio.get_reader(st.session_state['driving_video_path'])
    st.session_state['fps'] = reader.get_meta_data()['fps']
    st.session_state['duration'] = reader.get_meta_data()['duration']
    video_width = reader.get_meta_data()['source_size'][0]

    print("st.session_state['duration']="+str(st.session_state['duration']))
    st.session_state['driving_video'] = []
    print("reading video")

    # Calculate the number of frames
    estimated_frame_count = int(st.session_state['fps'] * st.session_state['duration'])
    print("estimated_frame_count="+str(estimated_frame_count))

    progress_bar.progress(0)
    try:
        video_frame_idx=0
        for im in reader:
            # im = np.clip(im * 255, 0, 255).astype(
            #     np.uint8) if im.max() <= 1 else im.astype(np.uint8)
            st.session_state['driving_video'].append(im)
            progress_percentage = (video_frame_idx) / (estimated_frame_count+1)
            # print(f"video_frame_idx = {video_frame_idx} estimated_frame_count={estimated_frame_count}")
            progress_bar.progress(progress_percentage)
            video_frame_idx = video_frame_idx+1
    except RuntimeError:
        pass
    reader.close()
    print("finished reading video")
    # st.session_state['driving_video'] = [resize(frame, (st.session_state['pixel'], st.session_state['pixel']))[..., :3]
    #                                      for frame in st.session_state['driving_video']]
    print("check resize width ="+str(video_width))
    # if video_width != 512:
    progress_bar.progress(0)
    # Process each frame, update progress bar along the way
    if video_width != 512:
        resized_frames = []
        num_frames=len(st.session_state['driving_video'])
        current_status_placeholder.write("resizing video")
        for i, frame in enumerate(st.session_state['driving_video']):
            if i == 0:
                print("frame.dtype="+str(frame.dtype))
            # Resize frame
            resized_frame = resize(frame, (st.session_state['pixel'], st.session_state['pixel']))[..., :3]
            if i == 0:
                print("resized_frame.dtype="+str(resized_frame.dtype))
            resized_frames.append(resized_frame)

            # Update progress bar
            progress_bar.progress((i + 1) / num_frames)

        # Update the session state with the resized frames
        st.session_state['driving_video'] = resized_frames
    else:
        for i, frame in enumerate(st.session_state['driving_video']):
            if frame.dtype != np.float64:
                # Convert to float64
                frame_float64 = frame.astype(np.float64)

                # Normalize the frame based on its original range
                if frame.dtype == np.uint8:
                    frame_normalized = frame_float64 / 255.0
                elif frame.dtype == np.uint16:
                    frame_normalized = frame_float64 / 65535.0
                elif frame.dtype == np.float32:
                    # Assuming float32 range is 0.0 to 1.0, similar normalization might not be needed
                    frame_normalized = frame_float64
                st.session_state['driving_video'][i] = frame_normalized
def add_animation_to_image():

    inference_status_placeholder.write("start inference")
    print("device=" + str(st.session_state['device']))
    predictions = make_animation(st.session_state['source_image'], st.session_state['driving_video'], st.session_state['inpainting'], st.session_state['kp_detector'], st.session_state['dense_motion_network'],
                                     st.session_state['avd_network'], device=st.session_state['device'], mode=st.session_state['predict_mode'], progress_bar=create_animation_progress_bar)
    inference_status_placeholder.write("inference done")
    # save resulting video
    st.session_state['output_video_path']='assets/generated_video.mp4'
    st.session_state['side_by_side_with_generated_video_path']='assets/src_image_driving_video_generated_video_side_by_side.mp4'
    inference_status_placeholder.write("saving generated video")
    # for i in range(len(predictions)):
    #     predictions[i] = np.clip(predictions[i] * 255, 0, 255).astype(
    #         np.uint8) if predictions[i].max() <= 1 else predictions[i].astype(np.uint8)
    #st.image(predictions[0], caption="predictions[0]")
    imageio.mimsave(st.session_state['output_video_path'], [img_as_ubyte(frame) for frame in predictions], fps=st.session_state['fps'])
    inference_status_placeholder.write("saving generated video done")
    print("side_by_side_with_generated_video_path="+st.session_state['side_by_side_with_generated_video_path'])
    inference_status_placeholder.write("creating side by side video")
    st.session_state['side_by_side_with_generated_video_path'] = create_image_video_side_by_side(st.session_state['source_image'],
                                                                                         st.session_state['driving_video'], generated=predictions,
                                                                                         output_file=st.session_state['side_by_side_with_generated_video_path'], fps=st.session_state['fps'],progress_bar=create_animation_progress_bar)
    inference_status_placeholder.write("creating side by side video done")
def is_new_src_image_upload(uploaded_file):
    if 'last_src_image_uploaded_file' in st.session_state:
        # Check if the newly uploaded file is different from the last one
        if (uploaded_file.name != st.session_state.last_src_image_uploaded_file['name'] or
                uploaded_file.size != st.session_state.last_src_image_uploaded_file['size']):
            st.session_state.last_src_image_uploaded_file = {'name': uploaded_file.name, 'size': uploaded_file.size}
            # st.write("A new src image file has been uploaded.")
            return True
        else:
            # st.write("The same src image file has been re-uploaded.")
            return False
    else:
        # st.write("This is the first file upload detected.")
        st.session_state.last_src_image_uploaded_file = {'name': uploaded_file.name, 'size': uploaded_file.size}
        return True
    # Store current file details in session state

def is_new_driving_video_upload(uploaded_file):
    if 'last_driving_video_uploaded_file' in st.session_state:
        # Check if the newly uploaded file is different from the last one
        if (uploaded_file.name != st.session_state.last_driving_video_uploaded_file['name'] or
                uploaded_file.size != st.session_state.last_driving_video_uploaded_file['size']):
            st.session_state.last_driving_video_uploaded_file = {'name': uploaded_file.name, 'size': uploaded_file.size}
            # st.write("A new driving video file has been uploaded.")
            return True
        else:
            # st.write("The same driving video file has been re-uploaded.")
            return False
    else:
        # st.write("This is the first file upload detected.")
        st.session_state.last_driving_video_uploaded_file = {'name': uploaded_file.name, 'size': uploaded_file.size}
        return True

big_text = """
    <div style='text-align: center;'>
        <h1 style='font-size: 30x;'>Add motions to still images</h1>
    </div>
    """
    # Display the styled text
st.markdown(big_text, unsafe_allow_html=True)
#st.markdown("<h1>Add motions to still images</h1>")
current_status_placeholder = st.empty()
init_progress_bar = st.progress(0)

if 'is_initialized' not in st.session_state:
    st.session_state['is_initialized'] = True
    #st.set_option('enableStaticServing ', True)
    print("init")
    warnings.filterwarnings("ignore")
    current_status_placeholder.write("checking CUDA availability")
    if torch.cuda.is_available():
        print("CUDA is available on the following devices:")
        # Loop through available CUDA devices
        for i in range(torch.cuda.device_count()):
            print(f"Device {i}: {torch.cuda.get_device_name(i)}")
            current_status_placeholder.write(f"Device {i}: {torch.cuda.get_device_name(i)}")
    else:
        print("CUDA is not available. Listing CPU only.")
        print("Device 0: CPU")
        current_status_placeholder.write("CUDA is not available. Listing CPU only.")
    st.session_state['device'] = torch.device('cuda:0')
    st.session_state['dataset_name'] = 'vox'  # ['vox', 'taichi', 'ted', 'mgif']
    st.session_state['source_image_path'] = 'assets/src_image.png'
    st.session_state['driving_video_path'] = 'assets/driving_video.mp4'
    st.session_state['side_by_side_video_path'] = 'assets/src_image_driving_video_side_by_side.mp4'
    st.session_state['uploaded_src_image_file']=False
    # st.session_state[
    #     'side_by_side_with_generated_video_path'] = 'assets/src_image_driving_video_generated_video_side_by_side.mp4'
#side_by_side_with_generated_video_path
    st.session_state['predict_mode'] = 'relative'  # ['standard', 'relative', 'avd']
    st.session_state['find_best_frame'] = False  # when use the relative mode to animate a face, use 'find_best_frame=True' can get better quality result
    config_path = 'config/vox-256.yaml'
    checkpoint_path = 'checkpoints/vox.pth.tar'
    st.session_state['pixel'] = 512  # for vox, taichi and mgif, the resolution is 256*256
    print("start loading model")
    current_status_placeholder.write("start loading model")
    st.session_state['inpainting'], st.session_state['kp_detector'], st.session_state['dense_motion_network'], st.session_state['avd_network'] = load_checkpoints(config_path=config_path,
                                                                                  checkpoint_path=checkpoint_path,
                                                                                  device= st.session_state['device'] )
    print("finished loading model")
    current_status_placeholder.write("finished loading model")
    current_status_placeholder.write("copying default src image")
    try:
        shutil.copy2('assets/default_src_image.png',
                     st.session_state['source_image_path'])
        print("File copied successfully.")
    except FileNotFoundError:
        print("The source file was not found.")
    except PermissionError:
        print("Permission denied.")
    except Exception as e:
        print(f"Error occurred: {e}")

    try:
        current_status_placeholder.write("copying default driving video")
        shutil.copy2('assets/default_driving_video.mp4',
                     st.session_state['driving_video_path'])
        print("File copied successfully.")
    except FileNotFoundError:
        print("The source file was not found.")
    except PermissionError:
        print("Permission denied.")
    except Exception as e:
        print(f"Error occurred: {e}")
    current_status_placeholder.write("reading src image")
    read_src_image()
    # st.session_state['thumb_source_image'] = resize(st.session_state['source_image'], (250, 250))[..., :3]
    current_status_placeholder.write("reading driving video")


    read_driving_video(init_progress_bar)
    if os.path.exists('assets/default_src_image_driving_video_side_by_side.mp4'):
        print("deafult side_by_side_video already exists")
        try:
            current_status_placeholder.write("copying side by side video")
            shutil.copy2('assets/default_src_image_driving_video_side_by_side.mp4', st.session_state['side_by_side_video_path'])
            print("File copied successfully.")
        except FileNotFoundError:
            print("The source file was not found.")
        except PermissionError:
            print("Permission denied.")
        except Exception as e:
            print(f"Error occurred: {e}")
    else:
        current_status_placeholder.write("creating side by side video")

        st.session_state['side_by_side_video_path'] = create_image_video_side_by_side(st.session_state['source_image'], st.session_state['driving_video'], output_file=st.session_state['side_by_side_video_path'], fps=st.session_state['fps'], progress_bar=init_progress_bar)
    current_status_placeholder.write("")
st.video(st.session_state['side_by_side_video_path'])
col1, col2 = st.columns(2)

with col1:
    uploaded_src_image_file = st.file_uploader("Upload a source image... image must be square dimension", type=['jpg', 'jpeg', 'png'])
    st.markdown(f'<a href="https://ikmtechnology.github.io/ikmtechnology/Kyla2.png" target="_blank">Sample 1 download and then upload to above</a>', unsafe_allow_html=True)
    st.markdown(f'<a href="https://ikmtechnology.github.io/ikmtechnology/Aude.png" target="_blank">Sample 2 download and then upload to above</a>', unsafe_allow_html=True)
with col2:
    uploaded_driving_video_file = st.file_uploader(
        "Upload a driving video... video must be square dimension... 512x512 recommended", type=['mp4'])
    st.markdown(
        f'<a href="https://ikmtechnology.github.io/ikmtechnology/jenny.mp4" target="_blank">Sample 1 download and then upload to above</a>',
        unsafe_allow_html=True)
    st.markdown(
        f'<a href="https://ikmtechnology.github.io/ikmtechnology/anna.mp4" target="_blank">Sample 2 download and then upload to above</a>',
        unsafe_allow_html=True)

if uploaded_src_image_file is not None:
    if is_new_src_image_upload(uploaded_src_image_file):
        current_status_placeholder.write("checking uploaded source image")
        save_path = './uploaded_images'
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        # Open the file in the desired location with write-binary ('wb') mode
        with open(os.path.join(save_path, "temp_"+uploaded_src_image_file.name), "wb") as f:
            f.write(uploaded_src_image_file.getbuffer())  # Write the file to the specified location
            # st.success(f'Saved file temp_{uploaded_src_image_file.name} in {save_path}')

        image = imageio.imread(os.path.join(save_path, "temp_"+uploaded_src_image_file.name))
        height, width = image.shape[:2]
        # To see details
        #file_details = {"FileName": uploaded_src_image_file.name, "FileType": uploaded_src_image_file.type, "FileSize": uploaded_src_image_file.size}
        #st.write(file_details)

        # Save the file
        if width == height:
            current_status_placeholder.write("saving uploaded image")
            save_image_from_upload(uploaded_src_image_file)
            current_status_placeholder.write("reading uploaded image")
            read_src_image()
            current_status_placeholder.write("creating side by side video")
            st.session_state['side_by_side_video_path'] = create_image_video_side_by_side(st.session_state['source_image'],
                                                                                      st.session_state['driving_video'],
                                                                                      output_file=st.session_state[
                                                                                          'side_by_side_video_path'],
                                                                                      fps=st.session_state['fps'], progress_bar=init_progress_bar)
            print("uploaded_src_image_file Done! ")
            st.rerun();

        else:
            st.error("Error: Image width and height must be equal.")
    # if not st.session_state['uploaded_src_image_file']:
    #     st.rerun()
    #     st.session_state['uploaded_src_image_file'] = True
    # st.video(st.session_state['side_by_side_video_path'])


if uploaded_driving_video_file is not None:
    if is_new_driving_video_upload(uploaded_driving_video_file):

        # To see details
        # file_details = {"FileName": uploaded_driving_video_file.name, "FileType": uploaded_driving_video_file.type, "FileSize": uploaded_driving_video_file.size}
        # st.write(file_details)
        current_status_placeholder.write("checking uploaded video")
        save_path = './uploaded_videos'
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        # Open the file in the desired location with write-binary ('wb') mode
        with open(os.path.join(save_path, "temp_uploaded_video.mp4"), "wb") as f:
            f.write(uploaded_driving_video_file.getbuffer())  # Write the file to the specified location
            # st.success(f'Saved file "temp_uploaded_video.mp4" in {save_path}')

        reader = imageio.get_reader(os.path.join(save_path, "temp_uploaded_video.mp4"))


        video_width = reader.get_meta_data()['source_size'][0]
        video_height = reader.get_meta_data()['source_size'][1]
        # Check if dimensions are not equal
        if video_width != video_height:
            st.error("Error: Video width and height must be equal.")
        else:
            # Display dimensions
            # st.write(f"Width: {video_width}px")
            # st.write(f"Height: {video_height}px")
            current_status_placeholder.write("saving uploaded video")
            save_video_from_upload(uploaded_driving_video_file)
            current_status_placeholder.write("reading uploaded video")
            read_driving_video(init_progress_bar)
            current_status_placeholder.write("creating side by side video")
            st.session_state['side_by_side_video_path'] = create_image_video_side_by_side(st.session_state['source_image'],
                                                                                      st.session_state['driving_video'],
                                                                                      output_file=st.session_state[
                                                                                          'side_by_side_video_path'],
                                                                                      fps=st.session_state['fps'], progress_bar=init_progress_bar)
            st.rerun()
    # st.video(st.session_state['side_by_side_video_path'])
# x = st.slider('Select a value')
# st.write(x, 'squared is', x * x)
# Display the video

print("st.session_state['side_by_side_video_path'] =" +st.session_state['side_by_side_video_path'])

# Create a button and check if the button is clicked
inference_status_placeholder = st.empty()

if 'run_button' in st.session_state and st.session_state.run_button == True:
    st.session_state.running = True
else:
    st.session_state.running = False
create_animation_progress_bar = st.progress(0)
if st.button('Add motion to Image',disabled=st.session_state.running, key='run_button'):
    add_animation_to_image()
    st.session_state['video_generated'] = True
    st.rerun()
    # What to do after the button is clicked

if 'video_generated' in st.session_state:
    st.video(st.session_state['side_by_side_with_generated_video_path'])

    st.video(st.session_state['output_video_path'])
    del st.session_state['video_generated']