import spaces
import gradio as gr
import time
import torch
import os
import json
import subprocess

from diffusers import (
    DDPMScheduler,
    AutoPipelineForText2Image,
    AutoencoderKL,
)

def runcmd(cmd, verbose = False, *args, **kwargs):

    process = subprocess.Popen(
        cmd,
        stdout = subprocess.PIPE,
        stderr = subprocess.PIPE,
        text = True,
        shell = True
    )
    std_out, std_err = process.communicate()
    if verbose:
        print(std_out.strip(), std_err)
    pass

# os.system("python3 -m pip --no-cache-dir install --pre nexfort -f https://github.com/siliconflow/nexfort_releases/releases/expanded_assets/torch2.4.1_cu121")
# os.system("git clone https://github.com/siliconflow/onediff.git")
# os.system("cd onediff && python3 -m pip install .")
# sys.path.append("/home/user/app/onediff/src")
# os.system("cd onediff/onediff_diffusers_extensions && python3 -m pip install .")
# sys.path.append("/home/user/app/onediff/onediff_diffusers_extensions/src")

# os.system("pip show nexfort")
# os.system("pip show onediff")
# os.system("pip show onediffx")

# from onediffx import compile_pipe, save_pipe

# def nexfort_compile(torch_module: torch.nn.Module):
#     options = json.loads('{"mode": "max-optimize:max-autotune:low-precision", "memory_format": "channels_last", "dynamic": true}')
#     return compile_pipe(torch_module, backend="nexfort", options=options, fuse_qkv_projections=True)

os.system("apt install -y nvidia-cuda-toolkit")
print(os.environ.get('CUDA_PATH'))
print(os.environ.get('CUDA_HOME'))
os.system("pip show torch")
os.system("nvcc --version")
os.system("which nvcc")

# os.system("CUDA_HOME=/usr/local/cuda-12.1 python3 -m pip install stable-fast") 

import xformers
import triton
from sfast.compilers.diffusion_pipeline_compiler import (compile, CompilationConfig)

BASE_MODEL = "stabilityai/sdxl-turbo"
device = "cuda"

vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix", 
    torch_dtype=torch.float16,
)
base_pipe = AutoPipelineForText2Image.from_pretrained(
    BASE_MODEL, 
    vae=vae,
    torch_dtype=torch.float16, 
    variant="fp16", 
    use_safetensors=True,
)
base_pipe.to(device)

# helper = DeepCacheSDHelper(pipe=base_pipe)
# helper.set_params(cache_interval=3, cache_branch_id=0)
# helper.enable()

# base_pipe = nexfort_compile(base_pipe)

ccnf = CompilationConfig.Default()

ccnf.enable_xformers = True
ccnf.enable_triton = True
ccnf.enable_cuda_graph = True

base_pipe = compile(base_pipe, ccnf)

from gfpgan.utils import GFPGANer
from basicsr.archs.srvgg_arch import SRVGGNetCompact
from realesrgan.utils import RealESRGANer

if not os.path.exists('GFPGANv1.4.pth'):
    runcmd("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
if not os.path.exists('realesr-general-x4v3.pth'):
    runcmd("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")

model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
model_path = 'realesr-general-x4v3.pth'
half = True if torch.cuda.is_available() else False
upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)

face_enhancer = GFPGANer(model_path='GFPGANv1.4.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)

def create_demo() -> gr.Blocks:

    @spaces.GPU(duration=30)
    def text_to_image(
        prompt:str,
        steps:int,
    ):
        run_task_time = 0
        time_cost_str = ''
        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
        generated_image = base_pipe(
            prompt=prompt,
            num_inference_steps=steps,
        ).images[0]
        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
        return generated_image, time_cost_str

    def get_time_cost(run_task_time, time_cost_str):
        now_time = int(time.time()*1000)
        if run_task_time == 0:
            time_cost_str = 'start'
        else:
            if time_cost_str != '': 
                time_cost_str += f'-->'
            time_cost_str += f'{now_time - run_task_time}'
        run_task_time = now_time
        return run_task_time, time_cost_str

    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(label="Prompt", placeholder="Write a prompt here", lines=2, value="A beautiful sunset over the city")
            with gr.Column():
                steps = gr.Slider(minimum=1, maximum=100, value=5, step=1, label="Num Steps")
                g_btn = gr.Button("Generate")
                
        with gr.Row():
            with gr.Column():
                generated_image = gr.Image(label="Generated Image", type="pil", interactive=False)
            with gr.Column():
                time_cost = gr.Textbox(label="Time Cost", lines=1, interactive=False)
        
        g_btn.click(
            fn=text_to_image,
            inputs=[prompt, steps],
            outputs=[generated_image, time_cost],
        )

    return demo