import torch
from pathlib import Path
from PIL.Image import Image
from diffusers import StableDiffusionXLPipeline, DDIMScheduler
from pipelines.models import TextToImageRequest
from torch import Generator
from cache_diffusion import cachify
from pipe.deploy import compile
from loss import SchedulerWrapper

generator = Generator(torch.device("cuda")).manual_seed(6969)
prompt = "Make submissions great again"
SDXL_DEFAULT_CONFIG = [
            {
                "wildcard_or_filter_func": lambda name: "down_blocks.2" not in name and"down_blocks.3" not in name and "up_blocks.2" not in name,
                "select_cache_step_func": lambda step: (step % 2 != 0) and (step >= 10),
                }]
def load_pipeline() -> StableDiffusionXLPipeline:
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "stablediffusionapi/newdream-sdxl-20",torch_dtype=torch.float16, use_safetensors=True
    ).to("cuda") 
    compile(
        pipe,
        onnx_path=Path("/home/sandbox/.cache/huggingface/hub/models--RobertML--edge-onnx/snapshots/d56fa8ea1dc675b87de08eece735bc5ec80a247f"),
        engine_path=Path("/home/sandbox/.cache/huggingface/hub/models--RobertML--edge-engine/snapshots/e0dd02be0c58057947801857c41839f76df2fc88"),
        batch_size=1,
    )
    cachify.prepare(pipe, SDXL_DEFAULT_CONFIG)
    cachify.enable(pipe)
    pipe.scheduler = SchedulerWrapper(DDIMScheduler.from_config(pipe.scheduler.config))
    with cachify.infer(pipe) as cached_pipe:
        for _ in range(4):
            pipe(prompt=prompt, num_inference_steps=20)
        pipe.scheduler.prepare_loss()
    cachify.disable(pipe)
    return pipe

def infer(request: TextToImageRequest, pipeline: StableDiffusionXLPipeline) -> Image:

    if request.seed is None:
        generator = None
    else:
        generator = Generator(pipeline.device).manual_seed(request.seed)
    cachify.enable(pipeline)
    with cachify.infer(pipeline) as cached_pipe:
        image = cached_pipe(
            prompt=request.prompt,
            negative_prompt=request.negative_prompt,
            width=request.width,
            height=request.height,
            generator=generator,
            num_inference_steps=13,
            ).images[0]
    return image