Create app.py

app.py

@@ -0,0 +1,194 @@
+import torch
+import spaces
+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, snapshot_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")
+}
+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,
+    beta_end=0.012,
+    beta_schedule="scaled_linear",
+    clip_sample=False,
+    set_alpha_to_one=False,
+    steps_offset=1,
+)
+vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)
+
+# 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
+@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:
+        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)
+            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:
+        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
+        )
+    else:
+        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
+        )
+    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)
+
+def remove_back_to_files():
+    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+
+css = '''
+h1{margin-bottom: 0 !important}
+footer{display:none !important}
+'''
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown("")
+    gr.Markdown("")
+    with gr.Row():
+        with gr.Column():
+            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 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")
+            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(
+                label="Model Choice",
+                choices=list(model_paths.keys()),
+                value="Realistic Vision V4.0"
+            )
+            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"
+                )
+        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, prompt, negative_prompt, preserve, face_strength, likeness_strength, nfaa_negative_prompts, model_choice],
+            outputs=gallery
+        )
+    
+    gr.Markdown("")
+
+demo.launch()