import torch
from diffusers import StableDiffusionXLPipeline, AutoencoderKL, StableDiffusionXLImg2ImgPipeline, AutoPipelineForText2Image, \
    DDIMScheduler, EulerAncestralDiscreteScheduler
from transformers import CLIPTextModel

from utils.get_pipeline_embeds import get_pipeline_embeds

pipeline = StableDiffusionXLPipeline.from_single_file(
    "local_model",
    torch_dtype=torch.float16).to("cuda")

text_encoder = pipeline.text_encoder
tokenizer = pipeline.tokenizer

pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config)

vae = AutoencoderKL.from_single_file("local_vae").to("cuda")


pipeline.scheduler.beta_start = 0.00085
pipeline.scheduler.beta_end = 0.012

pipeline_text2image = AutoPipelineForText2Image.from_pipe(
    pipeline
)

prompt = 25 * "a photo of an astronaut riding a horse on mars"
negative_prompt = 25 * "earth, donkey, plants, water, cars"
num_inference_steps = 20  # Number of steps
guidance_scale = 7.5  # CFG scale


prompt_embeds, negative_prompt_embeds = get_pipeline_embeds(pipeline, prompt, negative_prompt, "cuda")

pooled_prompt_embeds = prompt_embeds.mean(dim=1)
pooled_negative_prompt_embeds = negative_prompt_embeds.mean(dim=1)

print("Prompt Embeds Shape:", prompt_embeds.shape)
print("Negative Prompt Embeds Shape:", negative_prompt_embeds.shape)


image = pipeline_text2image(
    prompt_embeds=prompt_embeds,
    pooled_prompt_embeds=pooled_prompt_embeds,
    negative_prompt_embeds=negative_prompt_embeds,
    negative_pooled_prompt_embeds=pooled_negative_prompt_embeds,
    num_inference_steps=num_inference_steps,
    guidance_scale=guidance_scale,
    text_encoder=text_encoder,
).images[0]