Update app.py

app.py

@@ -1,88 +1,36 @@
-import torch
+rt torch
 import spaces
-from diffusers import DDIMScheduler, StableDiffusionXLPipeline
-import ipown
+from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL
+from transformers import AutoFeatureExtractor
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from ip_adapter.ip_adapter_faceid import IPAdapterFaceID, IPAdapterFaceIDPlus
 from huggingface_hub import hf_hub_download
 from insightface.app import FaceAnalysis
+from insightface.utils import face_align
 import gradio as gr
 import cv2
 
-# List of models for switching
-model_options = {
-    "CyberRealistic": "John6666/cyberrealistic-pony-v61-sdxl",
-    "StallionDreams": "John6666/stallion-dreams-pony-realistic-v1-sdxl",
-    "PonyRealism": "John6666/pony-realism-v21main-sdxl"
+base_model_paths = {
+    "RealisticVisionV4": "SG161222/Realistic_Vision_V4.0_noVAE",
+    "RealisticVisionV6": "SG161222/Realistic_Vision_V6.0_B1_noVAE",
+    "Deliberate": "Yntec/Deliberate",
+    "DeliberateV2": "Yntec/Deliberate2",
+    "Dreamshaper8": "Lykon/dreamshaper-8",
+    "EpicRealism": "emilianJR/epiCRealism"
 }
 
-# Full style list for applying styles to the prompt
-style_list = [
-    {
-        "name": "(No style)",
-        "prompt": "{prompt}",
-        "negative_prompt": "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
-    },
-    {
-        "name": "Cinematic",
-        "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
-        "negative_prompt": "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured",
-    },
-    {
-        "name": "3D Model",
-        "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting",
-        "negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting",
-    },
-    {
-        "name": "Anime",
-        "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime,  highly detailed",
-        "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast",
-    },
-    {
-        "name": "Digital Art",
-        "prompt": "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed",
-        "negative_prompt": "photo, photorealistic, realism, ugly",
-    },
-    {
-        "name": "Photographic",
-        "prompt": "cinematic photo {prompt} . 35mm photograph, film, bokeh, professional, 4k, highly detailed",
-        "negative_prompt": "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly",
-    },
-    {
-        "name": "Pixel art",
-        "prompt": "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics",
-        "negative_prompt": "sloppy, messy, blurry, noisy, highly detailed, ultra textured, photo, realistic",
-    },
-    {
-        "name": "Fantasy art",
-        "prompt": "ethereal fantasy concept art of  {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
-        "negative_prompt": "photographic, realistic, realism, 35mm film, dslr, cropped, frame, text, deformed, glitch, noise, noisy, off-center, deformed, cross-eyed, closed eyes, bad anatomy, ugly, disfigured, sloppy, duplicate, mutated, black and white",
-    },
-    {
-        "name": "Neonpunk",
-        "prompt": "neonpunk style {prompt} . cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional",
-        "negative_prompt": "painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured",
-    },
-    {
-        "name": "Manga",
-        "prompt": "manga style {prompt} . vibrant, high-energy, detailed, iconic, Japanese comic style",
-        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, realism, photorealistic, Western comic style",
-    },
-]
-
-# Styles dictionary to map style names to prompts and negative prompts
-styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
-STYLE_NAMES = list(styles.keys())
-DEFAULT_STYLE_NAME = "(No style)"
-
-# Function to apply the selected style
-def apply_style(style_name: str, positive: str, negative: str = "") -> tuple[str, str]:
-    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-    return p.replace("{prompt}", positive), n + negative
-
-# Download the necessary model component
-ip_ckpt = hf_hub_download(repo_id="h94/IP-Adapter-FaceID", filename="ip-adapter-faceid_sdxl.bin", repo_type="model")
+
+vae_model_path = "stabilityai/sd-vae-ft-mse"
+image_encoder_path = "laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
+ip_ckpt = hf_hub_download(repo_id="h94/IP-Adapter-FaceID", filename="ip-adapter-faceid_sd15.bin", repo_type="model")
+ip_plus_ckpt = hf_hub_download(repo_id="h94/IP-Adapter-FaceID", filename="ip-adapter-faceid-plusv2_sd15.bin", repo_type="model")
+
+safety_model_id = "CompVis/stable-diffusion-safety-checker"
+safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id)
+safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id)
+
 device = "cuda"
 
