import gradio as gr
from PIL import Image, ImageOps, ImageEnhance
import torch
from diffusers import DiffusionPipeline

device = "cuda" if torch.cuda.is_available() else "cpu"

if torch.cuda.is_available():
    torch.cuda.max_memory_allocated(device=device)
    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
    pipe.enable_xformers_memory_efficient_attention()
    pipe = pipe.to(device)
else:
    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
    pipe = pipe.to(device)

def edit_image(image, operation, *args):
    if operation == "rotate":
        angle = int(args[0])
        image = image.rotate(angle, expand=True)
    elif operation == "crop":
        left, top, right, bottom = map(int, args)
        image = image.crop((left, top, right, bottom))
    elif operation == "resize":
        width, height = map(int, args)
        image = image.resize((width, height))
    elif operation == "flip":
        if args[0] == "horizontal":
            image = ImageOps.mirror(image)
        else:
            image = ImageOps.flip(image)
    elif operation == "color":
        factor = float(args[0])
        image = ImageEnhance.Color(image).enhance(factor)
    
    return image

def create_demo():
    with gr.Blocks() as demo:
        with gr.Row():
            gr.Markdown("# Image Editor")

        with gr.Row():
            with gr.Column():
                edit_operation = gr.Dropdown(choices=["rotate", "crop", "resize", "flip", "color"], label="Edit Operation")
                input_args = gr.State([])

                angle = gr.Slider(0, 360, step=1, label="Angle", visible=False)
                left = gr.Slider(0, 500, step=1, label="Left", visible=False)
                top = gr.Slider(0, 500, step=1, label="Top", visible=False)
                right = gr.Slider(0, 500, step=1, label="Right", visible=False)
                bottom = gr.Slider(0, 500, step=1, label="Bottom", visible=False)
                width = gr.Slider(50, 1000, step=1, label="Width", visible=False)
                height = gr.Slider(50, 1000, step=1, label="Height", visible=False)
                flip_direction = gr.Dropdown(choices=["horizontal", "vertical"], label="Direction", visible=False)
                color_factor = gr.Slider(0.1, 2.0, step=0.1, label="Color Factor", visible=False)

                edit_button = gr.Button("Edit Image")
                uploaded_image = gr.Image(label="Upload Image", type="pil")
                edited_image = gr.Image(label="Edited Image", type="pil", interactive=True)

                def update_inputs(operation):
                    if operation == "rotate":
                        return [gr.update(visible=True), angle]
                    elif operation == "crop":
                        return [gr.update(visible=True), left, top, right, bottom]
                    elif operation == "resize":
                        return [gr.update(visible=True), width, height]
                    elif operation == "flip":
                        return [gr.update(visible=True), flip_direction]
                    elif operation == "color":
                        return [gr.update(visible=True), color_factor]
                    else:
                        return []

                edit_operation.change(fn=update_inputs, inputs=[edit_operation], outputs=[angle, left, top, right, bottom, width, height, flip_direction, color_factor])

                edit_button.click(
                    fn=lambda img_data, operation, *args: edit_image(img_data, operation, *args),
                    inputs=[uploaded_image, edit_operation, angle, left, top, right, bottom, width, height, flip_direction, color_factor],
                    outputs=[edited_image]
                )

    return demo

demo = create_demo()
demo.queue().launch()