import os
import io
import base64
import gradio as gr
from stability_sdk.api import Context
from stability_sdk.animation import AnimationArgs, Animator
from dotenv import load_dotenv

load_dotenv(".env") 

STABILITY_HOST = "grpc.stability.ai:443" 
#STABILITY_KEY = os.getenv("STABILITY_KEY")


def anim(f_promt, s_promt, stability_key):
    # Connect to Stability API
    context = Context(STABILITY_HOST, stability_key)
    # Test the connection
    context.get_user_info()
    print("Connection successfuly!")

    # Configure the animation
    args = AnimationArgs()
    args.interpolate_prompts = True
    args.locked_seed = True
    args.max_frames = 5
    args.seed = 42
    args.strength_curve = "0:(0)"
    args.diffusion_cadence_curve = "0:(4)"
    args.cadence_interp = "film"

    animation_prompts = {
        0: f_promt,
        2: s_promt,
    }
    negative_prompt = ""

    # Create Animator object to orchestrate the rendering
    animator = Animator(
        api_context=context,
        animation_prompts=animation_prompts,
        negative_prompt=negative_prompt,
        args=args
    )
    # Get root folder of the repository
    repo_dir = os.getenv('HF_COMET_ROOT_PATH', '/')

    # Define output folder path
    image_path = "/tmp/frames/"
    output_dir = os.path.join(image_path, "output")

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    # Render each frame of animation
    images = []
    for idx, frame in enumerate(animator.render()):
        file_path = os.path.join(output_dir, f"frame_{idx:05d}.png") 
        frame.save(file_path)
        print("Created frame at:"+file_path)
        images.append(file_path)
    
    return images
      
with gr.Blocks() as demo:
    gr.Markdown("Stability Animation")
    f_promt = gr.Textbox(label="First Prompt", value="a photo of a cute cat")
    s_promt = gr.Textbox(label="Second Prompt", value="a photo of a cute dog")
    stability_key = gr.Textbox(label="Stability Key", value="")
    
    outimg = gr.File(label="Generated Files")
    btn = gr.Button('Anim')

    btn.click(fn=anim, inputs=[f_promt, s_promt, stability_key], outputs=[outimg],api_name="AnimAPI")

demo.launch()