import gradio as gr
import json
from deforum_engine import DeforumRunner

runner = DeforumRunner(device="cpu")

def process(prompts_json, neg, frames, w, h, z, a, tx, ty, stre, noi, fps, steps, 
            cadence, color, border, init_img, model, lora, sched):
    try:
        p_dict = json.loads(prompts_json.replace("'", '"'))
        prompts = {int(k): v for k, v in p_dict.items()}
    except: return None, None, None
    
    yield from runner.render(
        prompts, neg, int(frames), 256, 256, 
        z, a, tx, ty, stre, noi, int(fps), int(steps),
        int(cadence), color, border, init_img,
        model, lora, sched
    )

def stop_gen():
    runner.stop()
    return gr.update(value="Stopping...")

css = "#col-container {max_width: 950px; margin-left: auto; margin-right: auto;}"
with gr.Blocks() as demo:
    gr.Markdown("# 🌀 Authentic Deforum (CPU/LCM)")
    
    with gr.Row(elem_id="col-container"):
        with gr.Column(scale=1):
            with gr.Accordion("⚙️ Model & Init", open=False):
                model = gr.Dropdown(label="Model", value="AlekseyCalvin/acs_model",
                                    choices=["AlekseyCalvin/acs_model", "runwayml/stable-diffusion-v1-5"])
                lora = gr.Dropdown(label="LoRA", value="latent-consistency/lcm-lora-sdv1-5",
                                   choices=["latent-consistency/lcm-lora-sdv1-5", "None"])
                sched = gr.Dropdown(label="Scheduler", value="LCM", choices=["LCM", "Euler A"])
                init_img = gr.Image(label="Init Image (Optional)", type="pil", height=150)

            prompts = gr.Code(label="Prompts (JSON)", language="json",
                              value='{\n "0": "a mystical forest, highly detailed, 8k",\n "40": "a forest fire, intricate details"\n}')
            neg = gr.Textbox(label="Negative", value="blur, low quality, watermark")
            
            with gr.Row():
                frames = gr.Number(label="Frames", value=120, step=10)
                steps = gr.Slider(2, 12, value=4, step=1, label="Steps (LCM: 4-8)")
                cadence = gr.Slider(1, 4, value=1, step=1, label="Cadence (Warp-only frames)")
                fps = gr.Number(label="FPS", value=12)

            with gr.Accordion("🎬 Motion & Coherence", open=True):
                with gr.Row():
                    color = gr.Dropdown(label="Color Match", value="LAB", choices=["None", "LAB", "HSV"])
                    border = gr.Dropdown(label="Border Mode", value="Reflect", choices=["Reflect", "Replicate", "Wrap", "Black"])
                
                zoom = gr.Textbox(label="Zoom", value="0:(1.01)")
                angle = gr.Textbox(label="Angle", value="0:(0.5*sin(t/10))")
                tx = gr.Textbox(label="TX", value="0:(1*cos(t/15))")
                ty = gr.Textbox(label="TY", value="0:(0)")
                stre = gr.Textbox(label="Strength (Decay)", value="0:(0.65)")
                noi = gr.Textbox(label="Noise (Grain)", value="0:(0.03)")
                
            with gr.Row():
                btn = gr.Button("Generate", variant="primary", scale=3)
                stop = gr.Button("Stop", variant="stop", scale=1)
            
        with gr.Column(scale=1):
            img = gr.Image(label="Live Preview")
            vid = gr.Video(label="Final Video")
            files = gr.File(label="Frames ZIP")
            stop_msg = gr.Markdown("", visible=True)

    inputs = [prompts, neg, frames, gr.State(256), gr.State(256), zoom, angle, tx, ty, stre, noi, fps, steps, cadence, color, border, init_img, model, lora, sched]
    btn.click(process, inputs=inputs, outputs=[img, vid, files])
    stop.click(stop_gen, outputs=stop_msg)

if __name__ == "__main__":
    demo.queue().launch(css=css, theme=gr.themes.Glass())