Update README.md

README.md

@@ -89,15 +89,16 @@ import os
 import random
 import gradio as gr
 
-
 import cv2
 import torch
 import numpy as np
 from PIL import Image
 
 from transformers import CLIPVisionModelWithProjection
-from diffusers.utils import load_image
 from diffusers.models import ControlNetModel
+
+from huggingface_hub import snapshot_download
+
 from insightface.app import FaceAnalysis
 
 import io
@@ -108,8 +109,8 @@ from pipeline_stable_diffusion_xl_instantid import StableDiffusionXLInstantIDPip
 import pandas as pd
 import json
 import requests
-from PIL import Image
 from io import BytesIO
+from huggingface_hub import hf_hub_download, HfApi
 
 
 def resize_img(input_image, max_side=1280, min_side=1024, size=None, 
@@ -152,33 +153,50 @@ def make_canny_condition(image, min_val=100, max_val=200, w_bilateral=True):
 
 default_negative_prompt = "Logo,Watermark,Text,Ugly,Morbid,Extra fingers,Poorly drawn hands,Mutation,Blurry,Extra limbs,Gross proportions,Missing arms,Mutated hands,Long neck,Duplicate,Mutilated,Mutilated hands,Poorly drawn face,Deformed,Bad anatomy,Cloned face,Malformed limbs,Missing legs,Too many fingers"
 
-# Load face detection and recognition package
-app = FaceAnalysis(name='antelopev2', root='./', providers=['CPUExecutionProvider'])
+# Download face encoder
+snapshot_download(
+    "fal/AuraFace-v1",
+    local_dir="models/auraface",
+)
+
+app = FaceAnalysis(
+    name="auraface",
+    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+    root=".",
+)
+
 app.prepare(ctx_id=0, det_size=(640, 640))
 
 
+# download checkpoints
+print("Downloading checkpoints")
+hf_hub_download(repo_id="briaai/ID_preservation_2.3_auraFaceEnc", filename="checkpoint_105000/controlnet/config.json", local_dir="./checkpoints")
+hf_hub_download(repo_id="briaai/ID_preservation_2.3_auraFaceEnc", filename="checkpoint_105000/controlnet/diffusion_pytorch_model.safetensors", local_dir="./checkpoints")
+hf_hub_download(repo_id="briaai/ID_preservation_2.3_auraFaceEnc", filename="checkpoint_105000/ip-adapter.bin", local_dir="./checkpoints")
+hf_hub_download(repo_id="briaai/ID_preservation_2.3_auraFaceEnc", filename="image_encoder/pytorch_model.bin", local_dir="./checkpoints")
+hf_hub_download(repo_id="briaai/ID_preservation_2.3_auraFaceEnc", filename="image_encoder/config.json", local_dir="./checkpoints")
 
-face_adapter = f"./ip-adapter.bin"
-controlnet_path = f"./controlnet"
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# ckpts paths
+face_adapter = f"./checkpoints/checkpoint_105000/ip-adapter.bin"
+controlnet_path = f"./checkpoints/checkpoint_105000/controlnet"
 base_model_path = f'briaai/BRIA-2.3'
 resolution = 1024
 
+# Load ControlNet models
 controlnet_lnmks = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16)
-
 controlnet_canny = ControlNetModel.from_pretrained("briaai/BRIA-2.3-ControlNet-Canny",
-                                                   torch_dtype=torch.float16)
-
+                                                torch_dtype=torch.float16)
+    
 controlnet = [controlnet_lnmks, controlnet_canny]
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
 
 image_encoder = CLIPVisionModelWithProjection.from_pretrained(
         f"./checkpoints/image_encoder",
         torch_dtype=torch.float16,
     )
-
 pipe = StableDiffusionXLInstantIDPipeline.from_pretrained(
         base_model_path,
         controlnet=controlnet,