import gradio as gr
import numpy as np
import os
import subprocess
from datetime import datetime

os.system('pip install ./gaussiansplatting/submodules/diff-gaussian-rasterization')
os.system('pip install ./gaussiansplatting/submodules/simple-knn')

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)
]



def main(prompt, CFG, seed):
    if [prompt] in example_inputs:
        return example_outputs_1[example_inputs.index([prompt])]
    seed = int(seed)
    print('==> User Prompt:', prompt)
    timestamp = datetime.now().strftime("@%Y%m%d-%H%M%S")
    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} use_timestamp=False timestamp="{timestamp}" '],
                shell=True)
    path= os.path.join("./outputs/gaussiandreamer-sd",f'{prompt.replace(" ","_")}{timestamp}',"save/it1200-test.mp4")
    print('==> Save path:', path)
    return gr.Video(value=path, autoplay=True)

with gr.Blocks() as demo:
    gr.Markdown("# <center>LucidDreamer: Towards High-Fidelity Text-to-3D Generation via Interval Score Matching</center>")
    gr.Markdown("This live demo allows you to generate high-quality 3D content using text prompts. The outputs are 360° rendered 3d gaussian video and training progress visualization.<br> \
                It is based on Stable Diffusion 2.1. Please check out our <strong><a href=https://github.com/EnVision-Research/LucidDreamer>Project Page</a> / <a href=https://arxiv.org/abs/2311.11284>Paper</a> / <a href=https://github.com/EnVision-Research/LucidDreamer>Code</a></strong> if you want to learn more about our method!<br> \
                Note that this demo is running on A10G, the running time might be longer than the reported 35 minutes (5000 iterations) on A100.<br> \
                &copy; This Gradio space was developed by Haodong LI.")
    gr.Interface(fn=main, inputs=[gr.Textbox(lines=2, value="A portrait of IRONMAN, white hair, head, photorealistic, 8K, HDR.", label="Your prompt"),
            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()