import gradio as gr
import numpy as np
import os
import subprocess
from datetime import datetime
os.system('git clone https://huggingface.co/camenduru/GaussianDreamer')
os.system('pip install ./gaussiansplatting/submodules/diff-gaussian-rasterization')
os.system('pip install ./GaussianDreamer/nerfacc-0.5.3-cp310-cp310-linux_x86_64.whl')
os.system('pip install  ./gaussiansplatting/submodules/simple-knn')
os.system('pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch')
os.system('git clone https://github.com/openai/shap-e.git')
os.system('pip install -e ./shap-e')
os.system('mv ./GaussianDreamer/shapE_finetuned_with_330kdata.pth ./load/shapE_finetuned_with_330kdata.pth')
example_inputs = [[
        "A fox."
    ], [
        "fries and a hamburger."
    ], [
        "Viking axe, fantasy, weapon, blender, 8k, HD."
    ], [
        "ferrari convertible, trending on artstation, ultra realistic, 4k, HD"
    ], [
        "flamethrower, with fire, scifi, cyberpunk, photorealistic, 8K, HD"
]]
example_outputs_1 = [
    gr.Video(value=os.path.join(os.path.dirname(__file__), 'example/a_fox.mp4'), autoplay=True),
    gr.Video(value=os.path.join(os.path.dirname(__file__), 'example/fries_and_a_hamburger.mp4'), autoplay=True),
    gr.Video(value=os.path.join(os.path.dirname(__file__), 'example/Viking_axe,_fantasy,_weapon,_blender,_8k,_HD.mp4'), autoplay=True),
    gr.Video(value=os.path.join(os.path.dirname(__file__), 'example/ferrari_convertible,_trending_on_artstation,_ultra_realistic,_4k,_HD.mp4'), autoplay=True),
    gr.Video(value=os.path.join(os.path.dirname(__file__), 'example/flamethrower,_with_fire,_scifi,_cyberpunk,_photorealistic,_8K,_HD.mp4'), autoplay=True)
]


subprocess.run([
                f'python shape.py'],
            shell=True)

def main(prompt, iteration,CFG, seed):
    if [prompt] in example_inputs:
        return example_outputs_1[example_inputs.index([prompt])]
    seed = int(seed)
    iteration = int(iteration)
    print('==> User Prompt:', prompt)
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    print('==> User Shell:', f'python launch.py --config configs/gaussiandreamer-sd.yaml --train --gpu 0 system.prompt_processor.prompt="{prompt}" seed={seed} system.guidance.guidance_scale={CFG} trainer.max_steps={iteration} use_timestamp=False timestamp="{timestamp}" ',)
    
    subprocess.run([
                    f'python launch.py --config configs/gaussiandreamer-sd.yaml --train --gpu 0 system.prompt_processor.prompt="{prompt}" seed={seed} system.guidance.guidance_scale={CFG} trainer.max_steps={iteration} use_timestamp=False timestamp="{timestamp}" '],
                shell=True)
    path= os.path.join("./outputs/gaussiandreamer-sd",f'{prompt.replace(" ","_")}{timestamp}',f"save/it{iteration}-test.mp4")
    print('==> Save path:', path)
    return gr.Video(value=path, autoplay=True)

with gr.Blocks() as demo:

    gr.Markdown("# <center>GaussianDreamer: Fast Generation from Text to 3D Gaussians by Bridging 2D and 3D Diffusion Models</center>")
    gr.Markdown("This live demo allows you to generate high-quality 3D content using text prompts. The outputs are 360° rendered 3d video.<br> \
                It is based on Stable Diffusion 2.1-base. Please check out our <strong><a href=https://taoranyi.com/gaussiandreamer/>Project Page</a> / <a href=https://arxiv.org/abs/2310.08529>Paper</a> / <a href=https://github.com/hustvl/GaussianDreamer>Code</a></strong> if you want to learn more about our method!<br> \
                Note that this demo is running on T4, the running time might be longer than the reported 15 minutes (1200 iterations) on RTx 3090.<br> \
                &copy; This Gradio space is developed by <a href=https://taoranyi.com/>Taoran Yi</a>.")
    gr.Interface(fn=main, inputs=[gr.Textbox(lines=2, value="fries and a hamburger.", label="Your prompt"),
            gr.Slider(0, 2000, value=1200, label="Number of iteration"),
            gr.Slider(80, 200, value=100, label="CFG"),
            gr.Number(value=0, label="Seed")], 
        outputs=["playable_video"],
        examples=example_inputs,
        cache_examples=True,
        concurrency_limit=1)
demo.launch()