# Fooocus GPT2 Expansion # Algorithm created by Lvmin Zhang at 2023, Stanford # Modified by power88 and GPT-4o for stable-diffusion-webui # If used inside Fooocus, any use is permitted. # If used outside Fooocus, only non-commercial use is permitted (CC-By NC 4.0). # This applies to the word list, vocab, model, and algorithm. import os import torch import math import gradio as gr import psutil from modules.scripts import basedir from transformers.generation.logits_process import LogitsProcessorList from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed from modules import scripts, shared, script_callbacks from modules.ui_components import FormRow, FormColumn, FormGroup, ToolButton def text_encoder_device(): if torch.cuda.is_available(): return torch.device(torch.cuda.current_device()) else: return torch.device("cpu") def text_encoder_offload_device(): if torch.cuda.is_available(): return torch.device(torch.cuda.current_device()) else: return torch.device("cpu") def get_free_memory(dev=None, torch_free_too=False): global directml_enabled if dev is None: dev = text_encoder_device() if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'): mem_free_total = psutil.virtual_memory().available mem_free_torch = mem_free_total else: if directml_enabled: mem_free_total = 1024 * 1024 * 1024 #TODO mem_free_torch = mem_free_total else: stats = torch.cuda.memory_stats(dev) mem_active = stats['active_bytes.all.current'] mem_reserved = stats['reserved_bytes.all.current'] mem_free_cuda, _ = torch.cuda.mem_get_info(dev) mem_free_torch = mem_reserved - mem_active mem_free_total = mem_free_cuda + mem_free_torch # limitation of np.random.seed(), called from transformers.set_seed() SEED_LIMIT_NUMPY = 2**32 neg_inf = - 8192.0 ext_dir = basedir() path_fooocus_expansion = os.path.join('.', "models", "prompt_expansion") def safe_str(x): x = str(x) for _ in range(16): x = x.replace(' ', ' ') return x.strip(",. \r\n") def remove_pattern(x, pattern): for p in pattern: x = x.replace(p, '') return x def should_use_fp16(device=None, model_params=0, prioritize_performance=True): if device is not None: if hasattr(device, 'type'): if device.type == 'cpu': return False return False if torch.cuda.is_bf16_supported(): return True props = torch.cuda.get_device_properties("cuda") if props.major < 6: return False fp16_works = False #FP16 is confirmed working on a 1080 (GP104) but it's a bit slower than FP32 so it should only be enabled #when the model doesn't actually fit on the card #TODO: actually test if GP106 and others have the same type of behavior nvidia_10_series = ["1080", "1070", "titan x", "p3000", "p3200", "p4000", "p4200", "p5000", "p5200", "p6000", "1060", "1050"] for x in nvidia_10_series: if x in props.name.lower(): fp16_works = True if fp16_works: free_model_memory = (get_free_memory() * 0.9 - (1024 * 1024 * 1024)) if (not prioritize_performance) or model_params * 4 > free_model_memory: return True if props.major < 7: return False #FP16 is just broken on these cards nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600", "MX550", "MX450", "CMP 30HX", "T2000", "T1000", "T1200"] for x in nvidia_16_series: if x in props.name: return False return True def is_device_mps(device): if hasattr(device, 'type'): if (device.type == 'mps'): return True return False class FooocusExpansion: def __init__(self): global load_model_device print(f'Loading models from {path_fooocus_expansion}') self.tokenizer = AutoTokenizer.from_pretrained(path_fooocus_expansion) positive_words = open(os.path.join(path_fooocus_expansion, 'positive.txt'), encoding='utf-8').read().splitlines() positive_words = ['Ġ' + x.lower() for x in positive_words if x != ''] self.logits_bias = torch.zeros((1, len(self.tokenizer.vocab)), dtype=torch.float32) + neg_inf debug_list = [] for k, v in self.tokenizer.vocab.items(): if k in positive_words: self.logits_bias[0, v] = 0 debug_list.append(k[1:]) print(f'Fooocus V2 Expansion: Vocab with {len(debug_list)} words.') self.model = AutoModelForCausalLM.from_pretrained(path_fooocus_expansion) self.model.eval() load_model_device = text_encoder_device() offload_device = text_encoder_offload_device() # MPS hack if is_device_mps(load_model_device): load_model_device = torch.device('cpu') offload_device = torch.device('cpu') use_fp16 = should_use_fp16(device=load_model_device) if use_fp16: self.model.half() self.model.to(load_model_device) # Ensure model is on the correct device print(f'Fooocus Expansion engine loaded for {load_model_device}, use_fp16 = {use_fp16}.') def unload_model(self): """Unload the model to free up memory.""" del self.model torch.cuda.empty_cache() print('Model unloaded and memory cleared.') @torch.no_grad() @torch.inference_mode() def logits_processor(self, input_ids, scores): assert scores.ndim == 2 and scores.shape[0] == 1 self.logits_bias = self.logits_bias.to(load_model_device) bias = self.logits_bias.clone().to(load_model_device) # Ensure bias is on the correct device bias[0, input_ids[0].to(load_model_device).long()] = neg_inf # Ensure input_ids are on the correct device bias[0, 11] = 0 return scores + bias.to(scores.device) # Ensure bias is on the same device as scores @torch.no_grad() @torch.inference_mode() def __call__(self, prompt, seed): if prompt == '': return '' seed = int(seed) % SEED_LIMIT_NUMPY set_seed(seed) prompt = safe_str(prompt) + ',' tokenized_kwargs = self.tokenizer(prompt, return_tensors="pt") tokenized_kwargs.data['input_ids'] = tokenized_kwargs.data['input_ids'].to(load_model_device) tokenized_kwargs.data['attention_mask'] = tokenized_kwargs.data['attention_mask'].to(load_model_device) current_token_length = int(tokenized_kwargs.data['input_ids'].shape[1]) max_token_length = 75 * int(math.ceil(float(current_token_length) / 75.0)) max_new_tokens = max_token_length - current_token_length features = self.model.generate(**tokenized_kwargs, top_k=100, max_new_tokens=max_new_tokens, do_sample=True, logits_processor=LogitsProcessorList([self.logits_processor])) response = self.tokenizer.batch_decode(features, skip_special_tokens=True) result = safe_str(response[0]) return result def createPositive(positive, seed): try: expansion = FooocusExpansion() positive = expansion(positive, seed=seed) expansion.unload_model() # Unload the model after use return positive except Exception as e: print(f"An error occurred: {str(e)}") class FooocusPromptExpansion(scripts.Script): def __init__(self) -> None: super().__init__() def title(self): return 'Fooocus Prompt Expansion' def show(self, is_img2img): return scripts.AlwaysVisible def ui(self, is_img2img): with gr.Group(): with gr.Accordion("Fooocus Expansion", open=True): is_enabled = gr.Checkbox( value=True, label="Enable Expansion", info="Enable Or Disable Expansion ") seed = gr.Number( value=0, maximum=63, label="Seed", info="Seed for random number generator") return [is_enabled, seed] def process(self, p, is_enabled, seed): if not is_enabled: return for i, prompt in enumerate(p.all_prompts): positivePrompt = createPositive(prompt, seed) p.all_prompts[i] = positivePrompt def after_component(self, component, **kwargs): # https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/7456#issuecomment-1414465888 helpfull link # Find the text2img textbox component if kwargs.get("elem_id") == "txt2img_prompt": # postive prompt textbox self.boxx = component # Find the img2img textbox component if kwargs.get("elem_id") == "img2img_prompt": # postive prompt textbox self.boxxIMG = component