remove space from requirements.txt

app_imagegen_new.py

@@ -0,0 +1,863 @@
+import spaces
+import torch
+from diffusers import UNet2DConditionModel, DDIMScheduler, StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler, DPMSolverSDEScheduler, DPMSolverMultistepScheduler, AutoencoderKL, AutoencoderTiny, StableDiffusionXLImg2ImgPipeline
+import ipown
+from huggingface_hub import hf_hub_download
+from insightface.app import FaceAnalysis
+import gradio as gr
+import cv2
+import helper
+import random
+from transformers import Qwen2ForSequenceClassification, AutoTokenizer
+from PIL import Image
+import base64
+import json
+import time
+import jwt
+import glob
+import uuid
+import re
+import copy
+from model import LoraHelper, ModelManager, Model
+import numpy as np
+from compel import Compel, ReturnedEmbeddingsType
+from hidiffusion import apply_hidiffusion, remove_hidiffusion
+import utils
+from datetime import datetime,timezone
+import os
+from auth import AuthHelper
+import download_for_imagegen
+from safety import AgePredictor, NSFWClassifier, Qwen3Analyzer
+from utils import generate_watermark, save_image
+
+MAX_SEED = np.iinfo(np.int32).max
+config = json.load(open("./config.json", "r"))
+
+local_config = {}
+try:
+    local_config = json.load(open("./local_config.json", "r"))
+except Exception as e:
+    print(e)
+
+local_config.update({
+    "lazy_load": os.environ.get("LAZY_LOAD", "False").lower() == "true",  # Convert to boolean
+    "max_loaded_models": int(os.environ.get("MAX_LOADED_MODELS", 1)),  # Convert to integer
+})
+
+
+
+base_model_path = "John6666/pony-realism-v21main-sdxl"
+config["r2"]["access_key"] = os.getenv("R2_ACCESS_KEY")
+config["r2"]["secret_key"] = os.getenv("R2_SECRET_KEY")
+
+#
+characters = utils.RemoteJson(config["characters"],1800, lambda lst: {item['key']: item for item in lst})
+vibes = {item['key']: item for item in json.load(open("./vibe.json", "r"))}#utils.RemoteJson(config["vibes"],1800, lambda lst: {item['key']: item for item in lst})
+styles = utils.RemoteJson(config["styles"],1800, lambda lst: {item['key']: item for item in lst})
+#ip_xl_ckpt = hf_hub_download(repo_id="h94/IP-Adapter-FaceID", filename="ip-adapter-faceid_sdxl.bin", repo_type="model")
+#_ = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+
+# download weights
+if not os.path.exists('realesr-general-x4v3.pth'):
+    os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
+
+device = "cuda"
+scheduler_config = {
+    "num_train_timesteps": 1000,
+    "beta_start": 0.00085,
+    "beta_end": 0.012,
+    "beta_schedule": "scaled_linear",
+    "set_alpha_to_one": False,
+    "steps_offset": 1,
+    "prediction_type": "epsilon",
+}
+
+noise_scheduler = DDIMScheduler(
+    num_train_timesteps=1000,
+    #beta_start=0.00015,
+    #beta_end=0.02,
+    beta_start=0.00085,
+    beta_end=0.012,
+    beta_schedule="scaled_linear",
+    clip_sample=False,
+    set_alpha_to_one=False,
+    steps_offset=1,
+)
+euler = EulerAncestralDiscreteScheduler(
+        num_train_timesteps = 1000,
+        beta_start = 0.00085,
+        beta_end = 0.012,
+        #beta_schedule="scaled_linear",
+        #device = 'cuda',
+        steps_offset = 1,
+)
+
+age_predictor = AgePredictor()
+#nsfw_classifier = NSFWClassifier()
+llm_analyzer = Qwen3Analyzer()
+lora_mgr = LoraHelper()
+# vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)
+#pipe.scheduler = euler
+
+class PromptGenerator:
+# Define the regular expression
+    child_related_regex = re.compile(
+        r'(child|children|kid|kids|baby|shota|loli|lolicon|babies|ll*oo*ll*ii*|miniature|toddler|infant|juvenile|minor|underage|preteen|adolescent|youngster|youth|kindergarten|preschool|young girl|young daughter|'
+        r'([1-9]|1[0-7])\s*year(s)?\s*old|'  # Matches 1 to 17 years old
+        r'little|small|tiny|short|new\s*born\s*(boy|girl|bro|brother|sis|sister|shota|lolita|lolli))|'
+        r'(flat[\s_-]*chest(?:ed|s)?|small[\s_-]*chest(?:ed|s)?|medium[\s_-]*chest(?:ed|s)?|tiny[\s_-]*chest(?:ed|s)?|petite[\s_-]*chest(?:ed|s)?|underdeveloped[\s_-]*chest(?:ed|s)?)',
+        re.IGNORECASE
+    )
+    def __init__(self):
+        self.load_tag_generator()
+
+    # Function to remove child-related content from a prompt
+    def remove_child_related_content(prompt):
+        cleaned_prompt = re.sub(PromptGenerator.child_related_regex, '', prompt)
+        return cleaned_prompt.strip()
+
+    # Function to check if a prompt contains child-related content
+    def contains_child_related_content(prompt):
+        if PromptGenerator.child_related_regex.search(prompt):
+            return True
+        return False
+
+    def load_tag_generator(self):
+        torch.set_grad_enabled(False)
+        print("loading tag generator model")
+        tag_config = config.get("tag_generator", {})
+        model = Qwen2ForSequenceClassification.from_pretrained(
+            tag_config.get("model_path"),
+            num_labels=9086,
+            device_map="cpu",
+            local_files_only=True,
+            
+        )
+        model.eval()
+        tokenizer = AutoTokenizer.from_pretrained(tag_config.get("tokenizer_path"), local_files_only=True)
+        allowed_tags = None
+        with open("tags_9083.json", "r") as file:
+            allowed_tags = json.load(file)
+        allowed_tags = sorted(allowed_tags)
+        allowed_tags.append("explicit")
+        allowed_tags.append("questionable")
+        allowed_tags.append("safe")
+
+        self.model = model
+        self.tokenizer = tokenizer
+        self.allowed_tags = allowed_tags
+        print("done")
+    
+    def create_danbooru_tags(self, prompt, threshold):
+        inputs = self.tokenizer(
+            prompt,
+            padding="do_not_pad",
+            max_length=512,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        output = self.model(**inputs).logits
+        output = torch.nn.functional.sigmoid(output)
+        indices = torch.where(output > threshold)
+        values = output[indices]
+        indices = indices[1]
+        values = values.squeeze()
+        skip = set(["simple_background", "multiview","safe", "questionable","mammal"])
+        temp = []
+        tag_score = dict()
+        for i in range(indices.size(0)):
+            temp.append([self.allowed_tags[indices[i]], values[i].item()])
+            tag_score[self.allowed_tags[indices[i]]] = values[i].item()
+        temp = [t[0] for t in temp if t[0] not in skip]
+        text_no_impl = ",".join(temp)
+        current_datetime = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+        print(f"{current_datetime}: finished.")
+        return text_no_impl, tag_score
+
+    def extract_prompt_elements(prompt):
+        chars = characters.get()
+        mentioned_chars = utils.extract_characters("@",prompt)
+        converted_chars = set()
+        chars_in_prompt = {}
+        boys = set()
+        girls = set()
+        for c in mentioned_chars:
+            char_info = chars.get(c["key"], {})
+            if char_info:
+                c["info"] = char_info
+                sample_prompts = char_info.get("samplePrompts", [])
+                gender = c.get("gender", "")
+                prompt = prompt.replace(c["raw"], char_info.get("name", char_info.get("key")).replace("_", " "))#TODO:replace with character name
+                c["long_prompt"] = (lambda s: s)(max(sample_prompts, key=len)).replace("1girl","").replace("1boy","").replace("_", " ")
+                c["short_prompt"] = (lambda s: s)(min(sample_prompts, key=len)).replace("1girl","").replace("1boy","").replace("_", " ")
+                if "girl" in " ".join(char_info.get("samplePrompts",[])):
+                    girls.add(c["key"])
+                    c["gender"] = "girl"
+                elif "boy" in " ".join(char_info.get("samplePrompts",[])):
+                    boys.add(c["key"])
+                    c["gender"] = "boy"
+                else:
+                    c["gender"] = "girl"
+                    girls.add(c["key"])
+            chars_in_prompt[c["key"]] = c
+
+        pattern = re.compile(
+            r'(\d+|one|two|three|multi|multiple)\s*(girl|girls|boy|boys|woman|women|man|man|male|female)',
+            re.IGNORECASE
+        )
+        matches = pattern.findall(prompt)
+
+        # Convert tuples to strings
+        char_counts = ",".join(["".join(match) for match in matches])
+        print(f"1 {char_counts}, {girls}, {boys}")
+        print(chars_in_prompt)
+        if not char_counts:
+            char_counts = ""
+            if boys:
+                if len(boys) == 1:
+                    char_counts += "1boy,"
+                elif len(boys)> 1:
+                    char_counts += "multiple boy,"
+            if girls:
+                if len(girls) == 1:
+                    char_counts += "1girl,"
+                elif len(girls)> 1:
+                    char_counts += "multiple girl,"
+            if len(chars_in_prompt) ==3:
+                char_counts = "threesome,"+char_counts
+            elif len(chars_in_prompt) >3:
+                char_counts = "gangbang,"+char_counts
+                
+        print(f"2 {char_counts}")
+        return chars_in_prompt, char_counts,prompt
+
+    def compose_prompt_il(self, prompt, chars_in_prompt, char_counts, styles, extract_danbooru_tags=False):
+        #char count
+        prompt = PromptGenerator.remove_child_related_content(prompt)
+
+        if "vibe" in styles:
+            vibe_config = styles["vibe"]
+            vibe_style = vibe_config["styles"][0]
+            vibe_prompt = vibe_config["prompt"]
+            prompt = vibe_prompt.replace("{prompt}", prompt).replace("{style}", vibe_style)
+            styles.pop("vibe")
+        
+        new_prompt = ""
+        count_desc = char_counts
+        if count_desc:
+            count_desc += ","
+        
+        char_desc = ""
+        for c in chars_in_prompt:
+            char = chars_in_prompt[c]
+            char_info = char.get("info", {})
+            if char_info:
+                char_desc += char.get("short_prompt") + ","
+                continue
+            char_desc += char.get("name", char.get("key")) + ","
+        if len(chars_in_prompt) > 1:
+            char_desc = char_desc + "side-by-side,"
+        
+        style_prompt = []
+        for key in styles:
+            #check if style need to be expanded
+            if styles[key]:
+                style_prompt.append(styles[key])
+        
+        style_prompt = ",".join(style_prompt)
+
+        
+        if extract_danbooru_tags:
+            tags, _ = self.create_danbooru_tags(prompt, 0.8)
+            if tags:
+                prompt += "," + tags + ","
+        prompt = count_desc + char_desc + prompt + style_prompt
+        return prompt
+    
+    def compose_prompt_pony(self, prompt, chars_in_prompt, char_counts, styles, extract_danbooru_tags=False):
+        #char count
+        prompt = PromptGenerator.remove_child_related_content(prompt)
+
+        if "vibe" in styles:
+            vibe_config = styles["vibe"]
+            vibe_style = vibe_config["styles"][0]
+            vibe_prompt = vibe_config["prompt"]
+            prompt = vibe_prompt.replace("{prompt}", prompt).replace("{style}", vibe_style)
+            styles.pop("vibe")
+        
+        new_prompt = ""
+        count_desc = char_counts
+        if count_desc:
+            count_desc += "{subject:" + count_desc + "},"
+        
+        char_desc = ""
+        for c in chars_in_prompt:
+            char = chars_in_prompt[c]
+            char_info = char.get("info", {})
+            if char_info:
+                gender = char.get("gender", "")
+                char_desc += "{subject}, " + ("male" if gender == "boy" else "female") + "," + char.get("long_prompt") + " \n "
+                continue
+            char_desc += char.get("name", char.get("key")) + " "
+        if char_desc:
+            char_desc = "\n" + char_desc
+        
+        style_prompt = []
+        for key in styles:
+            #check if style need to be expanded
+            if styles[key]:
+                style_prompt.append(styles[key])
+        
+        style_prompt = ",".join(style_prompt)
+
+        
+        if extract_danbooru_tags:
+            tags, _ = self.create_danbooru_tags(prompt, 0.8)
+            if tags:
+                prompt += "," + tags + ","
+        prompt = count_desc + prompt + style_prompt + char_desc
+        return prompt
+    
+    def enhance_prompt(self, prompt, method="pony", styles={}, extract_danbooru_tags=True):
+        chars_in_prompt, char_counts, replaced_prompt = PromptGenerator.extract_prompt_elements(prompt)
+        new_prompt = prompt
+        if method == "il":
+            new_prompt = self.compose_prompt_il(replaced_prompt, chars_in_prompt, char_counts,styles, extract_danbooru_tags)
+        elif method == "pony":
+            new_prompt = self.compose_prompt_pony(replaced_prompt, chars_in_prompt, char_counts,styles, extract_danbooru_tags)
+        return new_prompt
+
+
+prompt_generator = None
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+def common_upscale(samples: torch.Tensor, width: int, height: int, upscale_method: str) -> torch.Tensor:
+    return torch.nn.functional.interpolate(samples, size=(height, width), mode=upscale_method)
+
+def upscale(samples: torch.Tensor, upscale_method: str, scale_by: float) -> torch.Tensor:
+    width = round(samples.shape[3] * scale_by)
+    height = round(samples.shape[2] * scale_by)
+    return common_upscale(samples, width, height, upscale_method)
+
+@spaces.GPU(enable_queue=True)
+def generate_image(model_id, prompt, negative_prompt, width, height, styles={}, scheduler_name=None, use_hd=True, cfg=7.5, steps=30, seed=0, options={}, progress=gr.Progress(track_tqdm=True)):
+    # Clear GPU memory
+    torch.cuda.empty_cache()
+    model = mdoel_mgr.get_model_set_style(model_id)
+    model_name = styles.get("vibe",{}).get("key", "realism")
+    print(model.config)
+    prediction_type = model.config.get("prediction_type", "epsilon")
+    # Start the process
+    pipe = model.pipe
+    pipe.scheduler.config.prediction_type = prediction_type
+    print(pipe.scheduler.config.prediction_type)
+    samplers = {
+        "Euler a": EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config),
+        "DPM++ SDE Karras": DPMSolverSDEScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True),
+        "DPM2 a": DPMSolverMultistepScheduler.from_config(pipe.scheduler.config),
+        "DPM++ SDE": DPMSolverSDEScheduler.from_config(pipe.scheduler.config),
+        "DPM++ 2M SDE": DPMSolverSDEScheduler.from_config(pipe.scheduler.config, use_2m=True),
+        "DPM++ 2S a": DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_2s=True),
+        "NOISE": noise_scheduler,
+    }
+    prompt_analysis_keys = {"is_porn_involved": "boolean", "is_underage_involved": "boolean", "is_non_human_animal_involved": "boolean", "is_revenge_porn": "boolean", "detail_improved_prompt_in_natural_language": "string"}
+    prompt_analysis_result = llm_analyzer.analyze(f"Analyze the following prompt, determine if it contains porn instruction, underage characters, sexual interactino with non-human animals, revenge porn instruction, and generate a JSON with the given format: {json.dumps(prompt_analysis_keys)}\nPrompt: {prompt}", prompt_analysis_keys) or {}
+    print("options", options)
+    print("prompt_analysis_result", prompt_analysis_result)
+    pipe.to(device)
+    #pipe.enable_xformers_memory_efficient_attention()
+    config = model.config
+    total_negative_prompt = config.get("negative_prompt", "")
+    steps = steps or config.get("inference_steps", 30)
+    guidance_scale = cfg or config.get("guidance_scale", 6)
+    width = width or config.get("width", 512)
+    height = height or config.get("height", 512)
+    if not scheduler_name:
+        scheduler_name = config.get("sampler", "")
+        scheduler = samplers.get(scheduler_name, None)
+    else:
+        scheduler = samplers.get(scheduler_name, None)
+    
+    prompt_str = prompt_generator.enhance_prompt(prompt,method=model.config.get("model_version","pony"),styles=styles, extract_danbooru_tags=False)    #!Test!True
+    if config.get("upsample_prompt", True):
+        prompt_str = prompt_str + "," + prompt_analysis_result.get("detail_improved_prompt_in_natural_language", "")
+
+    prompt_str = config.get("prompt", "{prompt}").replace("{prompt}", prompt_str)
+                                                          
+    seed = seed or int(randomize_seed_fn(seed, True))
+    generator = torch.Generator(pipe.device).manual_seed(seed)
+    total_negative_prompt = total_negative_prompt + negative_prompt + vibes.get(styles.get("vibe", ""), {}).get("negative_prompt", "")
+    
+    compel = Compel(
+        tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
+        text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
+        returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+        requires_pooled=[False, True],
+        truncate_long_prompts=False
+    )
+    use_hd = options.get("use_hd", model.config.get("use_hd", True))
+    if use_hd == True or options.get("use_hd", False):
+        apply_hidiffusion(pipe, apply_window_attn=True)
+
+    '''
+    conditioning, pooled = compel(prompt_str)                         # get weighted embeddings for prompt
+    neg_cond, neg_pooled = compel(total_negative_prompt)        
+    
+    print(conditioning.dtype)  # Output: torch.bfloat16
+    print(neg_cond.dtype)  # Output: torch.bfloat16
+                   # embeddings for negative prompt
+    # Ensure the embedding sequences are the same length (pad if needed)
+    [conditioning, neg_cond] = compel.pad_conditioning_tensors_to_same_length([conditioning, neg_cond])
+    '''
+
+    #prompt = helper.get_embed_new(prompt, pipe, compel, only_convert_string=True)
+    #negative_prompt = helper.get_embed_new(total_negative_prompt, pipe, compel, only_convert_string=True)
+    #conditioning, pooled = compel([prompt, negative_prompt])
+    print(prompt_str)
+    if scheduler:
+        pipe.scheduler = scheduler
+    print(f"Generating SDXL, scheduler={scheduler_name}, usd_hd={use_hd}")
+    images = pipe(
+        #prompt_embeds=conditioning, pooled_prompt_embeds=pooled,
+        #negative_prompt_embeds=neg_cond, negative_pooled_prompt_embeds=neg_pooled,
+
+        prompt=prompt_str, negative_prompt=total_negative_prompt,
+        
+        #prompt_embeds=conditioning[0:1],
+        #pooled_prompt_embeds=pooled[0:1],
+        #negative_prompt_embeds=conditioning[1:2],
+        #negative_pooled_prompt_embeds=pooled[1:2],
+        width=width, height=height, guidance_scale=cfg, num_inference_steps=steps, 
+        num_images_per_prompt=1, generator=generator,eta=1.0,
+        #upscaling
+        #output_type="latent",
+    ).images
+    
+    #upscaling
+    '''
+    upscale_by = 1.5
+    upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
+    upscaled_latents = upscale(images, "nearest-exact", upscale_by)
+    images = upscaler_pipe(
+        prompt=prompt_str,
+        negative_prompt=total_negative_prompt,
+        image=upscaled_latents,
+        guidance_scale=cfg,
+        num_inference_steps=steps,
+        strength=0.5,
+        generator=generator,
+        output_type="pil",
+    ).images
+    '''
+    ##
+    
+    webp_images = []
+    nsfw_flags = ["a","h","n","p","s"]
+    for i in images:
+        flags = ""
+        print(i)
+        result = age_predictor.predict(i, 0.6)
+        print(result)
+        underage_flag = "y" if result.get("is_minor", False) else "n"
+        prompt_underage_flag = "y" if prompt_analysis_result.get("is_underage_involved", False) else "n"
+        nsfw_flag = ""
+        bestiality_flag = "n"
+        revenge_porn_flag = "n"
+
+        if prompt_analysis_result.get("is_porn_involved", False):
+            nsfw_flag = "y"
+        else:
+            nsfw_flag = "n"
+        if prompt_analysis_result.get("is_non_human_animal_involved", False):
+            if nsfw_flag == "y":
+                bestiality_flag = "y"
+            else:
+                bestiality_flag = "n"
+        else:
+            bestiality_flag = "n"
+        if prompt_analysis_result.get("is_revenge_porn", False):
+            revenge_porn_flag = "y"
+        else:
+            revenge_porn_flag = "n"
+        watermark = options.get("watermark", {})
+        file_uid = str(uuid.uuid4())
+        # Generate UTC timestamp
+        file_publish_time = int(datetime.now(timezone.utc).timestamp())
+        image_id = (hex(file_publish_time)[2:] + file_uid[-2:]).lower()
+        filename = f"hf_upif_{file_uid}_{file_publish_time}_{image_id}"
+        # image_id= hex of publish_time + last 4 char of file uid
+        should_publish = False
+        model_underage_check = model.config.get("underage_check", ["Teenager 13-20", "Child 0-12"])
+        underage_flag1 = result.get("age_group", "") in model_underage_check
+        if underage_flag == 'y':
+            if nsfw_flag == 'n':
+                underage_flag = 'n'
+            elif len(model_underage_check) < 2 and result.get("confidence", 0) > 0.4:
+                underage_flag = 'y' if underage_flag1 else 'n'
+            
+        if bestiality_flag == "n" and underage_flag == "n" and prompt_underage_flag == "n":
+            if watermark and watermark.get("url", ""):
+                watermark_url = watermark["url"].replace("{image_id}", image_id).replace("{seed}", str(seed)).replace("{prompt}", prompt)
+                i = generate_watermark(original_image=i, qr_code_url=watermark_url)
+                should_publish = True
+        #nsfw_result = nsfw_classifier.predict(i)
+        #print(nsfw_result)
+        #nsfw_lvl = nsfw_result.get("category", 2)
+        
+        #flags += nsfw_flags[nsfw_ lvl]
+        i = save_image(i, filename, f"{underage_flag}{nsfw_flag}{bestiality_flag}{revenge_porn_flag}{prompt_underage_flag}",
+                                      {
+                                        "seed": seed,
+                                        "domain": options.get("domain", ""),
+                                        "model_id": model.config.get("model_id", ""),
+                                        "model_name": model_name,
+                                        "prompt": prompt,
+                                        "publish_time": file_publish_time,
+                                        "file_uid": file_uid,
+                                        "image_id": image_id,
+                                        "underage_flag": underage_flag,
+                                        "nsfw_flag": nsfw_flag,
+                                        "bestiality_flag": bestiality_flag,
+                                        "revenge_porn_flag": revenge_porn_flag,
+                                        "prompt_underage_flag": prompt_underage_flag,
+                                      }
+                                )
+        webp_images.append({
+            "image": i,
+            "should_publish": should_publish
+        })
+    return webp_images
+
+auth = None
+
+@spaces.GPU(enable_queue=True)
+def generate_image_with_ipa(model_id, image, prompt, negative_prompt, width, height, styles={}, scheduler=None, face_strength=7.5, likeness_strength=0.1, steps=30, seed=0, options={}, progress=gr.Progress(track_tqdm=True)):
+    # Clear GPU memory
+    torch.cuda.empty_cache()
+    model = mdoel_mgr.get_model_set_style(model_id)
+    underage_check = age_predictor.predict(image, 0.6)
+    if underage_check.get("is_minor", False):
+        raise Exception("Uploaded image contains inappropriate content")
+
+    prediction_type = model.config.get("prediction_type", "epsilon")
+    # Start the process
+    pipe = model.pipe
+    pipe.scheduler.config.prediction_type = prediction_type
+    print(pipe.scheduler.config.prediction_type)
+    samplers = {
+        "Euler a": EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config),
+        "DPM++ SDE Karras": DPMSolverSDEScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True),
+        "DPM2 a": DPMSolverMultistepScheduler.from_config(pipe.scheduler.config),
+        "DPM++ SDE": DPMSolverSDEScheduler.from_config(pipe.scheduler.config),
+        "DPM++ 2M SDE": DPMSolverSDEScheduler.from_config(pipe.scheduler.config, use_2m=True),
+        "DPM++ 2S a": DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_2s=True),
+        "NOISE": noise_scheduler,
+    }
+    pipe.to(device)
+    #pipe.enable_xformers_memory_efficient_attention()
+    config = model.config
+    total_negative_prompt = config.get("negative_prompt", "")
+    steps = steps or config.get("inference_steps", 30)
+    likeness_strength = likeness_strength or config.get("likeness_strength", 0.1)
+    face_strength = face_strength or config.get("face_strength", 7.0)
+    prompt_analysis_keys = {"is_porn_involved": "boolean", "is_underage_involved": "boolean", "is_non_human_animal_involved": "boolean", "is_revenge_porn": "boolean", "detail_improved_prompt_in_natural_language": "string"}
+    prompt_analysis_result = llm_analyzer.analyze(f"Analyze the following prompt, determine if it contains porn instruction, underage characters, sexual interactino with non-human animals, revenge porn instruction, and generate a JSON with the given format: {json.dumps(prompt_analysis_keys)}\nPrompt: {prompt}", prompt_analysis_keys) or {}
+
+    width = width or config.get("width", 512)
+    height = height or config.get("height", 512)
+    if not scheduler:
+        scheduler = config.get("sampler", "")
+        scheduler = samplers.get(scheduler, None)
+    else:
+        scheduler = samplers.get(scheduler, None)
+    
+    prompt_str = prompt_generator.enhance_prompt(prompt,method=model.config.get("model_version","pony"),styles=styles, extract_danbooru_tags=False)    #!Test!True
+    prompt_str = config.get("prompt", "{prompt}").replace("{prompt}", prompt_str)
+    prompt_str = prompt_str + "," + prompt_analysis_result.get("detail_improved_prompt_in_natural_language", "")
+    seed = seed or int(randomize_seed_fn(seed, True))
+    generator = torch.Generator(pipe.device).manual_seed(seed)
+    total_negative_prompt = total_negative_prompt + negative_prompt + vibes.get(styles.get("vibe", ""), {}).get("negative_prompt", "")
+    
+    app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+    app.prepare(ctx_id=0, det_size=(512, 512))
+    
+    faceid_all_embeds = []
+    
+    face = cv2.imread(image)
+    faces = app.get(face)
+    faceid_embed = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0)
+    faceid_all_embeds.append(faceid_embed)
+
+    average_embedding = torch.mean(torch.stack(faceid_all_embeds, dim=0), dim=0)
+
+    #prompt = helper.get_embed_new(prompt, pipe, compel, only_convert_string=True)
+    #negative_prompt = helper.get_embed_new(total_negative_prompt, pipe, compel, only_convert_string=True)
+    #conditioning, pooled = compel([prompt, negative_prompt])
+    print(prompt_str)
+    if scheduler:
+        pipe.scheduler = scheduler
+        
+    print(f"Generating SDXL, scheduler={scheduler}, usd_hd={use_hd}")
+    images = model.ip_model.generate(
+        prompt=prompt, negative_prompt=total_negative_prompt, faceid_embeds=average_embedding,num_samples=1,num_images_per_prompt=1,
+        scale=likeness_strength, width=width, height=height, guidance_scale=face_strength, num_inference_steps=steps
+    )    
+    #upscaling
+    '''
+    upscale_by = 1.5
+    upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
+    upscaled_latents = upscale(images, "nearest-exact", upscale_by)
+    images = upscaler_pipe(
+        prompt=prompt_str,
+        negative_prompt=total_negative_prompt,
+        image=upscaled_latents,
+        guidance_scale=cfg,
+        num_inference_steps=steps,
+        strength=0.5,
+        generator=generator,
+        output_type="pil",
+    ).images
+    '''
+    ##
+    
+    webp_images = []
+    for i in images:
+        print(i)
+        result = age_predictor.predict(i, 0.6)
+        print(result)
+        webp_images.append(helper.save_image(i, "fua" if result.get("is_minor", False) else "fip"))
+        #webp_images.append(helper.save_image(i))
+    return webp_images
+
+mdoel_mgr = None
+
+
+# Example usage in generate function
+@spaces.GPU(enable_queue=True)
+def generate(model_id, image, prompt, negative_prompt, scheduler, use_hd, face_strength, likeness_strength, width, height, publish, seed=0, progress=gr.Progress(track_tqdm=True)):
+    try:
+        if image is not None:
+            #save the image to file
+            image_path = f"temp_image_{time.time()}.jpg"
+            image = Image.fromarray(image)
+            image.save(image_path)
+            try:
+                return generate_image_with_ipa(model_id, image_path, prompt, negative_prompt, width, height, {"vibe":vibes[model_id]}, 
+                                             scheduler=None, face_strength=face_strength, 
+                                             likeness_strength=likeness_strength, steps=30, seed=seed, options={"domain": "nsfwais.io"}, progress=progress)
+            finally:
+                os.remove(image_path)
+        else:
+            result = generate_image(model_id, prompt, negative_prompt, width, height, {"vibe":vibes[model_id]}, 
+                                scheduler, use_hd, cfg=face_strength, steps=30, seed=seed, options={"domain": "nsfwais.io"}, progress=progress)
+            display_image_path = result[0]["image"]
+            if not os.path.exists(display_image_path):
+                raise FileNotFoundError(f"Image file not found: {display_image_path}")
+
+            # Open it with PIL
+            display_image = Image.open(display_image_path)
+
+            # (Optional) convert mode if needed, e.g. to RGB
+            # display_image = display_image.convert("RGB")
+
+            if publish:
+                try:
+                    publish_result = utils.publish_url_sync(local_config.get("download_url",""), result[0]["image"], config.get("publish_url", ""))
+                    print(publish_result)
+                except Exception as e:
+                    print(e)
+            print(result)
+            
+            return [display_image]
+    finally:
+        # Force CUDA memory cleanup after generation
+        torch.cuda.empty_cache()
+generate.zerogpu=True
+
+def ipa_image_gen_api(source_image, prompt, negative_prompt, styles, options, token, request: gr.Request, progress=gr.Progress(track_tqdm=True)):
+    #auth.check_auth("/image_gen_api", request, token)
+    print(styles)
+    if isinstance(styles, str):
+        styles = json.loads(styles)
+    if isinstance(options, str):
+        options = json.loads(options)
+    if source_image is not None:
+        source_image = Image.fromarray(source_image)
+        #save the image to file
+        image_path = f"temp_image_{time.time()}.jpg"
+        source_image.save(image_path)
+        try:
+            return image_gen_api_entry("ipa", image_path, prompt, negative_prompt, styles, options, token, mdoel_mgr, progress)
+        finally:
+            os.remove(image_path)
+    else:
+        return image_gen_api_entry("image", None, prompt, negative_prompt, styles, options, token, mdoel_mgr, progress)
+ipa_image_gen_api.zerogpu=True
+
+def image_gen_api(prompt, negative_prompt, styles, options, token, request: gr.Request, progress=gr.Progress(track_tqdm=True)):
+    #auth.check_auth("/image_gen_api", request, token)
+    print(styles)
+    if isinstance(styles, str):
+        styles = json.loads(styles)
+    if isinstance(options, str):
+        options = json.loads(options)
+        if "domain" not in options:
+            options["domain"] = "nsfwais.io"
+    return image_gen_api_entry("image", None, prompt, negative_prompt, styles, options, token, mdoel_mgr, progress)
+image_gen_api.zerogpu=True
+
+def image_gen_api_entry(mode, source_image, prompt, negative_prompt, styles, options, token, use_model_mgr:ModelManager, progress=gr.Progress(track_tqdm=True)):
+    size_2_width_height = {
+        "1_1": [[1280, 1280], [2048, 2048]],  # Square (already multiple of 8)
+        "4_3": [[1200, 896], [2048, 1536]],  # Twitter/X (1200 rounded to 1192, 900 to 896)
+        "16_9": [[1280, 720], [1920, 1088]],  # Facebook (1280 is fine, 1080 rounded to 1088)
+        "4_5": [[1080, 1344], [2048, 2560]],  # Instagram (1080 rounded to 1080, 1350 to 1344)
+        "9_16": [[720, 1280], [1080, 1920]],  # Tiktok (720 is fine, 1280 is fine)
+    }
+    for s in styles:
+        if styles[s] == "Auto":
+            styles[s] = ""
+    vibe = styles.get("vibe","") or "realistic"
+    vibe_config = vibes.get(vibe, {})
+    if not vibe_config:
+        raise Exception(f"invalid vibe {vibe}")
+    style_name = vibe_config.get("models")[0]
+    styles["vibe"] = vibe_config
+    #lighting = styles.get("lighting", "")
+    #camera = styles.get("camera", "")
+    size = styles.get("size", "") or "4_3"
+    styles.pop("size",None)
+    steps = options.get("steps", 30)
+    seed = options.get("seed",0)
+    
+    width, height = size_2_width_height.get(size)[0]
+
+    #prompt = prompt_generator.enhance_prompt(prompt)
+    #source_image = Image.fromarray(source_image)
+
+
+    '''
+    chars = characters.get()
+    mentioned_chars = utils.extract_characters(prompt)
+    converted_chars = set()
+    for c in mentioned_chars:
+        char_info = chars.get(c["key"], {})
+        char_prompts = char_info.get("samplePrompts", [])
+        if char_prompts:
+            func = max
+            if c["key"] in converted_chars:
+                func = min
+            
+            char_prompt = (lambda s: s)(func(char_prompts, key=len))
+            if char_prompt:
+                prompt = prompt.replace(c["raw"], char_prompt)
+                converted_chars.add(c["key"])
+        else:
+            print("char "+c["key"]+" info not found, skip interpretting")
+    '''
+    
+    if mode == "image":
+        generated_image = generate_image(vibe, prompt, negative_prompt, width, height, styles, scheduler_name=None,use_hd=None, cfg=5, steps=steps, seed=seed, options=options, progress=progress)[0]
+        if generated_image.get("should_publish", False):
+            try:
+                result = utils.publish_url_sync(local_config.get("download_url",""), generated_image["image"], config.get("publish_url", ""), timeout=10)
+                if "url" in result:
+                    return generated_image["image"], result["url"]
+                else:
+                    print(f"error in result{result}")
+            except Exception as e:
+                print(e)
+        
+        return generated_image["image"], ""
+
+    elif mode == "ipa":
+        face_strength = options.get("face_strength", 11)
+        likeness_strength = options.get("likeness_strength", 0.5)
+        return generate_image_with_ipa(vibe, source_image, prompt, negative_prompt, width, height, styles, scheduler=None, face_strength=face_strength, likeness_strength=likeness_strength, steps=steps, seed=seed, options=options, progress=progress)[0]
+    else:
+        raise Exception(f"invalid mode {mode}")
+
+
+
+def swap_to_gallery(images):
+    return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
+
+def remove_back_to_files():
+    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+css = '''
+h1{margin-bottom: 0 !important}
+'''
+with gr.Blocks(css=css) as demo:
+    #gr.Markdown("# IP-Adapter-FaceID SDXL demo")
+    #gr.Markdown("A simple Demo for the [h94/IP-Adapter-FaceID SDXL model](https://huggingface.co/h94/IP-Adapter-FaceID) together with [Juggernaut XL v7](https://huggingface.co/stablediffusionapi/juggernaut-xl-v7). You should run this on at least 24 GB of VRAM.")
+    with gr.Row():
+        with gr.Column():
+            '''
+            files = gr.Files(
+                        label="Drag 1 or more photos of your face",
+                        file_types=["image"]
+                    )
+            '''
+            upload_image = gr.Image(label="Upload an image")
+            #uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=250)
+            #with gr.Column(visible=False) as clear_button:
+                #remove_and_reupload = gr.ClearButton(value="Remove files and upload new ones", components=files, size="sm")
+            model = gr.Textbox(label="",
+                        info="",
+                        placeholder="",
+                        value=config.get("default_model", ""))
+            
+            prompt = gr.Textbox(label="Prompt",
+                        info="Try something like 'a photo of a man/woman/person'",
+                        placeholder="A photo of a man/woman/person ...",
+                        lines=10,
+                        value="")
+            negative_prompt = gr.Textbox(label="Negative Prompt", info="What the model should NOT produce.",placeholder="low quality", value="worst quality, low quality")
+            # Text box for the width
+            width = gr.Number(value=config.get("width", 1280), label="Width", info="Width of the generated image")
+            # Text box for the height
+            height = gr.Number(value=config.get("height", 1280), label="Height", info="Height of the generated image")
+            scheduler = gr.Textbox(label="scheduler", info="",placeholder="low quality", value="Euler a")
+            use_hd = gr.Checkbox(label="use hidiffusion", value=config.get("use_hd", False))
+            publish = gr.Checkbox(label="publish", value=False, visible=config.get("publish_url", "") != "")
+            negative_prompt = gr.Textbox(label="Negative Prompt", info="What the model should NOT produce.",placeholder="low quality", value="worst quality, low quality")
+            seed = gr.Number(label="Seed", info="Seed for the generated image", value=0)
+            style = "Photorealistic"
+            face_strength = gr.Slider(label="Prompt Strength", info="How much the written prompt weighs into the generated images.", value=7.5, step=0.1, minimum=0, maximum=15)
+            likeness_strength = gr.Slider(label="Photo Embedding Strength", info="How much your uploaded files weigh into the generated images.", value=1.0, step=0.1, minimum=0, maximum=5)
+
+            submit = gr.Button("Submit", variant="primary")
+            submit1 = gr.Button("image_gen_api", variant="primary", visible=False)
+            submit2 = gr.Button("image_gen_api_ipa", variant="primary", visible=False)
+        with gr.Column():
+            gallery = gr.Gallery(label="Generated Images")
+        #files.upload(fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files])
+        #remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files])
+        submit.click(fn=generate,
+                    inputs=[model, upload_image,prompt,negative_prompt,scheduler, use_hd, face_strength, likeness_strength, width, height, publish, seed],
+                    outputs=gallery)
+        
+        submit1.click(fn=image_gen_api,
+                    inputs=[prompt,negative_prompt, negative_prompt, negative_prompt, negative_prompt],
+                    outputs=[gr.Image(visible=False), gr.Textbox(visible=False)])
+        submit2.click(fn=ipa_image_gen_api,
+                    inputs=[upload_image, prompt,negative_prompt, negative_prompt, negative_prompt, negative_prompt],
+                    outputs=gr.Image(visible=False))
+        
+    
+    # gr.Markdown("This demo includes extra features to mitigate the implicit bias of the model and prevent explicit usage of it to generate content with faces of people, including third parties, that is not safe for all audiences, including naked or semi-naked people.")
+
+if __name__ == "__main__":
+    # Read a local.json for local config
+
+    auth = AuthHelper()
+    download_for_imagegen.download_models()
+    prompt_generator = PromptGenerator()
+    mdoel_mgr = ModelManager(config, vibes, lazy_load=local_config.get("lazy_load", False), max_loaded_models=local_config.get("max_loaded_models", 1))
+    print("launching")
+    demo.launch()

requirements.txt

@@ -3,7 +3,6 @@
 torch==2.4.0
 
 # Core CUDA stack (these three versions are a matched set)
-spaces==0.30.0
 pydantic==2.10.6
 insightface==0.7.3
 diffusers