from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler, AutoencoderKL,DiffusionPipeline
import torch
from compel import Compel, DiffusersTextualInversionManager

model_id = "cagliostrolab/animagine-xl-3.1"
# model_id = "model/Counterfeit-V3.0_fp16.safetensors"
pipe = DiffusionPipeline.from_pretrained(
# pipe = StableDiffusionPipeline.from_pretrained(
# pipe = StableDiffusionPipeline.from_single_file(
    model_id,
    torch_dtype=torch.float16,
    use_safetensors=True
    # vae=AutoencoderKL.from_single_file("vae/Counterfeit-V2.5.vae.pt"),
    # load_safety_checker=False,
    # extract_ema=True
    )
# pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
# pipe.load_textual_inversion("embeddings/EasyNegative.safetensors", token="EasyNegative")
pipe.to("cuda")

prompt = "(high resolution)++, 8k+, attractive, highly detailed, pretty girl, short hair"
negative_prompt = "EasyNegative, (Worst Quality)++, (low quality)+"

# textual_inversion_manager = DiffusersTextualInversionManager(pipe)
# compel_proc = Compel(
#     tokenizer=pipe.tokenizer,
#     text_encoder=pipe.text_encoder,
#     textual_inversion_manager=textual_inversion_manager,
#     truncate_long_prompts=False)

# prompt_embeds = compel_proc([prompt])
# negative_prompt_embeds = compel_proc([negative_prompt])

seed = 20000
generator = torch.Generator(device="cuda").manual_seed(seed)
image = pipe(
    # prompt_embeds=prompt_embeds,
    # negative_prompt_embeds=negative_prompt_embeds,
    prompt=prompt,
    negative_prompt=negative_prompt,
    # generator=generator,
    num_inference_steps=30,
    guidance_scale=7.0,
    width=768,
    height=768,
    ).images[0]
image.save("result.png")