File size: 2,709 Bytes
a1b31a2 |
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 StableDiffusionXLPipeline
import torch
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
pipe = StableDiffusionXLPipeline.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 []
|