Upload 5 files
Browse files- config.toml +68 -0
- lpw_stable_diffusion_xl.py +0 -0
- requirements.txt +9 -6
- style.css +173 -0
- utils.py +179 -0
config.toml
ADDED
|
@@ -0,0 +1,68 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[model]
|
| 2 |
+
min_image_size = 1024
|
| 3 |
+
max_image_size = 2880
|
| 4 |
+
use_torch_compile = false
|
| 5 |
+
enable_cpu_offload = false
|
| 6 |
+
output_dir = "./outputs"
|
| 7 |
+
|
| 8 |
+
[[models]]
|
| 9 |
+
path = "https://huggingface.co/IbarakiDouji/WAI-NSFW-illustrious-SDXL/blob/main/waiNSFWIllustrious_v150.safetensors"
|
| 10 |
+
name = "v15"
|
| 11 |
+
|
| 12 |
+
[[models]]
|
| 13 |
+
path = "https://huggingface.co/IbarakiDouji/WAI-NSFW-illustrious-SDXL/blob/main/WAI-NSFW-illustrious-SDXL-v14.safetensors"
|
| 14 |
+
name = "v14"
|
| 15 |
+
|
| 16 |
+
[[models]]
|
| 17 |
+
path = "https://huggingface.co/IbarakiDouji/WAI-NSFW-illustrious-SDXL/blob/main/WAI-NSFW-illustrious-SDXL-v13.safetensors"
|
| 18 |
+
name = "v13"
|
| 19 |
+
|
| 20 |
+
[[models]]
|
| 21 |
+
path = "https://huggingface.co/IbarakiDouji/WAI-NSFW-illustrious-SDXL/blob/main/WAI-NSFW-illustrious-SDXL-v12.safetensors"
|
| 22 |
+
name = "v12"
|
| 23 |
+
|
| 24 |
+
[[models]]
|
| 25 |
+
path = "https://huggingface.co/IbarakiDouji/WAI-NSFW-illustrious-SDXL/blob/main/WAI-NSFW-illustrious-SDXL-v11.safetensors"
|
| 26 |
+
name = "v11"
|
| 27 |
+
|
| 28 |
+
[[models]]
|
| 29 |
+
path = "https://huggingface.co/IbarakiDouji/WAI-NSFW-illustrious-SDXL/blob/main/WAI-NSFW-illustrious-SDXL-v10.safetensors"
|
| 30 |
+
name = "v10"
|
| 31 |
+
|
| 32 |
+
[prompts]
|
| 33 |
+
quality_tags = "masterpiece, best quality, amazing quality, {prompt}"
|
| 34 |
+
default_negative = "sensitive, nsfw, explicit, bad quality, worst quality, worst detail, sketch, censor"
|
| 35 |
+
default_aspect_ratio = "1248 x 1824"
|
| 36 |
+
examples = [
|
| 37 |
+
"1girl, souryuu asuka langley, neon genesis evangelion, eyepatch, red plugsuit, sitting, on throne, crossed legs, head tilt, holding weapon, lance of longinus \\\\(evangelion\\\\), cowboy shot, depth of field, faux traditional media, painterly, impressionism, photo background",
|
| 38 |
+
"1boy, vash the stampede, trigun stampede, red jacket, sunglasses, gun, hand on own hip, aiming, standing, looking at viewer, upper body, desert, cliff, cowboy shot",
|
| 39 |
+
"1girl, vertin \\(reverse:1999\\), reverse:1999, black umbrella, headwear, suitcase, looking at viewer, rain, night, city, bridge, from side, dutch angle, upper body",
|
| 40 |
+
"1girl, 1boy, c.c., lelouch vi britannia, year 2024, code geass, ascot, bare shoulders, black choker, black hair, blue flower, blue rose, breasts, bush, choker, closed mouth, collarbone, couple, dress, flower, frills, green hair, hetero, long hair, long sleeves, looking at another, medium breasts, off-shoulder dress, off shoulder, parted lips, purple ascot, purple eyes, rose, short hair, sitting, straight hair, yellow eyes",
|
| 41 |
+
"1girl, hatsune miku, vocaloid, blue eyes, blue hair, bowl, can, chopsticks, collared shirt, detached sleeves, eating, elbow rest, fish \\(food\\), food, holding, holding chopsticks, katsudon \\(food\\), long hair, long sleeves, looking at viewer, meal, nail polish, necktie, noodles, onigiri, plate, ramen, sashimi, shirt, shrimp, shrimp tempura, sleeveless, sleeveless shirt, solo, spring onion, tempura, twintails",
|
| 42 |
+
"4girls, multiple girls, gotoh hitori, ijichi nijika, kita ikuyo, yamada ryo, bocchi the rock!, ahoge, black shirt, blank eyes, blonde hair, blue eyes, blue hair, brown sweater, collared shirt, cube hair ornament, detached ahoge, empty eyes, green eyes, hair ornament, hairclip, kessoku band, long sleeves, looking at viewer, medium hair, mole, mole under eye, one side up, pink hair, pink track suit, red eyes, red hair, sailor collar, school uniform, serafuku, shirt, shuka high school uniform, side ahoge, side ponytail, sweater, sweater vest, track suit, white shirt, yellow eyes, painterly, impressionism, faux traditional media, v, double v, waving",
|
| 43 |
+
"1other, solo, outdoors, sky, arm up, night, earth, helmet, outstretched arm, star \\(sky\\), night sky, full moon, floating, starry sky, reaching, jumping, space, cowboy shot, ambiguous gender, spacesuit, moonlight, space helmet, astronaut, horror, black and white, monochromatic, high contrast, abstract background, dutch angle, dark, depth of field, chromatic aberration, faux traditional media"
|
| 44 |
+
]
|
| 45 |
+
|
| 46 |
+
[samplers]
|
| 47 |
+
list = [
|
| 48 |
+
"DPM++ 2M Karras",
|
| 49 |
+
"DPM++ SDE Karras",
|
| 50 |
+
"DPM++ 2M SDE Karras",
|
| 51 |
+
"Euler",
|
| 52 |
+
"Euler a",
|
| 53 |
+
"DDIM"
|
| 54 |
+
]
|
| 55 |
+
|
| 56 |
+
[aspect_ratios]
|
| 57 |
+
list = [
|
| 58 |
+
"1536 x 1536",
|
| 59 |
+
"1728 x 1344",
|
| 60 |
+
"1344 x 1728",
|
| 61 |
+
"1824 x 1248",
|
| 62 |
+
"1248 x 1824",
|
| 63 |
+
"2016 x 1152",
|
| 64 |
+
"1152 x 2016",
|
| 65 |
+
"2304 x 960",
|
| 66 |
+
"960 x 2304",
|
| 67 |
+
"Custom"
|
| 68 |
+
]
|
lpw_stable_diffusion_xl.py
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
requirements.txt
CHANGED
|
@@ -1,6 +1,9 @@
|
|
| 1 |
-
accelerate
|
| 2 |
-
diffusers
|
| 3 |
-
|
| 4 |
-
|
| 5 |
-
|
| 6 |
-
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
accelerate>=1.2.1
|
| 2 |
+
diffusers>=0.32.1
|
| 3 |
+
gradio==4.44.1
|
| 4 |
+
hf-transfer>=0.1.9
|
| 5 |
+
spaces>=0.32.0
|
| 6 |
+
torch>=2.4.0
|
| 7 |
+
transformers>=4.48.0
|
| 8 |
+
tomli>=2.0.1
|
| 9 |
+
pydantic==2.10.6
|
style.css
ADDED
|
@@ -0,0 +1,173 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
:root {
|
| 2 |
+
--title-font-size: clamp(1.5rem, 6vw, 3rem);
|
| 3 |
+
--subtitle-font-size: clamp(1rem, 2vw, 1.2rem);
|
| 4 |
+
--text-color: #fff;
|
| 5 |
+
--font-family: 'Helvetica Neue', sans-serif;
|
| 6 |
+
--gradient-primary: linear-gradient(45deg, #c136eb, #4EACEF);
|
| 7 |
+
--primary-color: #1565c0;
|
| 8 |
+
--primary-hover: #1976d2;
|
| 9 |
+
--border-radius: 12px;
|
| 10 |
+
--box-shadow: 0 2px 6px rgba(0, 0, 0, 0.05);
|
| 11 |
+
}
|
| 12 |
+
|
| 13 |
+
body {
|
| 14 |
+
font-family: var(--font-family);
|
| 15 |
+
color: var(--text-color);
|
| 16 |
+
margin: 0;
|
| 17 |
+
padding: 0;
|
| 18 |
+
min-height: 100vh;
|
| 19 |
+
}
|
| 20 |
+
|
| 21 |
+
.header {
|
| 22 |
+
text-align: center;
|
| 23 |
+
padding: 1rem 0;
|
| 24 |
+
margin-bottom: 0.5rem;
|
| 25 |
+
}
|
| 26 |
+
|
| 27 |
+
.title {
|
| 28 |
+
font-size: var(--title-font-size);
|
| 29 |
+
font-weight: 700;
|
| 30 |
+
text-transform: uppercase;
|
| 31 |
+
margin-bottom: 0.25rem;
|
| 32 |
+
background-image: var(--gradient-primary);
|
| 33 |
+
-webkit-text-fill-color: transparent;
|
| 34 |
+
-webkit-background-clip: text;
|
| 35 |
+
background-clip: text;
|
| 36 |
+
display: inline-block;
|
| 37 |
+
}
|
| 38 |
+
|
| 39 |
+
.subtitle {
|
| 40 |
+
font-size: var(--subtitle-font-size);
|
| 41 |
+
color: #999;
|
| 42 |
+
margin-bottom: 0.5rem;
|
| 43 |
+
}
|
| 44 |
+
|
| 45 |
+
.status {
|
| 46 |
+
display: none;
|
| 47 |
+
}
|
| 48 |
+
|
| 49 |
+
#duplicate-button {
|
| 50 |
+
margin: 1.5rem auto 2.5rem;
|
| 51 |
+
color: #fff;
|
| 52 |
+
background: #1565c0;
|
| 53 |
+
border-radius: 100vh;
|
| 54 |
+
padding: 0.75rem 1.5rem;
|
| 55 |
+
font-weight: 500;
|
| 56 |
+
box-shadow: 0 2px 8px rgba(21, 101, 192, 0.25);
|
| 57 |
+
transition: all 0.2s ease;
|
| 58 |
+
display: block;
|
| 59 |
+
}
|
| 60 |
+
|
| 61 |
+
#duplicate-button:hover {
|
| 62 |
+
background: #c136eb;
|
| 63 |
+
box-shadow: 0 4px 12px rgba(21, 101, 192, 0.35);
|
| 64 |
+
transform: translateY(-1px);
|
| 65 |
+
}
|
| 66 |
+
|
| 67 |
+
.contain {
|
| 68 |
+
max-width: 80%;
|
| 69 |
+
margin: 2rem auto;
|
| 70 |
+
padding: 2rem 1.5rem;
|
| 71 |
+
}
|
| 72 |
+
|
| 73 |
+
/* Component styling */
|
| 74 |
+
.gr-box {
|
| 75 |
+
border-radius: var(--border-radius);
|
| 76 |
+
border: 1px solid #e0e0e0;
|
| 77 |
+
background: #ffffff;
|
| 78 |
+
box-shadow: var(--box-shadow);
|
| 79 |
+
transition: box-shadow 0.2s ease;
|
| 80 |
+
}
|
| 81 |
+
|
| 82 |
+
.gr-box:hover {
|
| 83 |
+
box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
|
| 84 |
+
}
|
| 85 |
+
|
| 86 |
+
.gr-button.primary {
|
| 87 |
+
background: var(--primary-color);
|
| 88 |
+
border-radius: var(--border-radius);
|
| 89 |
+
padding: 0.8rem 2rem;
|
| 90 |
+
font-weight: 500;
|
| 91 |
+
box-shadow: 0 2px 8px rgba(21, 101, 192, 0.25);
|
| 92 |
+
transition: all 0.2s ease;
|
| 93 |
+
text-transform: uppercase;
|
| 94 |
+
letter-spacing: 0.5px;
|
| 95 |
+
}
|
| 96 |
+
|
| 97 |
+
.gr-button.primary:hover {
|
| 98 |
+
background: var(--primary-hover);
|
| 99 |
+
box-shadow: 0 4px 12px rgba(21, 101, 192, 0.35);
|
| 100 |
+
transform: translateY(-1px);
|
| 101 |
+
}
|
| 102 |
+
|
| 103 |
+
/* Form elements */
|
| 104 |
+
.gr-form {
|
| 105 |
+
background: #fff;
|
| 106 |
+
padding: 1.5rem;
|
| 107 |
+
border-radius: var(--border-radius);
|
| 108 |
+
box-shadow: var(--box-shadow);
|
| 109 |
+
}
|
| 110 |
+
|
| 111 |
+
.gr-input, .gr-textarea {
|
| 112 |
+
border: 1px solid #e0e0e0;
|
| 113 |
+
border-radius: 8px;
|
| 114 |
+
padding: 0.8rem;
|
| 115 |
+
transition: all 0.2s ease;
|
| 116 |
+
}
|
| 117 |
+
|
| 118 |
+
.gr-input:focus, .gr-textarea:focus {
|
| 119 |
+
border-color: var(--primary-color);
|
| 120 |
+
box-shadow: 0 0 0 2px rgba(21, 101, 192, 0.1);
|
| 121 |
+
}
|
| 122 |
+
|
| 123 |
+
/* Accordion styling */
|
| 124 |
+
.gr-accordion {
|
| 125 |
+
border: none;
|
| 126 |
+
margin: 1rem 0;
|
| 127 |
+
}
|
| 128 |
+
|
| 129 |
+
.gr-accordion-header {
|
| 130 |
+
background: #f8f9fa;
|
| 131 |
+
border-radius: var(--border-radius);
|
| 132 |
+
padding: 1rem;
|
| 133 |
+
font-weight: 500;
|
| 134 |
+
}
|
| 135 |
+
|
| 136 |
+
/* Gallery styling */
|
| 137 |
+
.gr-gallery {
|
| 138 |
+
background: #fff;
|
| 139 |
+
padding: 1rem;
|
| 140 |
+
border-radius: var(--border-radius);
|
| 141 |
+
box-shadow: var(--box-shadow);
|
| 142 |
+
}
|
| 143 |
+
|
| 144 |
+
/* Tooltips */
|
| 145 |
+
.gr-form small {
|
| 146 |
+
color: #666;
|
| 147 |
+
font-size: 0.875rem;
|
| 148 |
+
margin-top: 0.25rem;
|
| 149 |
+
display: block;
|
| 150 |
+
}
|
| 151 |
+
|
| 152 |
+
/* Responsive layout */
|
| 153 |
+
@media (max-width: 768px) {
|
| 154 |
+
.contain {
|
| 155 |
+
max-width: 90%;
|
| 156 |
+
padding: 1rem;
|
| 157 |
+
}
|
| 158 |
+
|
| 159 |
+
.gr-box {
|
| 160 |
+
margin: 0.5rem 0;
|
| 161 |
+
}
|
| 162 |
+
|
| 163 |
+
.gr-button.primary {
|
| 164 |
+
width: 100%;
|
| 165 |
+
}
|
| 166 |
+
}
|
| 167 |
+
|
| 168 |
+
@media (min-width: 1200px) {
|
| 169 |
+
.contain {
|
| 170 |
+
max-width: 1400px;
|
| 171 |
+
padding: 2.5rem 2rem;
|
| 172 |
+
}
|
| 173 |
+
}
|
utils.py
ADDED
|
@@ -0,0 +1,179 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import gc
|
| 2 |
+
import os
|
| 3 |
+
import random
|
| 4 |
+
import numpy as np
|
| 5 |
+
import json
|
| 6 |
+
import torch
|
| 7 |
+
import uuid
|
| 8 |
+
from PIL import PngImagePlugin
|
| 9 |
+
from datetime import datetime
|
| 10 |
+
from dataclasses import dataclass
|
| 11 |
+
from typing import Callable, Dict, Optional, Tuple, Any
|
| 12 |
+
from diffusers import (
|
| 13 |
+
DDIMScheduler,
|
| 14 |
+
DPMSolverMultistepScheduler,
|
| 15 |
+
DPMSolverSinglestepScheduler,
|
| 16 |
+
EulerAncestralDiscreteScheduler,
|
| 17 |
+
EulerDiscreteScheduler,
|
| 18 |
+
AutoencoderKL,
|
| 19 |
+
)
|
| 20 |
+
from lpw_stable_diffusion_xl import SDXLLongPromptWeightingPipeline
|
| 21 |
+
import logging
|
| 22 |
+
|
| 23 |
+
MAX_SEED = np.iinfo(np.int32).max
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
@dataclass
|
| 27 |
+
class StyleConfig:
|
| 28 |
+
prompt: str
|
| 29 |
+
negative_prompt: str
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
|
| 33 |
+
if randomize_seed:
|
| 34 |
+
seed = random.randint(0, MAX_SEED)
|
| 35 |
+
return seed
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def seed_everything(seed: int) -> torch.Generator:
|
| 39 |
+
torch.manual_seed(seed)
|
| 40 |
+
torch.cuda.manual_seed_all(seed)
|
| 41 |
+
np.random.seed(seed)
|
| 42 |
+
generator = torch.Generator()
|
| 43 |
+
generator.manual_seed(seed)
|
| 44 |
+
return generator
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def parse_aspect_ratio(aspect_ratio: str) -> Optional[Tuple[int, int]]:
|
| 48 |
+
if aspect_ratio == "Custom":
|
| 49 |
+
return None
|
| 50 |
+
width, height = aspect_ratio.split(" x ")
|
| 51 |
+
return int(width), int(height)
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def aspect_ratio_handler(
|
| 55 |
+
aspect_ratio: str, custom_width: int, custom_height: int
|
| 56 |
+
) -> Tuple[int, int]:
|
| 57 |
+
if aspect_ratio == "Custom":
|
| 58 |
+
return custom_width, custom_height
|
| 59 |
+
else:
|
| 60 |
+
width, height = parse_aspect_ratio(aspect_ratio)
|
| 61 |
+
return width, height
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def get_scheduler(scheduler_config: Dict, name: str) -> Optional[Callable]:
|
| 65 |
+
scheduler_factory_map = {
|
| 66 |
+
"DPM++ 2M Karras": lambda: DPMSolverMultistepScheduler.from_config(
|
| 67 |
+
scheduler_config, use_karras_sigmas=True
|
| 68 |
+
),
|
| 69 |
+
"DPM++ SDE Karras": lambda: DPMSolverSinglestepScheduler.from_config(
|
| 70 |
+
scheduler_config, use_karras_sigmas=True
|
| 71 |
+
),
|
| 72 |
+
"DPM++ 2M SDE Karras": lambda: DPMSolverMultistepScheduler.from_config(
|
| 73 |
+
scheduler_config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++"
|
| 74 |
+
),
|
| 75 |
+
"Euler": lambda: EulerDiscreteScheduler.from_config(scheduler_config),
|
| 76 |
+
"Euler a": lambda: EulerAncestralDiscreteScheduler.from_config(
|
| 77 |
+
scheduler_config
|
| 78 |
+
),
|
| 79 |
+
"DDIM": lambda: DDIMScheduler.from_config(scheduler_config),
|
| 80 |
+
}
|
| 81 |
+
return scheduler_factory_map.get(name, lambda: None)()
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def free_memory() -> None:
|
| 85 |
+
"""Free up GPU and system memory."""
|
| 86 |
+
if torch.cuda.is_available():
|
| 87 |
+
torch.cuda.empty_cache()
|
| 88 |
+
torch.cuda.ipc_collect()
|
| 89 |
+
gc.collect()
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
def preprocess_prompt(
|
| 93 |
+
positive: str,
|
| 94 |
+
negative: str = "",
|
| 95 |
+
add_style: bool = True,
|
| 96 |
+
) -> Tuple[str, str]:
|
| 97 |
+
formatted_positive = positive
|
| 98 |
+
|
| 99 |
+
combined_negative = ""
|
| 100 |
+
if negative.strip():
|
| 101 |
+
if combined_negative:
|
| 102 |
+
combined_negative += ", " + negative
|
| 103 |
+
else:
|
| 104 |
+
combined_negative = negative
|
| 105 |
+
|
| 106 |
+
return formatted_positive, combined_negative
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
def common_upscale(
|
| 110 |
+
samples: torch.Tensor,
|
| 111 |
+
width: int,
|
| 112 |
+
height: int,
|
| 113 |
+
upscale_method: str,
|
| 114 |
+
) -> torch.Tensor:
|
| 115 |
+
return torch.nn.functional.interpolate(
|
| 116 |
+
samples, size=(height, width), mode=upscale_method
|
| 117 |
+
)
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def upscale(
|
| 121 |
+
samples: torch.Tensor, upscale_method: str, scale_by: float
|
| 122 |
+
) -> torch.Tensor:
|
| 123 |
+
width = round(samples.shape[3] * scale_by)
|
| 124 |
+
height = round(samples.shape[2] * scale_by)
|
| 125 |
+
return common_upscale(samples, width, height, upscale_method)
|
| 126 |
+
|
| 127 |
+
|
| 128 |
+
def preprocess_image_dimensions(width, height):
|
| 129 |
+
if width % 8 != 0:
|
| 130 |
+
width = width - (width % 8)
|
| 131 |
+
if height % 8 != 0:
|
| 132 |
+
height = height - (height % 8)
|
| 133 |
+
return width, height
|
| 134 |
+
|
| 135 |
+
|
| 136 |
+
def save_image(image, metadata, output_dir, is_colab):
|
| 137 |
+
if is_colab:
|
| 138 |
+
current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
|
| 139 |
+
filename = f"image_{current_time}.png"
|
| 140 |
+
else:
|
| 141 |
+
filename = str(uuid.uuid4()) + ".png"
|
| 142 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 143 |
+
filepath = os.path.join(output_dir, filename)
|
| 144 |
+
metadata_str = json.dumps(metadata)
|
| 145 |
+
info = PngImagePlugin.PngInfo()
|
| 146 |
+
info.add_text("parameters", metadata_str)
|
| 147 |
+
image.save(filepath, "PNG", pnginfo=info)
|
| 148 |
+
return filepath
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def is_google_colab():
|
| 152 |
+
try:
|
| 153 |
+
import google.colab
|
| 154 |
+
return True
|
| 155 |
+
except:
|
| 156 |
+
return False
|
| 157 |
+
|
| 158 |
+
|
| 159 |
+
def load_pipeline(model_name: str, device: torch.device, hf_token: Optional[str] = None, vae: Optional[AutoencoderKL] = None) -> Any:
|
| 160 |
+
"""Load the Stable Diffusion pipeline."""
|
| 161 |
+
try:
|
| 162 |
+
pipeline = (
|
| 163 |
+
SDXLLongPromptWeightingPipeline.from_single_file
|
| 164 |
+
if model_name.endswith(".safetensors")
|
| 165 |
+
else SDXLLongPromptWeightingPipeline.from_pretrained
|
| 166 |
+
)
|
| 167 |
+
|
| 168 |
+
pipe = pipeline(
|
| 169 |
+
model_name,
|
| 170 |
+
vae=vae,
|
| 171 |
+
torch_dtype=torch.float16,
|
| 172 |
+
use_safetensors=True,
|
| 173 |
+
add_watermarker=False
|
| 174 |
+
)
|
| 175 |
+
pipe.to(device)
|
| 176 |
+
return pipe
|
| 177 |
+
except Exception as e:
|
| 178 |
+
logging.error(f"Failed to load pipeline: {str(e)}", exc_info=True)
|
| 179 |
+
raise
|