Fooocus Prompt Expansion (#329)

* add vae approx download

* files

* files

* files

* i

* i

* i

* i

* i

* i

* i

* i

* i

* i
This commit is contained in:
lllyasviel
2023-09-09 17:22:32 -07:00
committed by GitHub
parent 09e0d1cb3a
commit 496766edd7
19 changed files with 150532 additions and 81 deletions
+73 -21
View File
@@ -15,7 +15,7 @@ def worker():
import modules.path
import modules.patch
from modules.sdxl_styles import apply_style, aspect_ratios
from modules.sdxl_styles import apply_style_negative, apply_style_positive, aspect_ratios
from modules.private_logger import log
try:
@@ -29,19 +29,69 @@ def worker():
def handler(task):
prompt, negative_prompt, style_selction, performance_selction, \
aspect_ratios_selction, image_number, image_seed, sharpness, base_model_name, refiner_model_name, \
aspect_ratios_selction, image_number, image_seed, sharpness, raw_mode, \
base_model_name, refiner_model_name, \
l1, w1, l2, w2, l3, w3, l4, w4, l5, w5 = task
loras = [(l1, w1), (l2, w2), (l3, w3), (l4, w4), (l5, w5)]
modules.patch.sharpness = sharpness
outputs.append(['preview', (1, 'Initializing ...', None)])
seed = image_seed
max_seed = int(1024 * 1024 * 1024)
if not isinstance(seed, int):
seed = random.randint(1, max_seed)
if seed < 0:
seed = - seed
seed = seed % max_seed
outputs.append(['preview', (3, 'Load models ...', None)])
pipeline.refresh_base_model(base_model_name)
pipeline.refresh_refiner_model(refiner_model_name)
pipeline.refresh_loras(loras)
pipeline.clean_prompt_cond_caches()
p_txt, n_txt = apply_style(style_selction, prompt, negative_prompt)
outputs.append(['preview', (5, 'Encoding negative text ...', None)])
n_txt = apply_style_negative(style_selction, negative_prompt)
n_cond = pipeline.process_prompt(n_txt)
tasks = []
if raw_mode:
outputs.append(['preview', (9, 'Encoding positive text ...', None)])
p_txt = apply_style_positive(style_selction, prompt)
p_cond = pipeline.process_prompt(p_txt)
for i in range(image_number):
tasks.append(dict(
prompt=prompt,
negative_prompt=negative_prompt,
seed=seed + i,
n_cond=n_cond,
p_cond=p_cond,
real_positive_prompt=p_txt,
real_negative_prompt=n_txt
))
else:
for i in range(image_number):
outputs.append(['preview', (9, f'Preparing positive text #{i + 1} ...', None)])
current_seed = seed + i
p_txt = pipeline.expand_txt(prompt, current_seed)
print(f'Expanded positive prompt: {p_txt}')
p_txt = apply_style_positive(style_selction, p_txt)
tasks.append(dict(
prompt=prompt,
negative_prompt=negative_prompt,
seed=current_seed,
n_cond=n_cond,
real_positive_prompt=p_txt,
real_negative_prompt=n_txt
))
for i, t in enumerate(tasks):
outputs.append(['preview', (12, f'Encoding positive text #{i + 1} ...', None)])
t['p_cond'] = pipeline.process_prompt(t['real_positive_prompt'])
if performance_selction == 'Speed':
steps = 30
@@ -53,45 +103,47 @@ def worker():
width, height = aspect_ratios[aspect_ratios_selction]
results = []
seed = image_seed
max_seed = int(1024*1024*1024)
if not isinstance(seed, int):
seed = random.randint(1, max_seed)
if seed < 0:
seed = - seed
seed = seed % max_seed
all_steps = steps * image_number
def callback(step, x0, x, total_steps, y):
done_steps = i * steps + step
done_steps = current_task_id * steps + step
outputs.append(['preview', (
int(100.0 * float(done_steps) / float(all_steps)),
int(15.0 + 85.0 * float(done_steps) / float(all_steps)),
f'Step {step}/{total_steps} in the {i}-th Sampling',
y)])
for i in range(image_number):
imgs = pipeline.process(p_txt, n_txt, steps, switch, width, height, seed, callback=callback)
outputs.append(['preview', (13, 'Starting tasks ...', None)])
for current_task_id, task in enumerate(tasks):
imgs = pipeline.process_diffusion(
positive_cond=task['p_cond'],
negative_cond=task['n_cond'],
steps=steps,
switch=switch,
width=width,
height=height,
image_seed=task['seed'],
callback=callback)
for x in imgs:
d = [
('Prompt', prompt),
('Negative Prompt', negative_prompt),
('Prompt', task['prompt']),
('Negative Prompt', task['negative_prompt']),
('Real Positive Prompt', task['real_positive_prompt']),
('Real Negative Prompt', task['real_negative_prompt']),
('Raw Mode', str(raw_mode)),
('Style', style_selction),
('Performance', performance_selction),
('Resolution', str((width, height))),
('Sharpness', sharpness),
('Base Model', base_model_name),
('Refiner Model', refiner_model_name),
('Seed', seed)
('Seed', task['seed'])
]
for n, w in loras:
if n != 'None':
d.append((f'LoRA [{n}] weight', w))
log(x, d)
seed += 1
results += imgs
outputs.append(['results', results])
+20 -38
View File
@@ -5,6 +5,7 @@ import modules.path
from comfy.model_base import SDXL, SDXLRefiner
from modules.patch import cfg_patched
from modules.expansion import FooocusExpansion
xl_base: core.StableDiffusionModel = None
@@ -43,7 +44,6 @@ def refresh_base_model(name):
xl_base_patched = xl_base
xl_base_patched_hash = ''
print(f'Base model loaded: {xl_base_hash}')
return
@@ -103,27 +103,24 @@ refresh_base_model(modules.path.default_base_model_name)
refresh_refiner_model(modules.path.default_refiner_model_name)
refresh_loras([(modules.path.default_lora_name, 0.5), ('None', 0.5), ('None', 0.5), ('None', 0.5), ('None', 0.5)])
positive_conditions_cache = None
negative_conditions_cache = None
positive_conditions_refiner_cache = None
negative_conditions_refiner_cache = None
expansion_model = FooocusExpansion()
def clean_prompt_cond_caches():
global positive_conditions_cache, negative_conditions_cache, \
positive_conditions_refiner_cache, negative_conditions_refiner_cache
positive_conditions_cache = None
negative_conditions_cache = None
positive_conditions_refiner_cache = None
negative_conditions_refiner_cache = None
return
def expand_txt(*args, **kwargs):
return expansion_model(*args, **kwargs)
def process_prompt(text):
base_cond = core.encode_prompt_condition(clip=xl_base_patched.clip, prompt=text)
if xl_refiner is not None:
refiner_cond = core.encode_prompt_condition(clip=xl_refiner.clip, prompt=text)
else:
refiner_cond = None
return base_cond, refiner_cond
@torch.no_grad()
def process(positive_prompt, negative_prompt, steps, switch, width, height, image_seed, callback):
global positive_conditions_cache, negative_conditions_cache, \
positive_conditions_refiner_cache, negative_conditions_refiner_cache
def process_diffusion(positive_cond, negative_cond, steps, switch, width, height, image_seed, callback):
if xl_base is not None:
xl_base.unet.model_options['sampler_cfg_function'] = cfg_patched
@@ -133,40 +130,27 @@ def process(positive_prompt, negative_prompt, steps, switch, width, height, imag
if xl_refiner is not None:
xl_refiner.unet.model_options['sampler_cfg_function'] = cfg_patched
positive_conditions = core.encode_prompt_condition(clip=xl_base_patched.clip, prompt=positive_prompt) if positive_conditions_cache is None else positive_conditions_cache
negative_conditions = core.encode_prompt_condition(clip=xl_base_patched.clip, prompt=negative_prompt) if negative_conditions_cache is None else negative_conditions_cache
positive_conditions_cache = positive_conditions
negative_conditions_cache = negative_conditions
empty_latent = core.generate_empty_latent(width=width, height=height, batch_size=1)
if xl_refiner is not None:
positive_conditions_refiner = core.encode_prompt_condition(clip=xl_refiner.clip, prompt=positive_prompt) if positive_conditions_refiner_cache is None else positive_conditions_refiner_cache
negative_conditions_refiner = core.encode_prompt_condition(clip=xl_refiner.clip, prompt=negative_prompt) if negative_conditions_refiner_cache is None else negative_conditions_refiner_cache
positive_conditions_refiner_cache = positive_conditions_refiner
negative_conditions_refiner_cache = negative_conditions_refiner
sampled_latent = core.ksampler_with_refiner(
model=xl_base_patched.unet,
positive=positive_conditions,
negative=negative_conditions,
positive=positive_cond[0],
negative=negative_cond[0],
refiner=xl_refiner.unet,
refiner_positive=positive_conditions_refiner,
refiner_negative=negative_conditions_refiner,
refiner_positive=positive_cond[1],
refiner_negative=negative_cond[1],
refiner_switch_step=switch,
latent=empty_latent,
steps=steps, start_step=0, last_step=steps, disable_noise=False, force_full_denoise=True,
seed=image_seed,
callback_function=callback
)
else:
sampled_latent = core.ksampler(
model=xl_base_patched.unet,
positive=positive_conditions,
negative=negative_conditions,
positive=positive_cond[0],
negative=negative_cond[0],
latent=empty_latent,
steps=steps, start_step=0, last_step=steps, disable_noise=False, force_full_denoise=True,
seed=image_seed,
@@ -174,7 +158,5 @@ def process(positive_prompt, negative_prompt, steps, switch, width, height, imag
)
decoded_latent = core.decode_vae(vae=xl_base_patched.vae, latent_image=sampled_latent)
images = core.image_to_numpy(decoded_latent)
return images
+23
View File
@@ -0,0 +1,23 @@
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, set_seed
from modules.path import fooocus_expansion_path
class FooocusExpansion:
def __init__(self):
self.tokenizer = AutoTokenizer.from_pretrained(fooocus_expansion_path)
self.model = AutoModelForCausalLM.from_pretrained(fooocus_expansion_path)
self.pipe = pipeline('text-generation',
model=self.model,
tokenizer=self.tokenizer,
device='cpu',
torch_dtype=torch.float32)
print('Fooocus Expansion engine loaded.')
def __call__(self, prompt, seed):
prompt = str(prompt).rstrip('\n')
seed = int(seed)
set_seed(seed)
response = self.pipe(prompt, max_length=len(prompt) + 256)
result = response[0]['generated_text'].rstrip('\n')
return result
+4
View File
@@ -2,8 +2,12 @@ import os
modelfile_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../models/checkpoints/'))
lorafile_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../models/loras/'))
vae_approx_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../models/vae_approx/'))
temp_outputs_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../outputs/'))
fooocus_expansion_path = os.path.abspath(os.path.join(os.path.dirname(__file__),
'../models/prompt_expansion/fooocus_expansion'))
os.makedirs(temp_outputs_path, exist_ok=True)
default_base_model_name = 'sd_xl_base_1.0_0.9vae.safetensors'
+1 -1
View File
@@ -22,7 +22,7 @@ def log(img, dic):
f.write(f"<p>{only_name}</p>\n")
i = 0
for k, v in dic:
if i < 2:
if i < 4:
f.write(f"<p>{k}: <b>{v}</b> </p>\n")
else:
if i % 2 == 0:
+10 -2
View File
@@ -959,6 +959,14 @@ SD_XL_BASE_RATIOS = {
aspect_ratios = {str(v[0]) + '×' + str(v[1]): v for k, v in SD_XL_BASE_RATIOS.items()}
def apply_style(style, positive, negative):
def apply_style_positive(style, txt):
p, n = styles.get(style, default_style)
return p.replace('{prompt}', positive), n + ', ' + negative
return p.replace('{prompt}', txt)
def apply_style_negative(style, txt):
p, n = styles.get(style, default_style)
if n == '':
return txt
else:
return n + ', ' + txt