Update app.py

app.py

@@ -1,36 +1,28 @@
 import torch
 import spaces
-from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL
+from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL, StableDiffusionXLPipeline
 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, snapshot_download
+from huggingface_hub import hf_hub_download
 from insightface.app import FaceAnalysis
 from insightface.utils import face_align
 import gradio as gr
 import cv2
 import os
 
-# Model paths
-model_paths = {
-    "Realistic Vision V4.0": "SG161222/Realistic_Vision_V4.0_noVAE",
-    "Pony Realism v21": snapshot_download(repo_id="John6666/pony-realism-v21main-sdxl"),
-    "Cyber Realistic Pony v61": snapshot_download(repo_id="John6666/cyberrealistic-pony-v61-sdxl"),
-    "Stallion Dreams Pony Realistic v1": snapshot_download(repo_id="John6666/stallion-dreams-pony-realistic-v1-sdxl")
-}
+base_model_path = "SG161222/Realistic_Vision_V4.0_noVAE"
 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 Checker Setup
 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"
 
-# Define the scheduler
 noise_scheduler = DDIMScheduler(
     num_train_timesteps=1000,
     beta_start=0.00085,
@@ -41,35 +33,39 @@ noise_scheduler = DDIMScheduler(
     steps_offset=1,
 )
 vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)
+pipe = StableDiffusionPipeline.from_pretrained(
+    base_model_path,
+    torch_dtype=torch.float16,
+    scheduler=noise_scheduler,
+    vae=vae,
+    feature_extractor=safety_feature_extractor,
+    safety_checker=None  # <--- Disable safety checker
+).to(device)
+
+ip_model = IPAdapterFaceID(pipe, ip_ckpt, device)
+ip_model_plus = IPAdapterFaceIDPlus(pipe, image_encoder_path, ip_plus_ckpt, device)
 
-# Face analysis setup
 app = FaceAnalysis(name="buffalo_l", providers=['CPUExecutionProvider'])
 app.prepare(ctx_id=0, det_size=(640, 640))
 
 cv2.setNumThreads(1)
 
-# Function to load the appropriate pipeline based on user selection
-def load_model(model_choice):
-    model_path = model_paths[model_choice]
-    pipeline = StableDiffusionPipeline.from_pretrained(
-        model_path,
-        torch_dtype=torch.float16,
-        scheduler=noise_scheduler,
-        vae=vae,
-        feature_extractor=safety_feature_extractor,
-        safety_checker=None
-    ).to(device)
-
-    # Load the IP Adapter models
-    ip_model = IPAdapterFaceID(pipeline, ip_ckpt, device)
-    ip_model_plus = IPAdapterFaceIDPlus(pipeline, image_encoder_path, ip_plus_ckpt, device)
-
-    return pipeline, ip_model, ip_model_plus
-
-# Gradio function to generate images
+# Download the SDXL model files
+ckpt_dir_pony = snapshot_download(repo_id="John6666/pony-realism-v21main-sdxl")
+ckpt_dir_cyber = snapshot_download(repo_id="John6666/cyberrealistic-pony-v61-sdxl")
+ckpt_dir_stallion = snapshot_download(repo_id="John6666/stallion-dreams-pony-realistic-v1-sdxl")
+
+# Load the SDXL models
+pipe_pony = StableDiffusionXLPipeline.from_pretrained(ckpt_dir_pony, torch_dtype=torch.float16)
+pipe_cyber = StableDiffusionXLPipeline.from_pretrained(ckpt_dir_cyber, torch_dtype=torch.float16)
+pipe_stallion = StableDiffusionXLPipeline.from_pretrained(ckpt_dir_stallion, torch_dtype=torch.float16)
+
+pipe_pony.to(device)
+pipe_cyber.to(device)
+pipe_stallion.to(device)
+
 @spaces.GPU(enable_queue=True)
 def generate_image(images, prompt, negative_prompt, preserve_face_structure, face_strength, likeness_strength, nfaa_negative_prompt, model_choice, progress=gr.Progress(track_tqdm=True)):
-    pipeline, ip_model, ip_model_plus = load_model(model_choice)
     faceid_all_embeds = []
     first_iteration = True
     for image in images:
@@ -77,43 +73,41 @@ def generate_image(images, prompt, negative_prompt, preserve_face_structure, fac
         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)
+        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)
+    
     total_negative_prompt = f"{negative_prompt} {nfaa_negative_prompt}"
-
-    if not preserve_face_structure:
+    
+    if model_choice == "Pony Realism v21":
+        pipe = pipe_pony
+    elif model_choice == "Cyber Realistic Pony v61":
+        pipe = pipe_cyber
+    else:  # "Stallion Dreams Pony Realistic v1"
+        pipe = pipe_stallion
+    
+    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=512,
-            height=512,
-            num_inference_steps=30
+            prompt=prompt, negative_prompt=total_negative_prompt, faceid_embeds=average_embedding,
+            scale=likeness_strength, width=512, height=512, num_inference_steps=30, pipe=pipe
         )
     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=512,
-            height=512,
-            num_inference_steps=30
+            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=512, height=512, num_inference_steps=30, pipe=pipe
         )
+    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)
+        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)
+        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)
@@ -132,63 +126,39 @@ with gr.Blocks(css=css) as demo:
     with gr.Row():
         with gr.Column():
             files = gr.Files(
-                label="Drag 1 or more photos of your face",
-                file_types=["image"]
-            )
+                        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 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]..."
-            )
+            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")
-            style = gr.Radio(
-                label="Generation type",
-                info="For stylized try prompts like 'a watercolor painting of a woman'",
-                choices=["Photorealistic", "Stylized"],
-                value="Photorealistic"
-            )
+            style = gr.Radio(label="Generation type", info="For stylized try prompts like 'a watercolor painting of a woman'", choices=["Photorealistic", "Stylized"], value="Photorealistic")
             model_choice = gr.Dropdown(
+                ["Pony Realism v21", "Cyber Realistic Pony v61", "Stallion Dreams Pony Realistic v1"],
                 label="Model Choice",
-                choices=list(model_paths.keys()),
-                value="Realistic Vision V4.0"
+                value="Pony Realism v21"
             )
             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
-                )
+                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"
-                )
+                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")    
         with gr.Column():
             gallery = gr.Gallery(label="Generated Images")
         style.change(fn=change_style,
-                     inputs=style,
-                     outputs=[preserve, face_strength, likeness_strength])
+                    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, prompt, negative_prompt, preserve, face_strength, likeness_strength, nfaa_negative_prompts, model_choice],
-            outputs=gallery
-        )
+        submit.click(fn=generate_image,
+                    inputs=[files,prompt,negative_prompt,preserve, face_strength, likeness_strength, nfaa_negative_prompts, model_choice],
+                    outputs=gallery)
     
     gr.Markdown("")
-
-demo.launch()
+    
+demo.launch()