from diffusers import DiffusionPipeline, AutoencoderKL, DDIMScheduler, ControlNetModel, StableDiffusionControlNetPipeline, UniPCMultistepScheduler
import torch
from PIL import Image
from safetensors import safe_open

generator = torch.Generator("cuda").manual_seed(0)

# Ensure the VAE is compatible and correctly downloaded
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)

# Load the main diffusion pipeline
pipe = DiffusionPipeline.from_pretrained(
    "stablediffusionapi/epicrealism-xl",
    vae=vae,
    torch_dtype=torch.float16,
    # variant="fp16",
    # use_safetensors=True
)

# Set the scheduler
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)

# Optionally, load LoRA weights if necessary
pipe.load_lora_weights("pytorch_lora_weights.safetensors", weight_name="pytorch_lora_weights.safetensors", adapter_name="satvik_sdxl_new")

# Enable VAE slicing for lower memory usage
pipe.enable_vae_slicing()

# Move the pipeline to CUDA
pipe.to("cuda")

prompt = "a photo of sssaaatvik person drinking beer in a bar, high resolution, best quality, 8k, realistic"
negative_prompt = "ugly, deformed, disfigured, poor details, bad anatomy, cartoon"
output_pipe1 = pipe(prompt=prompt, negative_prompt = negative_prompt, num_inference_steps=50,generator=generator, guidance_scale = 4.0).images[0]
output_pipe1.save("beer_12.jpg")