import gradio as gr
import numpy as np
import torch
import logging
import os
import cv2 
from huggingface_hub import login
from diffusers import AutoPipelineForInpainting, UNet2DConditionModel
import diffusers
import spaces # type: ignore
from PIL import Image

login(os.getenv("HF_TOKEN"))

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

pipe = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16").to(device, dtype=torch.float16)

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

@spaces.GPU()
def predict(image, mask=None, positive_prompt="", negative_prompt=None, guidance_scale=7.5, steps=20, strength=1.0, scheduler="EulerDiscreteScheduler"):
    try:
        logging.info("Starting inpainting process")
        
        scheduler_parts = scheduler.split("-")
        scheduler_class_name = scheduler_parts[0]
        add_kwargs = {}

        if len(scheduler_parts) > 1:
            add_kwargs["use_karras"] = True
        if len(scheduler_parts) > 2:
            add_kwargs["algorithm_type"] = "sde-dpmsolver++"
        
        logging.info(f"Using scheduler: {scheduler_class_name} with args {add_kwargs}")
        
        scheduler_cls = getattr(diffusers, scheduler_class_name)

        source_image = Image.fromarray(image["background"]).convert("RGB")
        
        if np.unique(mask["background"]).size == 1:
            alpha_channel = image["layers"][0][:, :, 3]
            mask = np.where(alpha_channel == 255,255,0).astype(np.uint8)
            mask_4ch = np.stack([mask] * 4, axis=2)
            mask_image = Image.fromarray(mask_4ch)
            logging.info(f"Mask was not uploaded")
        else :
            mask_image = Image.fromarray(mask["background"]).convert("RGBA")
            print(type(mask["background"]))
            print("Value of Mask : ", mask)
            print("Value of Mask : ", mask["background"])
            logging.info(f"Mask was uploaded")
        
        

        logging.info("Converted images to required formats")

        # Set the scheduler
        pipe.scheduler = scheduler_cls.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler", **add_kwargs)
        logging.info("Scheduler set successfully")

        h = (source_image.height//8) * 8
        w = (source_image.width//8) * 8
        # Run the inpainting model
        output = pipe(
            prompt=positive_prompt, 
            negative_prompt=negative_prompt, 
            image=source_image, 
            height=h,
            width=w,
            mask_image=mask_image,
            guidance_scale=guidance_scale, 
            num_inference_steps=int(steps), 
            strength=strength
        )
        logging.info("Model inference completed successfully")

            
        return output.images[0], mask_image
    
    except Exception as e:
        logging.error(f"Error occurred: {str(e)}", exc_info=True)
        return None


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            image = gr.ImageMask(
                type="numpy", label="Input Image", layers=False
            )
            with gr.Row():
                positive_prompt = gr.Textbox(placeholder="Your prompt", label="Positive Prompt")
                negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="What you don't want to see in the image")
            with gr.Row():    
                btn = gr.Button("Inpaint")
            with gr.Accordion(label="Advanced Settings", open=False):
                with gr.Row():
                    guidance_scale = gr.Number(value=7.5, minimum=1.0, maximum=20.0, step=0.1, label="Guidance Scale")
                    steps = gr.Number(value=50, minimum=10, maximum=100, step=1, label="Steps")
                    strength = gr.Number(value=0.80, minimum=0.01, maximum=1.0, step=0.01, label="Strength")
                    
                with gr.Row():
                    schedulers = ["DEISMultistepScheduler", "HeunDiscreteScheduler", "EulerDiscreteScheduler", "DPMSolverMultistepScheduler", "DPMSolverMultistepScheduler-Karras", "DPMSolverMultistepScheduler-Karras-SDE"]
                    scheduler = gr.Dropdown(label="Schedulers", choices=schedulers, value="DPMSolverMultistepScheduler-Karras-SDE")
                with gr.Row():
                    mask = gr.ImageMask(label="Mask", format="png", value=None, sources=["upload"])
                    mask_image = gr.Image(label="Mask Image", format="png")
        with gr.Column():
            image_out = gr.Image(label="Output", format="png")
    
    btn.click(fn=predict, inputs=[image, mask, positive_prompt, negative_prompt, guidance_scale, steps, strength, scheduler], outputs=[image_out, mask_image])
    
demo.launch(debug=True,show_error=True)