from logger import rich_logger as l
from wandb.integration.diffusers import autolog
from config import Project_Name
from clear_memory import clear_memory
import numpy as np
import torch
from PIL import Image
from mask_generator import invert_mask
from diffusers.utils import load_image
from pipeline import fetch_kandinsky_pipeline
from config import controlnet_adapter_model_name,controlnet_base_model_name,kandinsky_model_name
import cv2 
from diffusers import StableDiffusionInpaintPipeline, DPMSolverMultistepScheduler



autolog(init=dict(project=Project_Name))







def make_inpaint_condition(init_image, mask_image):
    init_image = np.array(init_image.convert("RGB")).astype(np.float32) / 255.0
    mask_image = np.array(mask_image.convert("L")).astype(np.float32) / 255.0

    assert init_image.shape[0:1] == mask_image.shape[0:1], "image and image_mask must have the same image size"
    init_image[mask_image > 0.5] = -1.0  # set as masked pixel
    init_image = np.expand_dims(init_image, 0).transpose(0, 3, 1, 2)
    init_image = torch.from_numpy(init_image)
    return init_image
    


    







def kandinsky_inpainting_inference(prompt, negative_prompt, image, mask_image):
    """
    Perform Kandinsky inpainting inference on the given image.

    Args:
        prompt (str): The prompt for the inpainting process.
        negative_prompt (str): The negative prompt for the inpainting process.
        image (PIL.Image.Image): The input image to be inpainted.
        mask_image (PIL.Image.Image): The mask image indicating the areas to be inpainted.

    Returns:
        PIL.Image.Image: The output inpainted image.
    """
    clear_memory()
    pipe = fetch_kandinsky_pipeline(controlnet_adapter_model_name, controlnet_base_model_name,kandinsky_model_name, image)
    output_image = pipe(prompt=prompt,negative_prompt=negative_prompt,image=image,mask_image=mask_image,num_inference_steps=800,strength=1.0,guidance_scale = 7.8,height = 1472, width = 2560).images[0]
    return output_image

    


  

def sd2_inpainting_inference(prompt, img, mask, repo_id="stabilityai/stable-diffusion-2-inpainting", revision="fp16"):
    """
    Generate an image based on a prompt using a pretrained model.

    Args:
        prompt (str): The prompt for the image generation.
        img_url (str): The URL of the initial image.
        mask_url (str): The URL of the mask image.
        repo_id (str, optional): The ID of the repository of the pretrained model. Defaults to "stabilityai/stable-diffusion-2-inpainting".
        revision (str, optional): The revision of the pretrained model. Defaults to "fp16".

    Returns:
        Image: The generated image.
    """
    init_image = load_image(img)
    mask_image = load_image(mask)
    pipe = StableDiffusionInpaintPipeline.from_pretrained(
    repo_id,
    torch_dtype=torch.float16)
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    pipe = pipe.to("cuda")
    image = pipe(prompt=prompt, image=init_image, mask_image=mask_image, num_inference_steps=400,guidence_scale=7.5).images[0]
    return image