Update app.py

app.py

@@ -2,6 +2,7 @@ import torch
 import spaces
 from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL, DiffusionPipeline
 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
@@ -9,12 +10,17 @@ from insightface.utils import face_align
 import gradio as gr
 import cv2
 
+pipeline = DiffusionPipeline.from_pretrained("fluently/Fluently-XL-v2")
+pipeline.load_lora_weights("ehristoforu/dalle-3-xl-v2")
+
 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"
-
 noise_scheduler = DDIMScheduler(
     num_train_timesteps=1000,
     beta_start=0.00085,
@@ -24,11 +30,7 @@ noise_scheduler = DDIMScheduler(
     set_alpha_to_one=False,
     steps_offset=1,
 )
-
 vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)
-
-pipeline = DiffusionPipeline.from_pretrained("fluently/Fluently-XL-v2")
-pipeline.load_lora_weights("ehristoforu/dalle-3-xl-v2")
 pipe = pipeline.to(device)
 
 ip_model = IPAdapterFaceID(pipe, ip_ckpt, device)
@@ -49,6 +51,7 @@ def generate_image(images, prompt, negative_prompt, preserve_face_structure, fac
         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)
     total_negative_prompt = f"{negative_prompt} {nfaa_negative_prompt}"