import gradio as gr
from gradio_client import Client
import os
import shutil
import tempfile
import zipfile

def generate_frame(client, prompt, frame_number, total_frames):
    adjusted_prompt = f"{prompt} - frame {frame_number} of {total_frames}"
    
    result = client.predict(
        prompt=adjusted_prompt,
        is_negative="(deformed, distorted, disfigured), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, misspellings, typos",
        steps=35,
        cfg_scale=7,
        sampler="DPM++ 2M Karras",
        seed=-1,
        strength=0.7,
        huggingface_api_key="",
        use_dev=True,
        enhance_prompt_style="Cartoon",
        enhance_prompt_option=True,
        nemo_enhance_prompt_style="Cartoon",
        use_mistral_nemo=True,
        api_name="/query"
    )
    return result

def generate_video_frames(prompt, num_frames, progress=gr.Progress()):
    client = Client("K00B404/FLUX.1-Dev-Serverless-darn-enhanced-prompt")
    
    with tempfile.TemporaryDirectory() as temp_dir:
        frame_paths = []
        for i in progress.tqdm(range(num_frames)):
            tmp_image_path, used_seed, enhanced_prompt = generate_frame(client, prompt, i+1, num_frames)
            
            destination_path = os.path.join(temp_dir, f"frame_{i:03d}.png")
            shutil.copy(tmp_image_path, destination_path)
            frame_paths.append(destination_path)
            
            yield f"Generated frame {i+1}/{num_frames}\nSeed: {used_seed}\nEnhanced prompt: {enhanced_prompt}"
        
        zip_path = os.path.join(temp_dir, "video_frames.zip")
        with zipfile.ZipFile(zip_path, 'w') as zipf:
            for file in frame_paths:
                zipf.write(file, os.path.basename(file))
        
        yield f"Generated {num_frames} frames. Downloading ZIP file..."
        return zip_path

def launch_interface(prompt, num_frames):
    output = generate_video_frames(prompt, num_frames)
    return output

# Define the Gradio interface
iface = gr.Interface(
    fn=launch_interface,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
        gr.Slider(minimum=1, maximum=50, step=1, label="Number of Frames", value=10)
    ],
    outputs=[
        gr.Textbox(label="Progress", lines=5),
        gr.File(label="Download Frames")
    ],
    title="Video Frame Generator",
    description="Generate a series of video frames using the FLUX.1 model.",
    allow_flagging="never"
)

# Launch the interface
iface.launch()