mirror of
https://github.com/lllyasviel/Fooocus.git
synced 2026-08-16 13:13:16 +02:00
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:
+73
-21
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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'
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user