import spaces
import gradio as gr
import torch
from diffusers import LCMScheduler, AutoPipelineForText2Image
from diffusers import AutoPipelineForInpainting, LCMScheduler
from diffusers import DiffusionPipeline, LCMScheduler
from PIL import Image, ImageEnhance
import io

@spaces.GPU
def generate_image(prompt, num_inference_steps, guidance_scale):
    model_id = "stabilityai/stable-diffusion-xl-base-1.0"
    adapter_id = "latent-consistency/lcm-lora-sdxl"

    pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch.float32, variant="fp16")
    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
    pipe.to("cuda")

    # Load and fuse lcm lora
    pipe.load_lora_weights(adapter_id)
    pipe.fuse_lora()

    # Generate the image
    image = pipe(prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
  
    return image

def inpaint_image(prompt, init_image, mask_image, num_inference_steps, guidance_scale):
    pipe = AutoPipelineForInpainting.from_pretrained(
        "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
        torch_dtype=torch.float32,
        variant="fp16",
    ).to("cuda")
    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
    pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl")
    pipe.fuse_lora()

    if init_image is not None:
        init_image_path = init_image.name  # Get the file path
        init_image = Image.open(init_image_path).resize((1024, 1024))
    else:
        raise ValueError("Initial image not provided or invalid")

    if mask_image is not None:
        mask_image_path = mask_image.name  # Get the file path
        mask_image = Image.open(mask_image_path).resize((1024, 1024))
    else:
        raise ValueError("Mask image not provided or invalid")

    # Generate the inpainted image
    generator = torch.manual_seed(42)
    image = pipe(
        prompt=prompt,
        image=init_image,
        mask_image=mask_image,
        generator=generator,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
    ).images[0]

    return image

def generate_image_with_adapter(prompt, num_inference_steps, guidance_scale):
    pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    variant="fp16",
    torch_dtype=torch.float32
    ).to("cuda")

    # set scheduler
    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)

    # Load and fuse lcm lora
    pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl", adapter_name="lcm")
    pipe.load_lora_weights("Pclanglais/Mickey-1928", weight_name="pytorch_lora_weights.safetensors", adapter_name="mickey")

    # Combine LoRAs
    pipe.set_adapters(["lcm", "mickey"], adapter_weights=[1.0, 0.8])
    pipe.fuse_lora()
    generator = torch.manual_seed(0)
    # Generate the image
    image = pipe(prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=generator).images[0]
    pipe.unfuse_lora()
    return image


def modify_image(image, brightness, contrast):
    # Function to modify brightness and contrast
    image = Image.open(io.BytesIO(image))
    enhancer = ImageEnhance.Brightness(image)
    image = enhancer.enhance(brightness)
    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(contrast)
    return image

with gr.Blocks(gr.themes.Soft()) as demo:
    with gr.Row():
        image_output = gr.Image(label="Generated Image")

    with gr.Row():
        with gr.Accordion(label="Configuration Options"):
            prompt_input = gr.Textbox(label="Prompt", placeholder="Self-portrait oil painting, a beautiful cyborg with golden hair, 8k")
            steps_input = gr.Slider(minimum=1, maximum=10, label="Inference Steps", value=4)
            guidance_input = gr.Slider(minimum=0, maximum=2, label="Guidance Scale", value=1)
            generate_button = gr.Button("Generate Image")
    with gr.Row():
        with gr.Accordion(label="Wiki-Mouse Image Generation"):
            adapter_prompt_input = gr.Textbox(label="Prompt", placeholder="papercut, a cute fox")
            adapter_steps_input = gr.Slider(minimum=1, maximum=10, label="Inference Steps", value=4)
            adapter_guidance_input = gr.Slider(minimum=0, maximum=2, label="Guidance Scale", value=1)
            adapter_generate_button = gr.Button("Generate Image with Adapter")

    

    generate_button.click(
        generate_image,
        inputs=[prompt_input, steps_input, guidance_input],
        outputs=image_output
    )

    modify_button.click(
        modify_image,
        inputs=[image_output, brightness_slider, contrast_slider],
        outputs=image_output
    )
    inpaint_button.click(
        inpaint_image,
        inputs=[inpaint_prompt_input, init_image_input, mask_image_input, inpaint_steps_input, inpaint_guidance_input],
        outputs=image_output
    )
    adapter_generate_button.click(
        generate_image_with_adapter,
        inputs=[adapter_prompt_input, adapter_steps_input, adapter_guidance_input],
        outputs=image_output
    )

demo.launch()