import gradio as gr
from PIL import Image
import tempfile
import os
import zipfile

def upscale_images(files):

    temp_dir = tempfile.mkdtemp()

    output_files = []

    for file in files:

        img = Image.open(file)

        width, height = img.size

        img = img.resize(
            (width * 2, height * 2),
            Image.LANCZOS
        )

        output_path = os.path.join(
            temp_dir,
            f"upscaled_{os.path.basename(file)}"
        )

        img.save(output_path)

        output_files.append(output_path)

    zip_path = os.path.join(
        temp_dir,
        "results.zip"
    )

    with zipfile.ZipFile(
        zip_path,
        "w"
    ) as zipf:

        for file in output_files:
            zipf.write(
                file,
                os.path.basename(file)
            )

    return output_files, zip_path

with gr.Blocks() as demo:

    gr.Markdown("# AI Image Upscaler")

    files = gr.File(
        file_count="multiple",
        label="画像を選択"
    )

    run_button = gr.Button(
        "高画質化開始"
    )

    gallery = gr.Gallery(
        label="結果"
    )

    zip_download = gr.File(
        label="ZIPダウンロード"
    )

    run_button.click(
        upscale_images,
        inputs=files,
        outputs=[
            gallery,
            zip_download
        ]
    )

demo.launch()