import gradio as gr
import os
import requests
from huggingface_hub import cached_download, hf_hub_url


MORE = """ ## TRY Other Models
        ### JARVIS: Your VOICE Assistant -> https://huggingface.co/spaces/KingNish/JARVIS
        ### Instant Image: 4k images in 5 Second -> https://huggingface.co/spaces/KingNish/Instant-Image
        """

# Function
def generate_image(prompt, base="Anime", motion="", step=8, progress=gr.Progress()):
    API_URL = "https://api-inference.huggingface.co/models/KingNish/Instant-Video"
    headers = {"Authorization": f"Bearer {os.environ.get('HF_API_TOKEN')}"}  # Замените на ваш токен Hugging Face API

    payload = {
        "inputs": prompt,
        "options": {
            "base": base,
            "motion": motion,
            "step": step
        }
    }

    response = requests.post(API_URL, headers=headers, json=payload)
    response.raise_for_status()

    video_url = response.json()[0]['video']
    video_path = cached_download(video_url)
    return video_path
        
# Gradio Interface
with gr.Blocks(css="style.css") as demo:
    
    with gr.Group():
        with gr.Row():
            prompt = gr.Textbox(
                label='Описание'
            )
        with gr.Row():
            select_base = gr.Dropdown(
                label='Стиль',
                choices=[
                    "Cartoon", 
                    "Realistic",
                    "3d",
                    "Anime",
                ],
                value="Anime",
                interactive=True
            )
            select_motion = gr.Dropdown(
                label='Движение',
                choices=[
                    ("Default", ""),
                    ("Zoom in", "guoyww/animatediff-motion-lora-zoom-in"),
                    ("Zoom out", "guoyww/animatediff-motion-lora-zoom-out"),
                    ("Tilt up", "guoyww/animatediff-motion-lora-tilt-up"),
                    ("Tilt down", "guoyww/animatediff-motion-lora-tilt-down"),
                    ("Pan left", "guoyww/animatediff-motion-lora-pan-left"),
                    ("Pan right", "guoyww/animatediff-motion-lora-pan-right"),
                    ("Roll left", "guoyww/animatediff-motion-lora-rolling-anticlockwise"),
                    ("Roll right", "guoyww/animatediff-motion-lora-rolling-clockwise"),
                ],
                value="",
                interactive=True
            )
            select_step = gr.Dropdown(
                label='Шаги вывода',
                choices=[
                    ('1-Step', 1), 
                    ('2-Step', 2),
                    ('4-Step', 4),
                    ('8-Step', 8),
                ],
                value=4,
                interactive=True
            )
            submit = gr.Button(
                scale=1,
                variant='primary'
            )
    video = gr.Video(
        label='Сгенерированое видео',
        autoplay=True,
        height=512,
        width=512,
        elem_id="video_output"
    )

    prompt.submit(
        fn=generate_image,
        inputs=[prompt, select_base, select_motion, select_step],
        outputs=video,
    )
    submit.click(
        fn=generate_image,
        inputs=[prompt, select_base, select_motion, select_step],
        outputs=video,
    )

demo.queue().launch()