from pathlib import Path
import gradio as gr
import numpy as np
from zipfile import ZipFile

def sepia(input_img, num_copies):
    sepia_filter = np.array([
        [0.393, 0.769, 0.189], 
        [0.349, 0.686, 0.168], 
        [0.272, 0.534, 0.131]
    ])
    input_img = input_img / 255.0
    sepia_img = np.dot(input_img[..., :3], sepia_filter.T)
    sepia_img = np.clip(sepia_img, 0, 1) * 255
    sepia_imgs = [(sepia_img).astype(np.uint8) for _ in range(num_copies)]
    return sepia_imgs

def zip_sepia_images(sepia_imgs):
    temp_dir = Path(tempfile.mkdtemp())  # Use Path for easier file handling
    zip_path = temp_dir / "sepia_images.zip"
    with ZipFile(zip_path, 'w') as zipf:
        for i, img in enumerate(sepia_imgs):
            img_path = temp_dir / f"sepia_image_{i}.png"
            gr.utils.save_image(str(img_path), img)  # Convert Path object to string for compatibility
            zipf.write(img_path, img_path.name)
    return str(zip_path)  # Convert Path object to string for return

with gr.Blocks() as demo:
    with gr.Row():
        input_img = gr.Image()
        num_copies = gr.Number(label="Number of Copies", value=1)
        gallery = gr.Gallery(label="Sepia Images")
        download_btn = gr.File(label="Download ZIP", visible=True)
        
    input_img.change(fn=lambda x: x, inputs=input_img, outputs=gallery)
    num_copies.change(fn=lambda x: x, inputs=num_copies, outputs=gallery)
    
    generate_btn = gr.Button("Generate Sepia Images")
    generate_btn.click(fn=sepia, inputs=[input_img, num_copies], outputs=gallery)
    generate_btn.click(fn=zip_sepia_images, inputs=gallery, outputs=download_btn)

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