import subprocess
import os
import gradio as gr

from PIL import Image

def add_new_image(prompt,image):
    image = Image.fromarray(image)
    height = image.height
    width = image.width
    lr = 0.5
    ip_img = "${.condition_images}"
    image.save("/content/img.png")
    img_settings = f"""- condition_images: /content/img.png
  eye_blinks_factor: 1.8
  height: {height}
  img_length_ratio: {lr}
  ipadapter_image: {ip_img}
  name: image
  prompt: {prompt}
  refer_image: {ip_img}
  video_path: null
  width: {width}"""
    print(img_settings)
    with open ("/content/MuseV/MiseV/configs/tasks/example.yaml","r+") as configs:
        configs.write(img_settings)
        configs.truncate()
        configs.seek(0)

def add_new_video(video):
    print(video)
    lr = 1.0
    ip_img = "${.condition_images}"
    img_settings = f"""- name: "dance2"
  prompt: "(best quality), ((masterpiece)), (highres), illustration, original, extremely detailed wallpaper"
  video_path:  ./MuseV/data/source_video/video1_girl_poseseq.mp4 
  condition_images: ./MuseV/data/images/cyber_girl.png
  refer_image: {ip_img}
  ipadapter_image: {ip_img}
  height: 960
  width: 512
  img_length_ratio: 1.0
  video_is_middle: True """
    print(img_settings)
    with open ("/content/MuseV/configs/tasks/example.yaml","r+") as configs:
        configs.write(img_settings)
        configs.truncate()
        configs.seek(0)     

def run(frames,fps):
    #subprocess.run(["python", "./MuseV/text2video.py", "--sd_model_name", "majicmixRealv6Fp16", "--unet_model_name", "musev", "-test_data_path", "./MuseV/configs/tasks/example.yaml", "--n_batch", "1", "--target_datas", "image", "--vae_model_path", "./MuseV/checkpoints/vae/sd-vae-ft-mse", "--motion_speed", "5.0", "--time_size", "12", "--fps", "12"])
    subprocess.run(["python", "./MuseV/text2video.py", "--sd_model_name", "majicmixRealv6Fp16", "--unet_model_name", "musev_referencenet", "--referencenet_model_name", "musev_referencenet", "--ip_adapter_model_name", "musev_referencenet", "-test_data_path", "./MuseV/configs/tasks/example.yaml", "--output_dir", "./MuseV",  "--n_batch", "1",  "--target_datas", "image",  "--vision_clip_extractor_class_name", "ImageClipVisionFeatureExtractor", "--vision_clip_model_path", "./MuseV/checkpoints/IP-Adapter/models/image_encoder", "--motion_speed", "5.0", "--vae_model_path", "./MuseV/checkpoints/vae/sd-vae-ft-mse", "--time_size", frames, "--fps", fps])
    return "./output.mp4"

def run_video():
    subprocess.run(["python", "./MuseV/video2video.py", "--sd_model_name", "fantasticmix_v10",  "--unet_model_name", "musev", "-test_data_path", "./MuseV/configs/tasks/example.yaml", "--output_dir", "./output",  "--n_batch", "1", "--controlnet_name", "dwpose_body_hand",  "--which2video", "video_middle",  "--target_datas",  "dance1", "--fps", "12", "--time_size", "12"])
    return "./output.mp4"
    
with gr.Blocks() as demo:
  title = gr.Markdown("""# MuseV Image2Vid & Vid2Vid """)
  subtitle1 = gr.Markdown("""Image2Vid""")
  image = gr.Image()
  button1 = gr.Button(value="Save Image")
  frames = gr.Number(value=12)
  fps = gr.Number(value=12)
  prompt = gr.Text(value="(masterpiece, best quality, highres:1),(1person, solo:1),(eye blinks:1.8),(head wave:1.3)")
  button1.click(fn=add_new_image,inputs=[prompt,image])
  button2 = gr.Button(value="Generate Img2Vid")
  video = gr.Video()
  button2.click(fn=run,inputs=[frames,fps],outputs=video)
  subtitle2 = gr.Markdown("""Vid2Vid""")
  video_in = gr.Video()
  button3 = gr.Button(value="Save Video")
  button3.click(fn=add_new_video,inputs=[video_in])
  button4 = gr.Button(value="Generate Vid2Vid")
  video_out = gr.Video()
  button4.click(fn=run_video,outputs=video_out)  

demo.launch(share=True)