import os
import io
import base64
import torch
from torch.cuda import amp
import numpy as np
from PIL import Image
from diffusers import AutoPipelineForText2Image, AutoencoderKL, DPMSolverMultistepScheduler

pipe = None

def load_model(_model = None, _vae = None, loras = []):
    global pipe

    _model = _model or 'cagliostrolab/animagine-xl-3.0'

    if torch.cuda.is_available():
        torch_dtype = torch.float16
    else:
        torch_dtype = torch.float32

    if _vae:
        # "stabilityai/sdxl-vae"
        vae = AutoencoderKL.from_pretrained(_vae, torch_dtype=torch_dtype)
        pipe = AutoPipelineForText2Image.from_pretrained(
            _model,
            torch_dtype=torch_dtype,
            vae=vae,
            )
    else:
        pipe = AutoPipelineForText2Image.from_pretrained(
            _model,
            torch_dtype=torch_dtype,
            )

    # DPM++ 2M Karras
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(
        pipe.scheduler.config,
        algorithm_type="sde-dpmsolver++",
        use_karras_sigmas=True
        )

    for lora in loras:
        pipe.load_lora_weights(".", weight_name=lora + ".safetensors")

    if torch.cuda.is_available():
        pipe.to("cuda")

    pipe.enable_vae_slicing()

def pil_to_webp(img):
    buffer = io.BytesIO()
    img.save(buffer, 'webp')

    return buffer.getvalue()

def bin_to_base64(bin):
    return base64.b64encode(bin).decode('ascii')

def run(prompt = None, negative_prompt = None, model = None, guidance_scale = None, steps = None, seed = None):
    global pipe
    
    if not pipe:
        load_model(model)

    _prompt = "masterpiece, best quality, 1girl, portrait"
    _negative_prompt = "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name"

    prompt = prompt or _prompt
    negative_prompt = negative_prompt or _negative_prompt
    guidance_scale = float(guidance_scale) if guidance_scale else 5.0
    steps = int(steps) if steps else 20
    seed = int(seed) if seed else -1

    generator = None
    if seed != -1:
        generator = torch.manual_seed(seed)

    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=steps,
        clip_skip=2,
        generator=generator,
        ).images[0]

    return image