Update app.py

app.py

@@ -22,7 +22,7 @@ import pandas as pd
 import json
 import requests
 from io import BytesIO
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, HfApi
 
 
 def resize_img(input_image, max_side=1280, min_side=1024, size=None, 
@@ -113,14 +113,14 @@ def process_image_by_bbox_larger(input_image, bbox_xyxy, min_bbox_ratio=0.2):
 
     return resized_image
 
-def calc_emb_cropped(image, app):
+def calc_emb_cropped(image, app, min_bbox_ratio=0.2):
     face_image = image.copy()
 
     face_info = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
 
     face_info = face_info[0]
 
-    cropped_face_image = process_image_by_bbox_larger(face_image, face_info["bbox"], min_bbox_ratio=0.2)
+    cropped_face_image = process_image_by_bbox_larger(face_image, face_info["bbox"], min_bbox_ratio=min_bbox_ratio)
 
     return cropped_face_image
 
@@ -144,7 +144,29 @@ default_negative_prompt = "Logo,Watermark,Text,Ugly,Morbid,Extra fingers,Poorly
 CURRENT_LORA_NAME = None
 
 # Load face detection and recognition package
-app = FaceAnalysis(name='antelopev2', root='./', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+# app = FaceAnalysis(name='antelopev2', root='./', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+
+
+# Define your repository, target directory, and local directory
+repo_id = "briaai/ID_preservation_2.3"
+target_dir = "models_aura/"  # The directory you want to download
+local_dir = "./checkpoints"  # Local directory to save files
+
+# Initialize the API
+api = HfApi()
+# List all files in the repository
+files = api.list_repo_files(repo_id=repo_id, repo_type="model")  # Use repo_type="space" for Spaces
+
+# Filter files that are in the target directory
+files_in_dir = [file for file in files if file.startswith(target_dir)]
+# Download each file in the target directory
+for file in files_in_dir:
+    local_path = os.path.join(local_dir, file)
+    os.makedirs(os.path.dirname(local_path), exist_ok=True)  # Ensure local directories exist
+    print(f"Downloading: {file}")
+    hf_hub_download(repo_id=repo_id, filename=file, local_dir=os.path.dirname(local_path))
+
+app = FaceAnalysis(name='auraface', root='./checkpoints/models_aura/', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
 app.prepare(ctx_id=0, det_size=(640, 640))
 
 
@@ -155,19 +177,14 @@ hf_hub_download(repo_id="briaai/ID_preservation_2.3", filename="controlnet/diffu
 hf_hub_download(repo_id="briaai/ID_preservation_2.3", filename="ip-adapter.bin", local_dir="./checkpoints")
 hf_hub_download(repo_id="briaai/ID_preservation_2.3", filename="image_encoder/pytorch_model.bin", local_dir="./checkpoints")
 hf_hub_download(repo_id="briaai/ID_preservation_2.3", filename="image_encoder/config.json", local_dir="./checkpoints")
-# Download Lora weights
-hf_hub_download(repo_id="briaai/ID_preservation_2.3", filename="LoRAs/3D_illustration/pytorch_lora_weights.safetensors", local_dir="./checkpoints")
-hf_hub_download(repo_id="briaai/ID_preservation_2.3", filename="LoRAs/Avatar_internlm/pytorch_lora_weights.safetensors", local_dir="./checkpoints")
-hf_hub_download(repo_id="briaai/ID_preservation_2.3", filename="LoRAs/Characters/pytorch_lora_weights.safetensors", local_dir="./checkpoints")
-hf_hub_download(repo_id="briaai/ID_preservation_2.3", filename="LoRAs/Storyboards/pytorch_lora_weights.safetensors", local_dir="./checkpoints")
-hf_hub_download(repo_id="briaai/ID_preservation_2.3", filename="LoRAs/Vangogh_Vanilla/pytorch_lora_weights.safetensors", local_dir="./checkpoints")
+
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # ckpts paths
 face_adapter = f"./checkpoints/ip-adapter.bin"
 controlnet_path = f"./checkpoints/controlnet"
-lora_base_path = "./checkpoints/LoRAs"
+# lora_base_path = "./checkpoints/LoRAs"
 base_model_path = f'briaai/BRIA-2.3'
 resolution = 1024
 
@@ -199,16 +216,16 @@ pipe.load_ip_adapter_instantid(face_adapter)
 
 clip_embeds=None
 
-Loras_dict = {
-    "":"",
-    "Vangogh_Vanilla": "bold, dramatic brush strokes, vibrant colors, swirling patterns, intense, emotionally charged paintings of",
-    "Avatar_internlm": "2d anime sketch avatar of",
-    "Storyboards": "Illustration style for storyboarding.",
-    "3D_illustration": "3D object illustration, abstract.",
-    "Characters": "gaming vector Art."
-    }
+# Loras_dict = {
+#     "":"",
+#     "Vangogh_Vanilla": "bold, dramatic brush strokes, vibrant colors, swirling patterns, intense, emotionally charged paintings of",
+#     "Avatar_internlm": "2d anime sketch avatar of",
+#     "Storyboards": "Illustration style for storyboarding.",
+#     "3D_illustration": "3D object illustration, abstract.",
+#     "Characters": "gaming vector Art."
+#     }
 
-lora_names = Loras_dict.keys()
+# lora_names = Loras_dict.keys()
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -217,8 +234,9 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
 
 
 @spaces.GPU
-def generate_image(image_path, prompt, num_steps, guidance_scale, seed, num_images, ip_adapter_scale, kps_scale, canny_scale, lora_name, lora_scale, progress=gr.Progress(track_tqdm=True)):
-    global CURRENT_LORA_NAME  # Use the global variable to track LoRA
+# def generate_image(image_path, prompt, num_steps, guidance_scale, seed, num_images, ip_adapter_scale, kps_scale, canny_scale, lora_name, lora_scale, progress=gr.Progress(track_tqdm=True)):
+def generate_image(image_path, prompt, num_steps, guidance_scale, seed, num_images, ip_adapter_scale, kps_scale, canny_scale, progress=gr.Progress(track_tqdm=True)):
+    # global CURRENT_LORA_NAME  # Use the global variable to track LoRA
 
     if image_path is None:
         raise gr.Error(f"Cannot find any input face image! Please upload a face image.")
@@ -272,31 +290,31 @@ def generate_image(image_path, prompt, num_steps, guidance_scale, seed, num_imag
                     
     generator = torch.Generator(device=device).manual_seed(seed)
     
-    if lora_name != CURRENT_LORA_NAME:  # Check if LoRA needs to be changed
-        if CURRENT_LORA_NAME is not None:  # If a LoRA is already loaded, unload it
-            pipe.disable_lora()
-            pipe.unfuse_lora()
-            pipe.unload_lora_weights()
-            print(f"Unloaded LoRA: {CURRENT_LORA_NAME}")
+    # if lora_name != CURRENT_LORA_NAME:  # Check if LoRA needs to be changed
+    #     if CURRENT_LORA_NAME is not None:  # If a LoRA is already loaded, unload it
+    #         pipe.disable_lora()
+    #         pipe.unfuse_lora()
+    #         pipe.unload_lora_weights()
+    #         print(f"Unloaded LoRA: {CURRENT_LORA_NAME}")
 
-        if lora_name != "":  # Load the new LoRA if specified
-            # pipe.enable_model_cpu_offload()
-            lora_path = os.path.join(lora_base_path, lora_name, "pytorch_lora_weights.safetensors")
-            pipe.load_lora_weights(lora_path)
-            pipe.fuse_lora(lora_scale)  
-            pipe.enable_lora()
+    #     if lora_name != "":  # Load the new LoRA if specified
+    #         # pipe.enable_model_cpu_offload()
+    #         lora_path = os.path.join(lora_base_path, lora_name, "pytorch_lora_weights.safetensors")
+    #         pipe.load_lora_weights(lora_path)
+    #         pipe.fuse_lora(lora_scale)  
+    #         pipe.enable_lora()
 
-            # lora_prefix = Loras_dict[lora_name]
+    #         # lora_prefix = Loras_dict[lora_name]
 
-            print(f"Loaded new LoRA: {lora_name}")
+    #         print(f"Loaded new LoRA: {lora_name}")
 
-        # Update the current LoRA name
-        CURRENT_LORA_NAME = lora_name
+    #     # Update the current LoRA name
+    #     CURRENT_LORA_NAME = lora_name
 
-    if lora_name != "":
-        full_prompt = f"{Loras_dict[lora_name]} + " " + {prompt}"
-    else:
-        full_prompt = prompt
+    # if lora_name != "":
+    #     full_prompt = f"{Loras_dict[lora_name]} + " " + {prompt}"
+    # else:
+    full_prompt = prompt
 
     print("Start inference...")    
     images = pipe(
@@ -360,7 +378,7 @@ with gr.Blocks(css=css) as demo:
             info="Describe what you want to generate or modify in the image."
             )
             
-            lora_name = gr.Dropdown(choices=lora_names, label="LoRA", value="",  info="Select a LoRA name from the list, not selecting any will disable LoRA.")
+            # lora_name = gr.Dropdown(choices=lora_names, label="LoRA", value="",  info="Select a LoRA name from the list, not selecting any will disable LoRA.")
 
             submit = gr.Button("Submit", variant="primary")
  
@@ -407,13 +425,13 @@ with gr.Blocks(css=css) as demo:
                     step=0.01,
                     value=0.4,
                 )
-                lora_scale = gr.Slider(
-                    label="lora_scale",
-                    minimum=0.0,
-                    maximum=1.0,
-                    step=0.01,
-                    value=0.7,
-                )
+                # lora_scale = gr.Slider(
+                #     label="lora_scale",
+                #     minimum=0.0,
+                #     maximum=1.0,
+                #     step=0.01,
+                #     value=0.7,
+                # )
                 seed = gr.Slider(
                     label="Seed",
                     minimum=0,