import gradio as gr
import tempfile
import json
from inference_color_temperature import load_models, run_inference, OmegaConf
import torch

# Initialize models once at startup
cfg = OmegaConf.load("configs/inference_genphoto/adv3_256_384_genphoto_relora_color_temperature.yaml")
pipeline, device = load_models(cfg)

def generate_video(base_scene, color_temperature_list):
    try:
        # Validate input
        if len(json.loads(color_temperature_list)) != 5:
            raise ValueError("Exactly 5 color_temperature values required")
            
        # Run inference
        video_path = run_inference(
            pipeline=pipeline,
            tokenizer=pipeline.tokenizer,
            text_encoder=pipeline.text_encoder,
            base_scene=base_scene,
            color_temperature_list=color_temperature_list,
            device=device
        )
        return video_path
    
    except Exception as e:
        raise gr.Error(f"Generation failed: {str(e)}")

# Example inputs
examples = [
    [
        "A beautiful blue sky with a mountain range in the background.",
        "[5455.0, 5155.0, 5555.0, 6555.0, 7555.0]"
    ],
    [
        "A red couch is situated in front of a window, which is filled with a variety of potted plants.",
        "[3500.0, 5500.0, 6500.0, 7500.0, 8500.0]"
    ]
]

with gr.Blocks(title="Color Temperature Effect Generator") as demo:
    gr.Markdown("# Dynamic Color Temperature Effect Generation")
    
    with gr.Row():
        with gr.Column():
            scene_input = gr.Textbox(
                label="Scene Description",
                placeholder="Describe the scene you want to generate..."
            )
            color_temperature_input = gr.Textbox(
                label="Color Temperature Values",
                placeholder="Enter 5 comma-separated values from 2000-10000 (e.g., [3001.3, 4000.2, 4400.34, 5488.23, 8888.82])"
            )
            submit_btn = gr.Button("Generate Video", variant="primary")
            
        with gr.Column():
            video_output = gr.Video(label="Generated Video")
            error_output = gr.Textbox(label="Error Messages", visible=False)
    
    gr.Examples(
        examples=examples,
        inputs=[scene_input, color_temperature_input],
        outputs=[video_output],
        fn=generate_video,
        cache_examples=True
    )
    
    submit_btn.click(
        fn=generate_video,
        inputs=[scene_input, color_temperature_input],
        outputs=[video_output],
    )

if __name__ == "__main__":
    demo.launch(share=True)