-# Configure the noise scheduler
 noise_scheduler = DDIMScheduler(
     num_train_timesteps=1000,
     beta_start=0.00085,
@@ -92,55 +40,71 @@ noise_scheduler = DDIMScheduler(
     set_alpha_to_one=False,
     steps_offset=1,
 )
+vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)
 
-# Function to initialize the pipeline with a selected model
-def get_pipeline(model_path):
-    return StableDiffusionXLPipeline.from_pretrained(
-        model_path,
+def load_model(base_model_path):
+    pipe = StableDiffusionPipeline.from_pretrained(
+        base_model_path,
         torch_dtype=torch.float16,
         scheduler=noise_scheduler,
-        use_safetensors=True,
-    )
+        vae=vae,
+        feature_extractor=safety_feature_extractor,
+        safety_checker=None  # <--- Disable safety checker
+    ).to(device)
+    return pipe
 
-# Initialize with a default model
-current_model = model_options["PonyRealism"]
-pipe = get_pipeline(current_model)
-ip_model = ipown.IPAdapterFaceIDXL(pipe, ip_ckpt, device)
+ip_model = None
+ip_model_plus = None
 
-@spaces.GPU()
-def generate_image(images, model_choice, style_name, prompt, negative_prompt, face_strength, likeness_strength, num_inference_steps, guidance_scale, width, height):
-    global current_model, pipe, ip_model
+app = FaceAnalysis(name="buffalo_l", providers=['CPUExecutionProvider'])
+app.prepare(ctx_id=0, det_size=(640, 640))
 
-    # Update the model if the choice has changed
-    if model_options[model_choice] != current_model:
-        current_model = model_options[model_choice]
-        pipe = get_pipeline(current_model)
-        ip_model = ipown.IPAdapterFaceIDXL(pipe, ip_ckpt, device)
+cv2.setNumThreads(1)
 
-    torch.cuda.empty_cache()
-
-    app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
-    app.prepare(ctx_id=0, det_size=(512, 512))
+@spaces.GPU(enable_queue=True)
+def generate_image(images, prompt, negative_prompt, preserve_face_structure, face_strength, likeness_strength, nfaa_negative_prompt, base_model, num_inference_steps, guidance_scale, width, height, progress=gr.Progress(track_tqdm=True)):
+    global ip_model, ip_model_plus
+    base_model_path = base_model_paths[base_model]
+    pipe = load_model(base_model_path)
+    ip_model = IPAdapterFaceID(pipe, ip_ckpt, device)
+    ip_model_plus = IPAdapterFaceIDPlus(pipe, image_encoder_path, ip_plus_ckpt, device)
 
     faceid_all_embeds = []
+    first_iteration = True
     for image in images:
         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)
-
+        if(first_iteration and preserve_face_structure):
+            face_image = face_align.norm_crop(face, landmark=faces[0].kps, image_size=224) # you can also segment the face
+            first_iteration = False
+            
     average_embedding = torch.mean(torch.stack(faceid_all_embeds, dim=0), dim=0)
     
-    # Apply the selected style
-    styled_prompt, styled_negative_prompt = apply_style(style_name, prompt, negative_prompt)
+    total_negative_prompt = f"{negative_prompt} {nfaa_negative_prompt}"
     
-    image = ip_model.generate(
-        prompt=styled_prompt, negative_prompt=styled_negative_prompt, faceid_embeds=average_embedding,
-        scale=likeness_strength, width=width, height=height, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps
-    )
-
+    if(not preserve_face_structure):
+        print("Generating normal")
+        image = ip_model.generate(
+            prompt=prompt, negative_prompt=total_negative_prompt, faceid_embeds=average_embedding,
+            scale=likeness_strength, width=width, height=height, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale
+        )
+    else:
+        print("Generating plus")
+        image = ip_model_plus.generate(
+            prompt=prompt, negative_prompt=total_negative_prompt, faceid_embeds=average_embedding,
+            scale=likeness_strength, face_image=face_image, shortcut=True, s_scale=face_strength, width=width, height=height, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale
+        )
+    print(image)
     return image
 
