import gradio as gr
import os
import tempfile
from PIL import Image
import time
import torch
temp_dir=tempfile.gettempdir()

def process_images(condition_images, input_images):
    start = time.time()
    output_images = []
    for img in input_images:
        if img is not None:
            output_images.append(img) 
        else:
            gr.Error("Please upload at least one  image")
    pth_path = os.path.join(temp_dir, "output.pth")
    temp_data = {"test": "test"}
    PTH_data={"test":"test"}
    torch.save(PTH_data,pth_path)
    end=time.time()
    process_time = f"{end-start:.2f} s"
    return output_images, pth_path, process_time

with gr.Blocks() as demo:
        gr.Markdown("Title")
        with gr.Row():
            with gr.Group():
                condition_inputs = gr.Files(label="Condition Img", file_types=[".png", ".jpg", ".jpeg"], type='filepath')
                input_images = gr.Files(label="Input Img", file_types=[".png", ".jpg", ".jpeg"], type='filepath')
            with gr.Group():
                output_gallery = gr.Gallery(label="Output Img", show_label=True, columns=2)
                pth_output = gr.File(label="Download PTH ")
                process_time = gr.Textbox(label="Processing Time",  type="text",interactive=False)
        button1 = gr.Button("Upload Images")
        def func1(condition_files, input_files):
            condition_imgs = [Image.open(f) for f in condition_files] if condition_files else []
            input_imgs = [Image.open(f) for f in input_files] if input_files else []
            return process_images(condition_imgs, input_imgs)
        button1.click(
            fn=func1,
            inputs=[condition_inputs, input_images],
            outputs=[output_gallery, pth_output, process_time]
        )

    
demo.launch()