import torch
from PIL.Image import Image
from onediffx.deep_cache import StableDiffusionXLPipeline

from pipelines.models import TextToImageRequest
from diffusers import DDIMScheduler
from torch import Generator
from loss import SchedulerWrapper

from onediffx import compile_pipe, save_pipe, load_pipe

def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
    if not pipeline:
        pipeline = StableDiffusionXLPipeline.from_pretrained(
            "./models/newdream-sdxl-20",
            torch_dtype=torch.float16,
            local_files_only=True,
        ).to("cuda")
    
    pipeline.scheduler = SchedulerWrapper(DDIMScheduler.from_config(pipeline.scheduler.config))

    pipeline = compile_pipe(pipeline)
    load_pipe(pipeline,dir="robert_pipe")
    for _ in range(2):
        deepcache_output = pipeline(prompt="telestereography, unstrengthen, preadministrator, copatroness, hyperpersonal, paramountness, paranoid, guaniferous", output_type="pil", cache_interval=2, cache_layer_id=1, cache_block_id=0, num_inference_steps=20)
    pipeline.scheduler.prepare_loss()
    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,
        cache_interval=1,
        cache_layer_id=1,
        cache_block_id=0,
        eta=1.0,
        guidance_scale = 5.0,
        guidance_rescale = 0.0,
    ).images[0]