Spaces:
Running
on
Zero
Running
on
Zero
File size: 7,099 Bytes
1a688bc cb5daed 4d6f2bc 7736f5f 48c31e7 4d6f2bc 5c4e8c1 4d6f2bc b7fd57e 4d6f2bc 48c31e7 dffd0bb 23f4f95 ca5a1e4 4d6f2bc cb5daed dffd0bb 1128e78 cb5daed 1128e78 53eff53 1a688bc 4d6f2bc b7fd57e 48c31e7 b7fd57e 4d6f2bc 48c31e7 4d6f2bc 48c31e7 4d6f2bc 48c31e7 4d6f2bc 053b3a4 1a688bc 053b3a4 1a688bc 053b3a4 1a688bc cb5daed 4d6f2bc 23f4f95 1a688bc 4d6f2bc 1a688bc 48c31e7 1128e78 1a688bc 4d6f2bc 1a688bc 48c31e7 c348e53 48c31e7 4d6f2bc b7fd57e 05246f1 5c4e8c1 4d6f2bc 1128e78 5c4e8c1 1128e78 1a688bc 48c31e7 22a0476 cb5daed 22a0476 48c31e7 22a0476 48c31e7 4d6f2bc 7736f5f 4d6f2bc 05246f1 cb5daed c348e53 cb5daed 05246f1 22a0476 cb5daed 4d6f2bc 23f4f95 4d6f2bc 48c31e7 22a0476 48c31e7 4d6f2bc 48c31e7 22a0476 4d6f2bc 48c31e7 dffd0bb 1a688bc dffd0bb 4d6f2bc 48c31e7 22a0476 1128e78 dffd0bb 1a688bc dffd0bb 4d6f2bc b7fd57e 22a0476 05246f1 22a0476 05246f1 22a0476 23f4f95 22a0476 4d6f2bc 1128e78 23f4f95 1128e78 cb5daed 05246f1 4d6f2bc |
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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 |
import json
import os
import re
import time
from contextlib import contextmanager
from datetime import datetime
from itertools import product
from typing import Callable
import spaces
import tomesd
import torch
from compel import Compel, DiffusersTextualInversionManager, ReturnedEmbeddingsType
from compel.prompt_parser import PromptParser
from huggingface_hub.utils import HFValidationError, RepositoryNotFoundError
from .loader import Loader
__import__("warnings").filterwarnings("ignore", category=FutureWarning, module="diffusers")
__import__("warnings").filterwarnings("ignore", category=FutureWarning, module="transformers")
__import__("transformers").logging.set_verbosity_error()
ZERO_GPU = (
os.environ.get("SPACES_ZERO_GPU", "").lower() == "true"
or os.environ.get("SPACES_ZERO_GPU", "") == "1"
)
with open("./data/styles.json") as f:
styles = json.load(f)
# applies tome to the pipeline
@contextmanager
def token_merging(pipe, tome_ratio=0):
try:
if tome_ratio > 0:
tomesd.apply_patch(pipe, max_downsample=1, sx=2, sy=2, ratio=tome_ratio)
yield
finally:
tomesd.remove_patch(pipe) # idempotent
# parse prompts with arrays
def parse_prompt(prompt: str) -> list[str]:
arrays = re.findall(r"\[\[(.*?)\]\]", prompt)
if not arrays:
return [prompt]
tokens = [item.split(",") for item in arrays]
combinations = list(product(*tokens))
prompts = []
for combo in combinations:
current_prompt = prompt
for i, token in enumerate(combo):
current_prompt = current_prompt.replace(f"[[{arrays[i]}]]", token.strip(), 1)
prompts.append(current_prompt)
return prompts
def apply_style(prompt, style_id, negative=False):
global styles
if not style_id or style_id == "None":
return prompt
for style in styles:
if style["id"] == style_id:
if negative:
return prompt + " . " + style["negative_prompt"]
else:
return style["prompt"].format(prompt=prompt)
return prompt
@spaces.GPU(duration=40)
def generate(
positive_prompt,
negative_prompt="",
embeddings=[],
style=None,
seed=None,
model="runwayml/stable-diffusion-v1-5",
scheduler="PNDM",
width=512,
height=512,
guidance_scale=7.5,
inference_steps=50,
num_images=1,
karras=False,
taesd=False,
freeu=False,
clip_skip=False,
truncate_prompts=False,
increment_seed=True,
deepcache_interval=1,
tome_ratio=0,
scale=1,
Info: Callable[[str], None] = None,
Error=Exception,
):
if not torch.cuda.is_available():
raise Error("CUDA not available")
# https://pytorch.org/docs/stable/generated/torch.manual_seed.html
if seed is None or seed < 0:
seed = int(datetime.now().timestamp() * 1_000_000) % (2**64)
DEVICE = torch.device("cuda")
DTYPE = (
torch.bfloat16
if torch.cuda.is_available() and torch.cuda.get_device_properties(DEVICE).major >= 8
else torch.float16
)
EMBEDDINGS_TYPE = (
ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NORMALIZED
if clip_skip
else ReturnedEmbeddingsType.LAST_HIDDEN_STATES_NORMALIZED
)
with torch.inference_mode():
start = time.perf_counter()
loader = Loader()
pipe, upscaler = loader.load(
model,
scheduler,
karras,
taesd,
freeu,
deepcache_interval,
scale,
DTYPE,
DEVICE,
)
# load embeddings and append to negative prompt
embeddings_dir = os.path.join(os.path.dirname(__file__), "..", "embeddings")
embeddings_dir = os.path.abspath(embeddings_dir)
for embedding in embeddings:
try:
pipe.load_textual_inversion(
pretrained_model_name_or_path=f"{embeddings_dir}/{embedding}.pt",
token=f"<{embedding}>",
)
negative_prompt = (
f"{negative_prompt}, {embedding}" if negative_prompt else embedding
)
except (EnvironmentError, HFValidationError, RepositoryNotFoundError):
raise Error(f"Invalid embedding: {embedding}")
# prompt embeds
compel = Compel(
textual_inversion_manager=DiffusersTextualInversionManager(pipe),
dtype_for_device_getter=lambda _: DTYPE,
returned_embeddings_type=EMBEDDINGS_TYPE,
truncate_long_prompts=truncate_prompts,
text_encoder=pipe.text_encoder,
tokenizer=pipe.tokenizer,
device=pipe.device,
)
images = []
current_seed = seed
try:
styled_negative_prompt = apply_style(negative_prompt, style, negative=True)
neg_embeds = compel(styled_negative_prompt)
except PromptParser.ParsingException:
raise Error("ParsingException: Invalid negative prompt")
for i in range(num_images):
# seeded generator for each iteration
generator = torch.Generator(device=pipe.device).manual_seed(current_seed)
try:
all_positive_prompts = parse_prompt(positive_prompt)
prompt_index = i % len(all_positive_prompts)
pos_prompt = all_positive_prompts[prompt_index]
styled_pos_prompt = apply_style(pos_prompt, style)
pos_embeds = compel(styled_pos_prompt)
pos_embeds, neg_embeds = compel.pad_conditioning_tensors_to_same_length(
[pos_embeds, neg_embeds]
)
except PromptParser.ParsingException:
raise Error("ParsingException: Invalid prompt")
with token_merging(pipe, tome_ratio=tome_ratio):
try:
image = pipe(
output_type="np" if scale > 1 else "pil",
num_inference_steps=inference_steps,
negative_prompt_embeds=neg_embeds,
guidance_scale=guidance_scale,
prompt_embeds=pos_embeds,
generator=generator,
height=height,
width=width,
).images[0]
if scale > 1:
image = upscaler.predict(image)
images.append((image, str(current_seed)))
finally:
if not ZERO_GPU:
pipe.unload_textual_inversion()
torch.cuda.empty_cache()
if increment_seed:
current_seed += 1
if ZERO_GPU:
# spaces always start fresh
loader.pipe = None
loader.upscaler = None
diff = time.perf_counter() - start
if Info:
Info(f"Generated {len(images)} image{'s' if len(images) > 1 else ''} in {diff:.2f}s")
return images
|