import torch
from PIL.Image import Image
from diffusers import StableDiffusionXLPipeline
from pipelines.models import TextToImageRequest
from torch import Generator
from DeepCache import DeepCacheSDHelper

'''
def load_pipeline() -> StableDiffusionXLPipeline:
    pipeline = StableDiffusionXLPipeline.from_pretrained(
        "stablediffusionapi/newdream-sdxl-20",
        torch_dtype=torch.float16,
        local_files_only=True,
    ).to("cuda")

    pipeline(prompt="")

    return pipeline


def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> Image:
    if request.seed is None:
        generator = None
    else:
        generator = Generator(pipeline.device).manual_seed(request.seed)

    return pipeline(
        prompt=request.prompt,
        negative_prompt=request.negative_prompt,
        width=request.width,
        height=request.height,
        generator=generator,
        num_inference_steps=20,
    ).images[0]
'''

def load_pipeline() -> dict: #StableDiffusionXLPipeline, AutoPipelineForImage2Image:
    
    pipeline_dict = {}
    pipeline = StableDiffusionXLPipeline.from_pretrained(
        "stablediffusionapi/newdream-sdxl-20",
        torch_dtype=torch.float16,
        #local_files_only=True,
        use_safetensors=True,
        variant='fp16',        
    ).to("cuda")

    refiner = AutoPipelineForImage2Image.from_pretrained(
        'stabilityai/stable-diffusion-xl-refiner-1.0',
        use_safetensors=True,
        torch_dtype=torch.float16,
        variant='fp16',
        ).to('cuda')

    
    helper = DeepCacheSDHelper(pipe=pipeline)
    helper.set_params(cache_interval=3, cache_branch_id=0)
    helper.enable()    

    refiner_helper = DeepCacheSDHelper(pipe=refiner)
    refiner_helper.set_params(cache_interval=3, cache_branch_id=0)
    refiner_helper.enable()    
    

    for _ in range(5):
        pipeline(prompt="")

    pipeline_dict = {
        'base_pipeline': pipeline,
        'refiner': refiner
    }
    return pipeline_dict #base_pipeline, refiner


def infer(request: TextToImageRequest, pipeline_dict: dict) -> Image: #pipeline: StableDiffusionXLPipeline, refiner: AutoPipelineForImage2Image) -> Image:
    if request.seed is None:
        generator = None
    else:
        generator = Generator(pipeline_dict['base_pipeline'].device).manual_seed(request.seed)


    image = pipeline_dict['base_pipeline'](
        prompt=request.prompt,
        negative_prompt=request.negative_prompt,
        width=request.width,
        height=request.height,
        generator=generator,
        num_inference_steps=27,
        denoising_end=0.8,
        output_type='latent',
    ).images

    return pipeline_dict['refiner'](
        prompt=request.prompt,
        negative_prompt=request.negative_prompt,
        width=request.width,
        height=request.height,
        generator=generator,
        num_inference_steps=27,
    denoising_start=0.8,
    image=image,
  ).images[0]