import concurrent.futures import random import gradio as gr import requests import io, base64, json, os import spaces from PIL import Image from .models import IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, VIDEO_GENERATION_MODELS, MUSEUM_UNSUPPORTED_MODELS, DESIRED_APPEAR_MODEL, load_pipeline from .fetch_museum_results import draw_from_imagen_museum, draw2_from_imagen_museum, draw_from_videogen_museum, draw2_from_videogen_museum from .pre_download import pre_download_all_models, pre_download_video_models from transformers import AutoTokenizer, AutoModelForCausalLM import torch import re def debug_packages(): import pkg_resources installed_packages = pkg_resources.working_set for package in installed_packages: print(f"{package.key}=={package.version}") def fetch_unsafe_words(file_path): """ Loads unsafe words from a file and returns them as a list. """ try: with open(file_path, 'r') as file: # Read lines from file and strip any extra whitespace unsafe_words = [line.strip() for line in file.readlines()] # Remove any empty strings that may result from empty lines unsafe_words = [word for word in unsafe_words if word] return unsafe_words except Exception as e: print(f"Error loading file: {e}. Using default unsafe words.") # Default unsafe words list return [ "anal", "anus", "arse", "ass", "ballsack", "bastard", "bdsm", "bitch", "bimbo", "blow job", "blowjob", "blue waffle", "boob", "booobs", "breasts", "booty call", "boner", "bondage", "bullshit", "busty", "butthole", "cawk", "chink", "clit", "cnut", "cock", "cokmuncher", "cowgirl", "crap", "crotch", "cum", "cunt", "damn", "dick", "dildo", "dink", "deepthroat", "deep throat", "dog style", "doggie style", "doggy style", "doosh", "douche", "duche", "ejaculate", "ejaculating", "ejaculation", "ejakulate", "erotic", "erotism", "fag", "fatass", "femdom", "fingering", "footjob", "foot job", "fuck", "fcuk", "fingerfuck", "fistfuck", "fook", "fooker", "fuk", "gangbang", "gang bang", "gaysex", "handjob", "hand job", "hentai", "hooker", "hoer", "homo", "horny", "incest", "jackoff", "jack off", "jerkoff", "jerk off", "jizz", "masturbate", "mofo", "mothafuck", "motherfuck", "milf", "muff", "nigga", "nigger", "nipple", "nob", "numbnuts", "nutsack", "nude", "orgy", "orgasm", "panty", "panties", "penis", "playboy", "porn", "pussy", "pussies", "rape", "raping", "rapist", "rectum", "retard", "rimming", "sadist", "sadism", "scrotum", "sex", "semen", "shemale", "she male", "shit", "slut", "spunk", "strip club", "stripclub", "tit", "threesome", "three some", "throating", "twat", "viagra", "vagina", "wank", "whore", "whoar", "xxx" ] def check_prompt_safety(prompt, unsafe_words_file='./profanity_words.txt'): """ Checking prompt safety. Returns boolean (Not Safe = False, Safe = True) """ # Load unsafe words from the provided file or use default if loading fails unsafe_words = fetch_unsafe_words(unsafe_words_file) # Convert input string to lowercase to ensure case-insensitive matching prompt = prompt.lower() # Check if any unsafe word is in the input string for word in unsafe_words: # Use regex to match whole words only if re.search(r'\b' + re.escape(word) + r'\b', prompt): return False return True class ModelManager: def __init__(self, enable_nsfw=False, do_pre_download=False, do_debug_packages=False): self.model_ig_list = IMAGE_GENERATION_MODELS self.model_ie_list = IMAGE_EDITION_MODELS self.model_vg_list = VIDEO_GENERATION_MODELS self.excluding_model_list = MUSEUM_UNSUPPORTED_MODELS self.desired_model_list = DESIRED_APPEAR_MODEL self.enable_nsfw = enable_nsfw self.load_guard(enable_nsfw) self.loaded_models = {} if do_pre_download: pre_download_all_models() if do_debug_packages: debug_packages() def load_model_pipe(self, model_name): if not model_name in self.loaded_models: pipe = load_pipeline(model_name) self.loaded_models[model_name] = pipe else: pipe = self.loaded_models[model_name] return pipe def load_guard(self, enable_nsfw=True): model_id = "meta-llama/Llama-Guard-3-8B" device = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.bfloat16 token = os.getenv("HF_TOKEN") or os.getenv("HF_GUARD") if enable_nsfw: self.guard_tokenizer = AutoTokenizer.from_pretrained(model_id) self.guard = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, device_map=device) else: self.guard_tokenizer = None self.guard = None def NSFW_filter_simple(self, prompt): is_safe = check_prompt_safety(prompt) if is_safe: return "safe" else: return "unsafe" def NSFW_filter(self, prompt): chat = [{"role": "user", "content": prompt}] input_ids = self.guard_tokenizer.apply_chat_template(chat, return_tensors="pt").to('cuda') self.guard.cuda() if self.guard: @spaces.GPU(duration=30) def _generate(): return self.guard.generate(input_ids=input_ids, max_new_tokens=100, pad_token_id=0) output = _generate() output = self.guard.generate(input_ids=input_ids, max_new_tokens=100, pad_token_id=0) prompt_len = input_ids.shape[-1] result = self.guard_tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True) return result else: # guard is disabled return "safe" @spaces.GPU(duration=120) def generate_image_ig(self, prompt, model_name): # if 'unsafe' not in self.NSFW_filter(prompt): print('The prompt is safe') pipe = self.load_model_pipe(model_name) result = pipe(prompt=prompt) # else: # print(f'The prompt "{prompt}" is not safe') # result = '' return result def generate_image_ig_api(self, prompt, model_name): # if 'unsafe' not in self.NSFW_filter(prompt): print('The prompt is safe') pipe = self.load_model_pipe(model_name) result = pipe(prompt=prompt) # else: # print(f'The prompt "{prompt}" is not safe') # result = '' return result def generate_image_ig_museum(self, model_name): model_name = model_name.split('_')[1] result_list = draw_from_imagen_museum("t2i", model_name) image_link = result_list[0] prompt = result_list[1] return image_link, prompt def generate_image_ig_parallel_anony(self, prompt, model_A, model_B): # Using list comprehension to get the difference between two lists picking_list = [item for item in self.model_ig_list if item not in self.excluding_model_list] if model_A == "" and model_B == "": model_names = random.sample([model for model in picking_list], 2) else: model_names = [model_A, model_B] with concurrent.futures.ThreadPoolExecutor() as executor: futures = [executor.submit(self.generate_image_ig, prompt, model) if model.startswith("imagenhub") else executor.submit(self.generate_image_ig_api, prompt, model) for model in model_names] results = [future.result() for future in futures] return results[0], results[1], model_names[0], model_names[1] def generate_image_ig_museum_parallel_anony(self, model_A, model_B): # Using list comprehension to get the difference between two lists picking_list = [item for item in self.model_ig_list if item not in self.excluding_model_list] if model_A == "" and model_B == "": model_names = random.sample([model for model in picking_list], 2) else: model_names = [model_A, model_B] with concurrent.futures.ThreadPoolExecutor() as executor: model_1 = model_names[0].split('_')[1] model_2 = model_names[1].split('_')[1] result_list = draw2_from_imagen_museum("t2i", model_1, model_2) image_links = result_list[0] prompt_list = result_list[1] return image_links[0], image_links[1], model_names[0], model_names[1], prompt_list[0] def generate_image_ig_parallel(self, prompt, model_A, model_B): model_names = [model_A, model_B] with concurrent.futures.ThreadPoolExecutor() as executor: futures = [executor.submit(self.generate_image_ig, prompt, model) if model.startswith("imagenhub") else executor.submit(self.generate_image_ig_api, prompt, model) for model in model_names] results = [future.result() for future in futures] return results[0], results[1] def generate_image_ig_museum_parallel(self, model_A, model_B): with concurrent.futures.ThreadPoolExecutor() as executor: model_1 = model_A.split('_')[1] model_2 = model_B.split('_')[1] result_list = draw2_from_imagen_museum("t2i", model_1, model_2) image_links = result_list[0] prompt_list = result_list[1] return image_links[0], image_links[1], prompt_list[0] @spaces.GPU(duration=200) def generate_image_ie(self, textbox_source, textbox_target, textbox_instruct, source_image, model_name): # if 'unsafe' not in self.NSFW_filter(" ".join([textbox_source, textbox_target, textbox_instruct])): pipe = self.load_model_pipe(model_name) result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct) # else: # result = '' return result def generate_image_ie_museum(self, model_name): model_name = model_name.split('_')[1] result_list = draw_from_imagen_museum("tie", model_name) image_links = result_list[0] prompt_list = result_list[1] # image_links = [src, model] # prompt_list = [source_caption, target_caption, instruction] return image_links[0], image_links[1], prompt_list[0], prompt_list[1], prompt_list[2] def generate_image_ie_parallel(self, textbox_source, textbox_target, textbox_instruct, source_image, model_A, model_B): model_names = [model_A, model_B] with concurrent.futures.ThreadPoolExecutor() as executor: futures = [ executor.submit(self.generate_image_ie, textbox_source, textbox_target, textbox_instruct, source_image, model) for model in model_names] results = [future.result() for future in futures] return results[0], results[1] def generate_image_ie_museum_parallel(self, model_A, model_B): model_names = [model_A, model_B] with concurrent.futures.ThreadPoolExecutor() as executor: model_1 = model_names[0].split('_')[1] model_2 = model_names[1].split('_')[1] result_list = draw2_from_imagen_museum("tie", model_1, model_2) image_links = result_list[0] prompt_list = result_list[1] # image_links = [src, model_A, model_B] # prompt_list = [source_caption, target_caption, instruction] return image_links[0], image_links[1], image_links[2], prompt_list[0], prompt_list[1], prompt_list[2] def generate_image_ie_parallel_anony(self, textbox_source, textbox_target, textbox_instruct, source_image, model_A, model_B): # Using list comprehension to get the difference between two lists picking_list = [item for item in self.model_ie_list if item not in self.excluding_model_list] if model_A == "" and model_B == "": model_names = random.sample([model for model in picking_list], 2) else: model_names = [model_A, model_B] with concurrent.futures.ThreadPoolExecutor() as executor: futures = [executor.submit(self.generate_image_ie, textbox_source, textbox_target, textbox_instruct, source_image, model) for model in model_names] results = [future.result() for future in futures] return results[0], results[1], model_names[0], model_names[1] def generate_image_ie_museum_parallel_anony(self, model_A, model_B): # Using list comprehension to get the difference between two lists picking_list = [item for item in self.model_ie_list if item not in self.excluding_model_list] if model_A == "" and model_B == "": model_names = random.sample([model for model in picking_list], 2) else: model_names = [model_A, model_B] with concurrent.futures.ThreadPoolExecutor() as executor: model_1 = model_names[0].split('_')[1] model_2 = model_names[1].split('_')[1] result_list = draw2_from_imagen_museum("tie", model_1, model_2) image_links = result_list[0] prompt_list = result_list[1] # image_links = [src, model_A, model_B] # prompt_list = [source_caption, target_caption, instruction] return image_links[0], image_links[1], image_links[2], prompt_list[0], prompt_list[1], prompt_list[2], model_names[0], model_names[1] @spaces.GPU(duration=150) def generate_video_vg(self, prompt, model_name): # if 'unsafe' not in self.NSFW_filter(prompt): pipe = self.load_model_pipe(model_name) result = pipe(prompt=prompt) # else: # result = '' return result def generate_video_vg_api(self, prompt, model_name): # if 'unsafe' not in self.NSFW_filter(prompt): pipe = self.load_model_pipe(model_name) result = pipe(prompt=prompt) # else: # result = '' return result def generate_video_vg_museum(self, model_name): model_name = model_name.split('_')[1] result_list = draw_from_videogen_museum("t2v", model_name) video_link = result_list[0] prompt = result_list[1] return video_link, prompt def generate_video_vg_parallel_anony(self, prompt, model_A, model_B): # Using list comprehension to get the difference between two lists picking_list = [item for item in self.model_vg_list if item not in self.excluding_model_list] if model_A == "" and model_B == "": model_names = random.sample([model for model in picking_list], 2) else: model_names = [model_A, model_B] with concurrent.futures.ThreadPoolExecutor() as executor: futures = [executor.submit(self.generate_video_vg, prompt, model) if model.startswith("videogenhub") else executor.submit(self.generate_video_vg_api, prompt, model) for model in model_names] results = [future.result() for future in futures] return results[0], results[1], model_names[0], model_names[1] def generate_video_vg_museum_parallel_anony(self, model_A, model_B): # Using list comprehension to get the difference between two lists picking_list = [item for item in self.model_vg_list if item not in self.excluding_model_list] #picking_list = [item for item in picking_list if item not in self.desired_model_list] if model_A == "" and model_B == "": model_names = random.sample([model for model in picking_list], 2) #override the random selection #model_names[random.choice([0, 1])] = random.choice(self.desired_model_list) else: model_names = [model_A, model_B] with concurrent.futures.ThreadPoolExecutor() as executor: model_1 = model_names[0].split('_')[1] model_2 = model_names[1].split('_')[1] result_list = draw2_from_videogen_museum("t2v", model_1, model_2) video_links = result_list[0] prompt_list = result_list[1] return video_links[0], video_links[1], model_names[0], model_names[1], prompt_list[0] def generate_video_vg_parallel(self, prompt, model_A, model_B): model_names = [model_A, model_B] with concurrent.futures.ThreadPoolExecutor() as executor: futures = [executor.submit(self.generate_video_vg, prompt, model) if model.startswith("videogenhub") else executor.submit(self.generate_video_vg_api, prompt, model) for model in model_names] results = [future.result() for future in futures] return results[0], results[1] def generate_video_vg_museum_parallel(self, model_A, model_B): model_names = [model_A, model_B] with concurrent.futures.ThreadPoolExecutor() as executor: model_1 = model_A.split('_')[1] model_2 = model_B.split('_')[1] result_list = draw2_from_videogen_museum("t2v", model_1, model_2) video_links = result_list[0] prompt_list = result_list[1] return video_links[0], video_links[1], prompt_list[0]