app.py

import os
import torch
import subprocess
import spaces
import gradio as gr
from PIL import Image
import numpy as np
from RealESRGAN import RealESRGAN

os.environ['CUDA_HOME'] = "/usr/local/cuda"
os.environ['PATH'] = '/usr/local/lib/python3.10/site-packages/nvidia/cuda_nvrtc/lib:' + os.environ.get('PATH', '')
os.environ['LD_LIBRARY_PATH'] = '/usr/local/lib/python3.10/site-packages/nvidia/cuda_nvrtc/lib:' + os.environ.get('LD_LIBRARY_PATH', '')

print(os.environ['PATH'])
print(os.environ['LD_LIBRARY_PATH'])

subprocess.run(["git", "clone", "https://github.com/sniklaus/3d-ken-burns.git"])
subprocess.run(["mv", "./network-disparity.pytorch", "./3d-ken-burns/models/disparity-estimation.pytorch"])
subprocess.run(["mv", "./network-refinement.pytorch", "./3d-ken-burns/models/disparity-refinement.pytorch"])
subprocess.run(["mv", "./network-inpainting.pytorch", "./3d-ken-burns/models/pointcloud-inpainting.pytorch"])

# with open ('./3d-ken-burns/autozoom.py','r+') as az:
#     script = az.read()
#    script = script.replace("'fltShift': 100.0","'fltShift': 0.0")
#    script = script.replace("'fltZoom': 1.25","'fltZoom': 1.35")
#    script = script.replace("'fltSteps': numpy.linspace(0.0, 1.0, 75).tolist(),","'fltSteps': numpy.linspace(0.0, 1.0, 250).tolist(),")
#    az.write(script)

@spaces.GPU
def upscale(image):
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  model = RealESRGAN(device, scale=4)
  model.load_weights('weights/RealESRGAN_x4.pth', download=True)
  sr_image = model.predict(image)
  sr_image.save('sr_image.png')
  sr_image.resize((1280,720))
  return sr_image

@spaces.GPU
def generate_video(image):
    subprocess.run(["nvidia-smi"])
    subprocess.run(["find", "/", "-name", "cuda"])
    image = Image.fromarray(image)
    image.save("/tmp/img.png")
    os.chdir("./3d-ken-burns")
    subprocess.run(["python", "autozoom.py", "--in", "/tmp/img.png", "--out", "/tmp/autozoom.mp4"]) 
    os.chdir("..")
    return "/temp/autozoom.mp4"

with gr.Blocks() as demo:
    gr.Markdown("""# Ken Burns Video Generator!""")
    image = gr.Image()
    submit_image = gr.Button(value="Generate Video")
    video = gr.Video()
    submit_image.click(fn=generate_video,inputs=image,outputs=video)
    
demo.launch()