+def change_style(style):
+    if style == "Photorealistic":
+        return(gr.update(value=True), gr.update(value=1.3), gr.update(value=1.0))
+    else:
+        return(gr.update(value=True), gr.update(value=0.1), gr.update(value=0.8))
+
 def swap_to_gallery(images):
     return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
 
@@ -149,34 +113,48 @@ def remove_back_to_files():
 
 css = '''
 h1{margin-bottom: 0 !important}
+footer{display:none !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).")
+    gr.Markdown("")
+    gr.Markdown("")
     with gr.Row():
         with gr.Column():
-            model_dropdown = gr.Dropdown(label="Select Model", choices=list(model_options.keys()), value="PonyRealism")
-            style_dropdown = gr.Dropdown(label="Style", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)
-            files = gr.Files(label="Drag 1 or more photos of your face", file_types=["image"])
-            uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=250)
+            files = gr.Files(
+                        label="Drag 1 or more photos of your face",
+                        file_types=["image"]
+                    )
+            uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=125)
             with gr.Column(visible=False) as clear_button:
-                remove_and_reupload = gr.ClearButton(value="Remove files and upload new ones", components=files, size="sm")
-            prompt = gr.Textbox(label="Prompt", placeholder="A photo of a man/woman/person ...")
+                remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=files, size="sm")
+            prompt = gr.Textbox(label="Prompt",
+                       info="Try something like 'a photo of a man/woman/person'",
+                       placeholder="A photo of a [man/woman/person]...")
             negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="low quality")
-            face_strength = gr.Slider(label="Face Strength", value=7.5, step=0.1, minimum=0, maximum=15)
-            likeness_strength = gr.Slider(label="Likeness Strength", value=1.0, step=0.1, minimum=0, maximum=5)
-            with gr.Accordion("Advanced Options", open=False):
+            style = gr.Radio(label="Generation type", info="For stylized try prompts like 'a watercolor painting of a woman'", choices=["Photorealistic", "Stylized"], value="Photorealistic")
+            base_model = gr.Dropdown(label="Base Model", choices=list(base_model_paths.keys()), value="Realistic_Vision_V4.0_noVAE")
+            submit = gr.Button("Submit")
+            with gr.Accordion(open=False, label="Advanced Options"):
+                preserve = gr.Checkbox(label="Preserve Face Structure", info="Higher quality, less versatility (the face structure of your first photo will be preserved). Unchecking this will use the v1 model.", value=True)
+                face_strength = gr.Slider(label="Face Structure strength", info="Only applied if preserve face structure is checked", value=1.3, step=0.1, minimum=0, maximum=3)
+                likeness_strength = gr.Slider(label="Face Embed strength", value=1.0, step=0.1, minimum=0, maximum=5)
+                nfaa_negative_prompts = gr.Textbox(label="Appended Negative Prompts", info="Negative prompts to steer generations towards safe for all audiences outputs", value="naked, bikini, skimpy, scanty, bare skin, lingerie, swimsuit, exposed, see-through")
                 num_inference_steps = gr.Slider(label="Number of Inference Steps", value=30, step=1, minimum=10, maximum=100)
                 guidance_scale = gr.Slider(label="Guidance Scale", value=7.5, step=0.1, minimum=1, maximum=20)
                 width = gr.Slider(label="Width", value=512, step=64, minimum=256, maximum=1024)
                 height = gr.Slider(label="Height", value=512, step=64, minimum=256, maximum=1024)
-            submit = gr.Button("Submit", variant="primary")
         with gr.Column():
             gallery = gr.Gallery(label="Generated Images")
-        
+        style.change(fn=change_style,
+                    inputs=style,
+                    outputs=[preserve, face_strength, likeness_strength])
         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_image, inputs=[files, model_dropdown, style_dropdown, prompt, negative_prompt, face_strength, likeness_strength, num_inference_steps, guidance_scale, width, height], outputs=gallery)
-
+        submit.click(fn=generate_image,
+                    inputs=[files,prompt,negative_prompt,preserve, face_strength, likeness_strength, nfaa_negative_prompts, base_model, num_inference_steps, guidance_scale, width, height],
+                    outputs=gallery)
+    
+    gr.Markdown("")
+    
 demo.launch()