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

#from sfast.compilers.diffusion_pipeline_compiler import (compile,
#                                                        CompilationConfig)


#from onediff.schedulers import EulerDiscreteScheduler
from onediffx import compile_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,
        )
  #  pipeline.scheduler = EulerDiscreteScheduler.from_config(
   # pipeline.scheduler.config,)

    pipeline.to("cuda")
    pipeline = compile_pipe(pipeline)
    for _ in range(4):
        pipeline(prompt="Donald Trump kissed Kamala Harris", num_inference_steps=15)

    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,
        end_cfg=0.473,
        num_inference_steps=21,
    ).images[0]