Update app.py

app.py

@@ -1,16 +1,32 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionXLPipeline, AutoencoderKL
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 from share_btn import community_icon_html, loading_icon_html, share_js
 from cog_sdxl_dataset_and_utils import TokenEmbeddingsHandler
+
 import lora
 import copy
 import json
 import gc
 import random
 from urllib.parse import quote
+import gdown
+import os
+
+import diffusers
+from diffusers.utils import load_image
+from diffusers.models import ControlNetModel
+from diffusers import AutoencoderKL, DPMSolverMultistepScheduler
+import cv2
+import torch
+import numpy as np
+from PIL import Image
+
+from insightface.app import FaceAnalysis
+from pipeline_stable_diffusion_xl_instantid_img2img import StableDiffusionXLInstantIDImg2ImgPipeline, draw_kps
+from controlnet_aux import ZoeDetector
+
 with open("sdxl_loras.json", "r") as file:
     data = json.load(file)
     sdxl_loras_raw = [
@@ -52,16 +68,60 @@ sdxl_loras_raw_new = [item for item in sdxl_loras_raw if item.get("new") == True
 
 sdxl_loras_raw = [item for item in sdxl_loras_raw if item.get("new") != True]
     
-lcm_lora_id = "lcm-sd/lcm-sdxl-base-1.0-lora"
+# download models
+hf_hub_download(
+    repo_id="InstantX/InstantID",
+    filename="ControlNetModel/config.json",
+    local_dir="/data/checkpoints",
+)
+hf_hub_download(
+    repo_id="InstantX/InstantID",
+    filename="ControlNetModel/diffusion_pytorch_model.safetensors",
+    local_dir="/data/checkpoints",
+)
+hf_hub_download(
+    repo_id="InstantX/InstantID", filename="ip-adapter.bin", local_dir="/data/checkpoints"
+)
+hf_hub_download(
+    repo_id="latent-consistency/lcm-lora-sdxl",
+    filename="pytorch_lora_weights.safetensors",
+    local_dir="/data/checkpoints",
+)
+# download antelopev2
+gdown.download(url="https://drive.google.com/file/d/18wEUfMNohBJ4K3Ly5wpTejPfDzp-8fI8/view?usp=sharing", output=".", quiet=False, fuzzy=True)
+# unzip antelopev2.zip
+os.system("unzip .antelopev2.zip -d /data/models/")
+
+app = FaceAnalysis(name='antelopev2', root='./data', providers=['CPUExecutionProvider'])
+app.prepare(ctx_id=0, det_size=(640, 640))
+
+# prepare models under ./checkpoints
+face_adapter = f'/data/checkpoints/ip-adapter.bin'
+controlnet_path = f'/data/checkpoints/ControlNetModel'
+
+# load IdentityNet
+identitynet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16)
+zoedepthnet = ControlNetModel.from_pretrained("diffusers/controlnet-zoe-depth-sdxl-1.0",torch_dtype=torch.float16)
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained("rubbrband/albedobaseXL_v21",
+                                                                 vae=vae,
+                                                                 controlnet=[identitynet, zoedepthnet],
+                                                                 torch_dtype=torch.float16)
+pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)
+pipe.load_ip_adapter_instantid(face_adapter)
+pipe.set_ip_adapter_scale(0.8)
+zoe = ZoeDetector.from_pretrained("lllyasviel/Annotators")
+zoe.to("cuda")
 
 vae = AutoencoderKL.from_pretrained(
     "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
 )
-pipe = StableDiffusionXLPipeline.from_pretrained(
+pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
     vae=vae,
     torch_dtype=torch.float16,
 )
+
 original_pipe = copy.deepcopy(pipe)
 pipe.to(device)
 
@@ -162,9 +222,22 @@ def merge_incompatible_lora(full_path_lora, lora_scale):
                 del lora_model
                 gc.collect()
 
-def run_lora(prompt, negative, lora_scale, selected_state, sdxl_loras, sdxl_loras_new, progress=gr.Progress(track_tqdm=True)):
+def run_lora(face_image, prompt, negative, lora_scale, selected_state, sdxl_loras, sdxl_loras_new, progress=gr.Progress(track_tqdm=True)):
     global last_lora, last_merged, last_fused, pipe
-    print("Index when running ", selected_state.index)
+    
+    face_info = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
+    face_info = sorted(face_info, key=lambda x:(x['bbox'][2]-x['bbox'][0])*x['bbox'][3]-x['bbox'][1])[-1] # only use the maximum face
+    face_emb = face_info['embedding']
+    face_kps = draw_kps(face_image, face_info['kps'])
+    
+    #prepare face zoe
+    with torch.no_grad():
+        image_zoe = zoe(face_image)
+    
+    width, height = face_kps.size
+    images = [face_kps, image_zoe.resize((height, width))]
+    
+    
     if(selected_state.index < 0):
         if(selected_state.index == -9999):
             selected_state.index = 0
@@ -194,19 +267,22 @@ def run_lora(prompt, negative, lora_scale, selected_state, sdxl_loras, sdxl_lora
         if(is_pivotal):
             #Add the textual inversion embeddings from pivotal tuning models
             text_embedding_name = sdxl_loras[selected_state.index]["text_embedding_weights"]
-            text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
-            tokenizers = [pipe.tokenizer, pipe.tokenizer_2]
-            embedding_path = hf_hub_download(repo_id=repo_name, filename=text_embedding_name, repo_type="model")
-            embhandler = TokenEmbeddingsHandler(text_encoders, tokenizers)
-            embhandler.load_embeddings(embedding_path)
+            state_dict_embedding = load_file(text_embedding_name)
+            pipe.load_textual_inversion(state_dict_embedding["clip_l"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder, tokenizer=pipe.tokenizer)
+            pipe.load_textual_inversion(state_dict_embedding["clip_g"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder_2, tokenizer=pipe.tokenizer_2)
             
     image = pipe(
         prompt=prompt,
         negative_prompt=negative,
         width=1024,
         height=1024,
+        image_embeds=face_emb,
+        image=face_image,
+        strength=0.85,
+        control_image=images,
         num_inference_steps=20,
-        guidance_scale=7.5,
+        guidance_scale = 7,
+        controlnet_conditioning_scale=[0.8, 0.8],
     ).images[0]
     last_lora = repo_name
     gc.collect()
@@ -332,7 +408,7 @@ with gr.Blocks(css="custom.css") as demo:
         show_progress=False
     ).success(
         fn=run_lora,
-        inputs=[prompt, negative, weight, selected_state, gr_sdxl_loras, gr_sdxl_loras_new],
+        inputs=[photo, prompt, negative, weight, selected_state, gr_sdxl_loras, gr_sdxl_loras_new],
         outputs=[result, share_group],
     )
     button.click(