File size: 2,709 Bytes
80e6c51
 
 
 
45c0511
80e6c51
 
45c0511
1a22f37
66a1189
 
80e6c51
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
66a1189
 
80e6c51
66a1189
 
80e6c51
 
 
 
45c0511
 
80e6c51
 
66a1189
 
80e6c51
66a1189
80e6c51
 
 
 
 
 
 
 
66a1189
80e6c51
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
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 []