import spaces


def load_pipeline():
    from diffusers import AutoPipelineForText2Image
    import torch
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    pipe = AutoPipelineForText2Image.from_pretrained(
        "John6666/rae-diffusion-xl-v2-sdxl-spo-pcm",
        #custom_pipeline="lpw_stable_diffusion_xl",
        custom_pipeline="nyanko7/sdxl_smoothed_energy_guidance",
        torch_dtype=torch.float16,
    )
    pipe.to(device)
    return pipe


def save_image(image, metadata, output_dir):
    import os
    import uuid
    import json
    from PIL import PngImagePlugin
    filename = str(uuid.uuid4()) + ".png"
    os.makedirs(output_dir, exist_ok=True)
    filepath = os.path.join(output_dir, filename)
    metadata_str = json.dumps(metadata)
    info = PngImagePlugin.PngInfo()
    info.add_text("metadata", metadata_str)
    image.save(filepath, "PNG", pnginfo=info)
    return filepath


pipe = load_pipeline()


@spaces.GPU
def generate_image(prompt, neg_prompt):
    metadata = {
        "prompt": prompt + ", anime, masterpiece, best quality, very aesthetic, absurdres",
        "negative_prompt": neg_prompt + ", bad hands, bad feet, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract], photo, deformed, disfigured, low contrast, photo, deformed, disfigured, low contrast",
        "resolution": f"{1024} x {1024}",
        "guidance_scale": 7.0,
        "num_inference_steps": 28,
        "sampler": "Euler",
    }
    try: 
        images = pipe(
            prompt=prompt + ", anime, masterpiece, best quality, very aesthetic, absurdres",
            negative_prompt=neg_prompt + ", bad hands, bad feet, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract], photo, deformed, disfigured, low contrast, photo, deformed, disfigured, low contrast",
            width=1024,
            height=1024,
            guidance_scale=7.0, seg_scale=3.0, seg_applied_layers=["mid"],
            num_inference_steps=28,
            output_type="pil",
            #clip_skip=1,
        ).images
        if images:
            image_paths = [
                save_image(image, metadata, "./outputs")
                for image in images
            ]
        return image_paths
    except Exception as e:
        print(e)
        return []