import gradio as gr
import numpy as np
from PIL import Image
import io
from zipfile import ZipFile

def sepia_and_zip(input_img_path, num_copies):
    input_img = Image.open(input_img_path)
    input_img_np = np.array(input_img)
    
    sepia_filter = np.array([
        [0.393, 0.769, 0.189],
        [0.349, 0.686, 0.168],
        [0.272, 0.534, 0.131]
    ])
    input_img_np = input_img_np / 255.0
    sepia_img_np = np.dot(input_img_np[..., :3], sepia_filter.T)
    sepia_img_np = np.clip(sepia_img_np, 0, 1) * 255
    sepia_imgs = [(sepia_img_np).astype(np.uint8) for _ in range(num_copies)]
    
    # Convert numpy images to PIL images for gallery
    pil_images = [Image.fromarray(img) for img in sepia_imgs]
    
    # Prepare ZIP file
    zip_bytes_io = io.BytesIO()
    with ZipFile(zip_bytes_io, 'w') as zip_file:
        for i, img in enumerate(pil_images):
            img_byte_arr = io.BytesIO()
            img.save(img_byte_arr, format='PNG')
            img_byte_arr = img_byte_arr.getvalue()
            zip_file.writestr(f"sepia_image_{i}.png", img_byte_arr)
    zip_bytes_io.seek(0)
    
    # Return a gallery of images and a ZIP file for download
    return pil_images, zip_bytes_io

title = "Sepia Tone Generator with Downloadable ZIP"
description = "Upload an image and select the number of sepia-toned copies you want. View the images in a gallery and download them all as a ZIP file."
examples = []

iface = gr.Interface(fn=sepia_and_zip,
                     inputs=[gr.Image(type="filepath", label="Input Image"),
                             gr.Number(label="Number of Copies")],
                     outputs=[gr.Gallery(label="Sepia Images"),
                              gr.File(label="Download ZIP")],
                     title=title,
                     description=description,
                     examples=examples,
                     allow_flagging="never",
                     cache_examples=False)

iface.launch(debug=True, enable_queue=False)