from diffusers import DiffusionPipeline
import gradio as gr
import torch
import cv2
import os

MY_SECRET_TOKEN=os.environ.get('HF_TOKEN_SD')

device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    use_auth_token=MY_SECRET_TOKEN,
    #revision='fp16',
    #torch_dtype=torch.float16,
    safety_checker=None,  # Very important for videos...lots of false positives while interpolating
    custom_pipeline="interpolate_stable_diffusion",
).to(device)
pipe.enable_attention_slicing()

def run(prompt1, seed1, prompt2, seed2, prompt3, seed3):
    
    frame_filepaths = pipe.walk(
        prompts=[prompt1, prompt2, prompt3],
        seeds=[seed1, seed2, seed3],
        num_interpolation_steps=16,
        output_dir='./dreams',
        batch_size=4,
        height=512,
        width=512,
        guidance_scale=8.5,
        num_inference_steps=50,
    )
    print(frame_filepaths)
    
    frame = cv2.imread(frame_filepaths[0])
    height, width, layers = frame.shape
    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
    video = cv2.VideoWriter("out.mp4", fourcc, 24, (width,height))
    for image in frame_filepaths:
        #print(image)
        video.write(cv2.imread(image))
    
    video.release()
    cv2.destroyAllWindows()
    #print(video)
    
    return "out.mp4", frame_filepaths

with gr.Blocks() as demo:
  with gr.Column():
    gr.HTML('''
        <h1 style='font-size: 2em;text-align:center;font-weigh:900;'>
            Stable Diffusion Interpolation • Community pipeline
        </h1>
        <p style='text-align: center;'><br />
            This community pipeline returns a list of images saved under the folder as defined in output_dir. <br />
            You can use these images to create videos of stable diffusion.
        </p>
        
        <p style='text-align: center;'>
            This demo can be run on a GPU of at least 8GB VRAM and should take approximately 5 minutes.<br />
            —         
        </p>

    ''')
    with gr.Row():
      with gr.Column():
        with gr.Column():
          with gr.Row():
            intpol_prompt_1 = gr.Textbox(lines=1, label="prompt 1")
            seed1 = gr.Slider(label = "Seed 1", minimum = 0, maximum = 2147483647, step = 1, randomize = True)
          with gr.Row():
            intpol_prompt_2 = gr.Textbox(lines=1, label="prompt 2")
            seed2 = gr.Slider(label = "Seed 2", minimum = 0, maximum = 2147483647, step = 1, randomize = True)
          with gr.Row():
            intpol_prompt_3 = gr.Textbox(lines=1, label="prompt 3")
            seed3 = gr.Slider(label = "Seed 3", minimum = 0, maximum = 2147483647, step = 1, randomize = True)
          intpol_run = gr.Button("Run Interpolation")
      
      with gr.Column():
        video_output = gr.Video(label="Generated video", show_label=True)
        gallery_output = gr.Gallery(label="Generated images", show_label=False).style(grid=2, height="auto")
    
    intpol_run.click(run, inputs=[intpol_prompt_1, seed1, intpol_prompt_2, seed2, intpol_prompt_3, seed3], outputs=[video_output, gallery_output])
    
demo.launch()