Update app.py

app.py

@@ -38,7 +38,9 @@ from compel import Compel, ReturnedEmbeddingsType
 
 from gradio_imageslider import ImageSlider
 
-# Load LoRA configurations
+
+#from gradio_imageslider import ImageSlider
+
 with open("sdxl_loras.json", "r") as file:
     data = json.load(file)
     sdxl_loras_raw = [
@@ -61,11 +63,12 @@ with open("sdxl_loras.json", "r") as file:
 
 with open("defaults_data.json", "r") as file:
     lora_defaults = json.load(file)
+    
 
-device = "cuda"
+device = "cuda" 
 
-# Cache for LoRA state dicts
 state_dicts = {}
+
 for item in sdxl_loras_raw:
     saved_name = hf_hub_download(item["repo"], item["weights"])
     
@@ -80,8 +83,8 @@ for item in sdxl_loras_raw:
     }
 
 sdxl_loras_raw = [item for item in sdxl_loras_raw if item.get("new") != True]
-
-# Download models
+    
+# download models
 hf_hub_download(
     repo_id="InstantX/InstantID",
     filename="ControlNetModel/config.json",
@@ -100,158 +103,70 @@ hf_hub_download(
     filename="pytorch_lora_weights.safetensors",
     local_dir="/data/checkpoints",
 )
+# download antelopev2
+#if not os.path.exists("/data/antelopev2.zip"):
+#    gdown.download(url="https://drive.google.com/file/d/18wEUfMNohBJ4K3Ly5wpTejPfDzp-8fI8/view?usp=sharing", output="/data/", quiet=False, fuzzy=True)
+#    os.system("unzip /data/antelopev2.zip -d /data/models/")
 
-# Download antelopev2
 antelope_download = snapshot_download(repo_id="DIAMONIK7777/antelopev2", local_dir="/data/models/antelopev2")
 print(antelope_download)
 app = FaceAnalysis(name='antelopev2', root='/data', providers=['CPUExecutionProvider'])
-app.prepare(ctx_id=0, det_size=(768, 768))
+app.prepare(ctx_id=0, det_size=(640, 640))
 
-# Prepare models
+# prepare models under ./checkpoints
 face_adapter = f'/data/checkpoints/ip-adapter.bin'
 controlnet_path = f'/data/checkpoints/ControlNetModel'
 
+# load IdentityNet
 st = time.time()
 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)
+zoedepthnet = ControlNetModel.from_pretrained("diffusers/controlnet-zoe-depth-sdxl-1.0",torch_dtype=torch.float16)
 et = time.time()
-print('Loading ControlNet took: ', et - st, 'seconds')
-
+elapsed_time = et - st
+print('Loading ControlNet took: ', elapsed_time, 'seconds')
 st = time.time()
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 et = time.time()
-print('Loading VAE took: ', et - st, 'seconds')
-
+elapsed_time = et - st
+print('Loading VAE took: ', elapsed_time, 'seconds')
 st = time.time()
-pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained(
-    "SG161222/RealVisXL_V5.0", 
-    vae=vae,
-    controlnet=[identitynet, zoedepthnet],
-    torch_dtype=torch.float16
-)
 
+#pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained("stablediffusionapi/albedobase-xl-v21",
+pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained("frankjoshua/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.9)
+pipe.set_ip_adapter_scale(0.8)
 et = time.time()
-print('Loading pipeline took: ', et - st, 'seconds')
-
+elapsed_time = et - st
+print('Loading pipeline took: ', elapsed_time, 'seconds')
 st = time.time()
-compel = Compel(
-    tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
-    text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
-    returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
-    requires_pooled=[False, True]
-)
+compel = Compel(tokenizer=[pipe.tokenizer, pipe.tokenizer_2] , text_encoder=[pipe.text_encoder, pipe.text_encoder_2], returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, requires_pooled=[False, True])
 et = time.time()
-print('Loading Compel took: ', et - st, 'seconds')
+elapsed_time = et - st
+print('Loading Compel took: ', elapsed_time, 'seconds')
 
 st = time.time()
 zoe = ZoeDetector.from_pretrained("lllyasviel/Annotators")
 et = time.time()
-print('Loading Zoe took: ', et - st, 'seconds')
+elapsed_time = et - st
+print('Loading Zoe took: ', elapsed_time, 'seconds')
 zoe.to(device)
 pipe.to(device)
 
 last_lora = ""
 last_fused = False
+js = '''
+var button = document.getElementById('button');
+// Add a click event listener to the button
+button.addEventListener('click', function() {
+    element.classList.add('selected');
+});
+'''
 lora_archive = "/data"
 
-# Improved face detection with multi-face support
-def detect_faces(face_image, use_multiple_faces=False):
-    """
-    Detect faces in the image
-    Returns: list of face info dictionaries, or empty list if no faces
-    """
-    try:
-        face_info_list = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
-        
-        if not face_info_list or len(face_info_list) == 0:
-            print("No faces detected")
-            return []
-        
-        # Sort faces by size (largest first)
-        face_info_list = sorted(
-            face_info_list, 
-            key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]),
-            reverse=True
-        )
-        
-        if use_multiple_faces:
-            print(f"Detected {len(face_info_list)} faces")
-            return face_info_list
-        else:
-            print(f"Using largest face (detected {len(face_info_list)} total)")
-            return [face_info_list[0]]
-            
-    except Exception as e:
-        print(f"Face detection error: {e}")
-        return []
-
-def process_face_embeddings(face_info_list):
-    """
-    Process face embeddings - average multiple faces or return single face
-    """
-    if not face_info_list:
-        return None
-    
-    if len(face_info_list) == 1:
-        return face_info_list[0]['embedding']
-    
-    # Average embeddings for multiple faces
-    embeddings = [face_info['embedding'] for face_info in face_info_list]
-    avg_embedding = np.mean(embeddings, axis=0)
-    return avg_embedding
-
-def create_face_kps_image(face_image, face_info_list):
-    """
-    Create keypoints image from face info
-    """
-    if not face_info_list:
-        return face_image
-    
-    # For multiple faces, draw all keypoints
-    if len(face_info_list) > 1:
-        return draw_multiple_kps(face_image, [f['kps'] for f in face_info_list])
-    else:
-        return draw_kps(face_image, face_info_list[0]['kps'])
-
-def draw_multiple_kps(image_pil, kps_list, color_list=[(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255)]):
-    """
-    Draw keypoints for multiple faces
-    """
-    stickwidth = 4
-    limbSeq = np.array([[0, 2], [1, 2], [3, 2], [4, 2]])
-    
-    w, h = image_pil.size
-    out_img = np.zeros([h, w, 3])
-    
-    for kps in kps_list:
-        kps = np.array(kps)
-        
-        for i in range(len(limbSeq)):
-            index = limbSeq[i]
-            color = color_list[index[0]]
-            
-            x = kps[index][:, 0]
-            y = kps[index][:, 1]
-            length = ((x[0] - x[1]) ** 2 + (y[0] - y[1]) ** 2) ** 0.5
-            angle = np.degrees(np.arctan2(y[0] - y[1], x[0] - x[1]))
-            polygon = cv2.ellipse2Poly(
-                (int(np.mean(x)), int(np.mean(y))), (int(length / 2), stickwidth), int(angle), 0, 360, 1
-            )
-            out_img = cv2.fillConvexPoly(out_img.copy(), polygon, color)
-        
-        out_img = (out_img * 0.6).astype(np.uint8)
-        
-        for idx_kp, kp in enumerate(kps):
-            color = color_list[idx_kp]
-            x, y = kp
-            out_img = cv2.circle(out_img.copy(), (int(x), int(y)), 10, color, -1)
-    
-    out_img_pil = Image.fromarray(out_img.astype(np.uint8))
-    return out_img_pil
-
 def update_selection(selected_state: gr.SelectData, sdxl_loras, face_strength, image_strength, weight, depth_control_scale, negative, is_new=False):
     lora_repo = sdxl_loras[selected_state.index]["repo"]
     new_placeholder = "Type a prompt to use your selected LoRA"
@@ -260,9 +175,9 @@ def update_selection(selected_state: gr.SelectData, sdxl_loras, face_strength, i
 
     for lora_list in lora_defaults:
         if lora_list["model"] == sdxl_loras[selected_state.index]["repo"]:
-            face_strength = lora_list.get("face_strength", 0.9)
-            image_strength = lora_list.get("image_strength", 0.2)
-            weight = lora_list.get("weight", 0.95)
+            face_strength = lora_list.get("face_strength", 0.85)
+            image_strength = lora_list.get("image_strength", 0.15)
+            weight = lora_list.get("weight", 0.9)
             depth_control_scale = lora_list.get("depth_control_scale", 0.8)
             negative = lora_list.get("negative", "")
     
@@ -283,82 +198,157 @@ def update_selection(selected_state: gr.SelectData, sdxl_loras, face_strength, i
         selected_state
     )
 
+def center_crop_image_as_square(img):
+    square_size = min(img.size)
+    
+    left = (img.width - square_size) / 2
+    top = (img.height - square_size) / 2
+    right = (img.width + square_size) / 2
+    bottom = (img.height + square_size) / 2
+    
+    img_cropped = img.crop((left, top, right, bottom))
+    return img_cropped
+    
 def check_selected(selected_state, custom_lora):
     if not selected_state and not custom_lora:
         raise gr.Error("You must select a style")
 
-def resize_image_aspect_ratio(img, max_dim=1280):
-    width, height = img.size
-    aspect_ratio = width / height
-
-    if aspect_ratio >= 1:  # Landscape or square
-        new_width = min(max_dim, width)
-        new_height = int(new_width / aspect_ratio)
-    else:  # Portrait
-        new_height = min(max_dim, height)
-        new_width = int(new_height * aspect_ratio)
-
-    new_width = (new_width // 8) * 8
-    new_height = (new_height // 8) * 8
-
-    return img.resize((new_width, new_height), Image.LANCZOS)
-
-
-def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_strength, image_strength, 
-             guidance_scale, depth_control_scale, sdxl_loras, custom_lora, use_multiple_faces=False, 
-             progress=gr.Progress(track_tqdm=True)):
-    """
-    Enhanced run_lora with support for:
-    - No faces (landscape mode)
-    - Multiple faces
-    - Improved results
-    """
-    print("Custom LoRA:", custom_lora)
-    custom_lora_path = custom_lora[0] if custom_lora else None
-    selected_state_index = selected_state.index if selected_state else -1
+def merge_incompatible_lora(full_path_lora, lora_scale):
+    for weights_file in [full_path_lora]:
+                if ";" in weights_file:
+                    weights_file, multiplier = weights_file.split(";")
+                    multiplier = float(multiplier)
+                else:
+                    multiplier = lora_scale
+
+                lora_model, weights_sd = lora.create_network_from_weights(
+                    multiplier,
+                    full_path_lora,
+                    pipe.vae,
+                    pipe.text_encoder,
+                    pipe.unet,
+                    for_inference=True,
+                )
+                lora_model.merge_to(
+                    pipe.text_encoder, pipe.unet, weights_sd, torch.float16, "cuda"
+                )
+                del weights_sd
+                del lora_model
 
+@spaces.GPU(duration=100)
+def generate_image(prompt, negative, face_emb, face_image, face_kps, image_strength, guidance_scale, face_strength, depth_control_scale, repo_name, loaded_state_dict, lora_scale, sdxl_loras, selected_state_index, st):
+    print(loaded_state_dict)
+    et = time.time()
+    elapsed_time = et - st
+    print('Getting into the decorated function took: ', elapsed_time, 'seconds')
+    global last_fused, last_lora
+    print("Last LoRA: ", last_lora)
+    print("Current LoRA: ", repo_name)
+    print("Last fused: ", last_fused)
+    #prepare face zoe
     st = time.time()
-    face_image = resize_image_aspect_ratio(face_image)
-
-    # Enhanced face detection
-    face_info_list = detect_faces(face_image, use_multiple_faces)
-    face_detected = len(face_info_list) > 0
-    
-    if face_detected:
-        face_emb = process_face_embeddings(face_info_list)
-        face_kps = create_face_kps_image(face_image, face_info_list)
-        print(f"Processing with {len(face_info_list)} face(s)")
+    with torch.no_grad():
+        image_zoe = zoe(face_image)
+    width, height = face_kps.size
+    images = [face_kps, image_zoe.resize((height, width))]
+    et = time.time()
+    elapsed_time = et - st
+    print('Zoe Depth calculations took: ', elapsed_time, 'seconds')
+    if last_lora != repo_name:
+        if(last_fused):
+            st = time.time()
+            pipe.unfuse_lora()
+            pipe.unload_lora_weights()
+            pipe.unload_textual_inversion()
+            et = time.time()
+            elapsed_time = et - st
+            print('Unfuse and unload LoRA took: ', elapsed_time, 'seconds')
+        st = time.time()
+        pipe.load_lora_weights(loaded_state_dict)
+        pipe.fuse_lora(lora_scale)
+        et = time.time()
+        elapsed_time = et - st
+        print('Fuse and load LoRA took: ', elapsed_time, 'seconds')
+        last_fused = True
+        is_pivotal = sdxl_loras[selected_state_index]["is_pivotal"]
+        if(is_pivotal):
+            #Add the textual inversion embeddings from pivotal tuning models
+            text_embedding_name = sdxl_loras[selected_state_index]["text_embedding_weights"]
+            embedding_path = hf_hub_download(repo_id=repo_name, filename=text_embedding_name, repo_type="model")
+            state_dict_embedding = load_file(embedding_path)
+            pipe.load_textual_inversion(state_dict_embedding["clip_l" if "clip_l" in state_dict_embedding else "text_encoders_0"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder, tokenizer=pipe.tokenizer)
+            pipe.load_textual_inversion(state_dict_embedding["clip_g" if "clip_g" in state_dict_embedding else "text_encoders_1"], token=["<s0>", "<s1>"], text_encoder=pipe.text_encoder_2, tokenizer=pipe.tokenizer_2)
+            
+    print("Processing prompt...")
+    st = time.time()
+    conditioning, pooled = compel(prompt)
+    if(negative):
+        negative_conditioning, negative_pooled = compel(negative)
     else:
-        face_emb = None
-        face_kps = face_image
-        print("No faces detected - using landscape/depth mode only")
-
+        negative_conditioning, negative_pooled = None, None
+    et = time.time()
+    elapsed_time = et - st
+    print('Prompt processing took: ', elapsed_time, 'seconds')
+    print("Processing image...")
+    st = time.time()
+    image = pipe(
+        prompt_embeds=conditioning,
+        pooled_prompt_embeds=pooled,
+        negative_prompt_embeds=negative_conditioning,
+        negative_pooled_prompt_embeds=negative_pooled,
+        width=1024,
+        height=1024,
+        image_embeds=face_emb,
+        image=face_image,
+        strength=1-image_strength,
+        control_image=images,
+        num_inference_steps=36,
+        guidance_scale = guidance_scale,
+        controlnet_conditioning_scale=[face_strength, depth_control_scale],
+    ).images[0]
+    et = time.time()
+    elapsed_time = et - st
+    print('Image processing took: ', elapsed_time, 'seconds')
+    last_lora = repo_name
+    return image
+    
+def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, sdxl_loras, custom_lora, progress=gr.Progress(track_tqdm=True)):
+    print("Custom LoRA: ", custom_lora)
+    custom_lora_path = custom_lora[0] if custom_lora else None
+    selected_state_index = selected_state.index if selected_state else -1
+    st = time.time()
+    face_image = center_crop_image_as_square(face_image)
+    try:
+        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'])
+    except: 
+        raise gr.Error("No face found in your image. Only face images work here. Try again")
     et = time.time()
-    print('Face processing took:', et - st, 'seconds')
+    elapsed_time = et - st
+    print('Cropping and calculating face embeds took: ', elapsed_time, 'seconds')
 
     st = time.time()
 
-    # Enhanced prompt processing
-    if custom_lora_path and custom_lora[1]:
+    if(custom_lora_path and custom_lora[1]):
         prompt = f"{prompt} {custom_lora[1]}"
     else:
         for lora_list in lora_defaults:
             if lora_list["model"] == sdxl_loras[selected_state_index]["repo"]:
                 prompt_full = lora_list.get("prompt", None)
-                if prompt_full:
+                if(prompt_full):
                     prompt = prompt_full.replace("<subject>", prompt)
 
-    print("Prompt:", prompt)
-    if prompt == "":
-        prompt = "a beautiful scene" if not face_detected else "a person"
+    print("Prompt:", prompt) 
+    if(prompt == ""):
+        prompt = "a person"
     print(f"Executing prompt: {prompt}")
-
+    #print("Selected State: ", selected_state_index)
+    #print(sdxl_loras[selected_state_index]["repo"])
     if negative == "":
-        # Enhanced negative prompt for better quality
-        negative = "worst quality, low quality, blurry, distorted, deformed" if not face_detected else None
-    
-    print("Custom Loaded LoRA:", custom_lora_path)
-
+        negative = None
+    print("Custom Loaded LoRA: ", custom_lora_path)
     if not selected_state and not custom_lora_path:
         raise gr.Error("You must select a style")
     elif custom_lora_path:
@@ -366,130 +356,21 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
         full_path_lora = custom_lora_path
     else:
         repo_name = sdxl_loras[selected_state_index]["repo"]
+        weight_name = sdxl_loras[selected_state_index]["weights"]
         full_path_lora = state_dicts[repo_name]["saved_name"]
-
-    repo_name = repo_name.rstrip("/").lower()
-
-    print("Full path LoRA", full_path_lora)
-
+    print("Full path LoRA ", full_path_lora)
+    #loaded_state_dict = copy.deepcopy(state_dicts[repo_name]["state_dict"])
+    cross_attention_kwargs = None
     et = time.time()
-    print('Prompt processing took:', et - st, 'seconds')
-
-    # Adjust parameters based on face detection
-    if not face_detected:
-        # For landscape/no face mode, reduce face strength and increase depth control
-        face_strength = 0.0
-        depth_control_scale = max(depth_control_scale, 0.9)
-        image_strength = min(image_strength, 0.4)
-        print("Adjusted parameters for no-face mode")
+    elapsed_time = et - st
+    print('Small content processing took: ', elapsed_time, 'seconds')
 
     st = time.time()
-    image = generate_image(
-        prompt, negative, face_emb, face_image, face_kps, image_strength,
-        guidance_scale, face_strength, depth_control_scale, repo_name,
-        full_path_lora, lora_scale, sdxl_loras, selected_state_index, face_detected, st
-    )
-
+    image = generate_image(prompt, negative, face_emb, face_image, face_kps, image_strength, guidance_scale, face_strength, depth_control_scale, repo_name, full_path_lora, lora_scale, sdxl_loras, selected_state_index, st)
     return (face_image, image), gr.update(visible=True)
 
 run_lora.zerogpu = True
 
-
-@spaces.GPU(duration=75)
-def generate_image(prompt, negative, face_emb, face_image, face_kps, image_strength, guidance_scale, 
-                   face_strength, depth_control_scale, repo_name, loaded_state_dict, lora_scale, 
-                   sdxl_loras, selected_state_index, face_detected, st):
-    global last_fused, last_lora
-
-    print("Loaded state dict:", loaded_state_dict)
-    print("Last LoRA:", last_lora, "| Current LoRA:", repo_name)
-
-    # Prepare control images and scales based on face detection
-    if face_detected:
-        control_images = [face_kps, zoe(face_image)]
-        control_scales = [face_strength, depth_control_scale]
-    else:
-        # Only use depth control for landscapes
-        control_images = [zoe(face_image)]
-        control_scales = [depth_control_scale]
-
-    # Handle custom LoRA from HuggingFace
-    if repo_name.startswith("https://huggingface.co"):
-        repo_id = repo_name.split("huggingface.co/")[-1]
-        fs = HfFileSystem()
-        files = fs.ls(repo_id, detail=False)
-        safetensors_files = [f for f in files if f.endswith(".safetensors")]
-
-        if not safetensors_files:
-            raise gr.Error("No .safetensors file found in this Hugging Face repository.")
-
-        weight_file = safetensors_files[0]
-        full_path_lora = hf_hub_download(repo_id=repo_id, filename=weight_file, repo_type="model")
-    else:
-        full_path_lora = loaded_state_dict
-
-    # Improved LoRA loading and caching
-    if last_lora != repo_name:
-        if last_fused:
-            pipe.unfuse_lora()
-            pipe.unload_lora_weights()
-            pipe.unload_textual_inversion()
-
-        # Load LoRA with better error handling
-        try:
-            pipe.load_lora_weights(full_path_lora)
-            pipe.fuse_lora(lora_scale)
-            last_fused = True
-            
-            # Handle pivotal tuning embeddings
-            is_pivotal = sdxl_loras[selected_state_index]["is_pivotal"]
-            if is_pivotal:
-                text_embedding_name = sdxl_loras[selected_state_index]["text_embedding_weights"]
-                embedding_path = hf_hub_download(repo_id=repo_name, filename=text_embedding_name, repo_type="model")
-                state_dict_embedding = load_file(embedding_path)
-                pipe.load_textual_inversion(
-                    state_dict_embedding["clip_l" if "clip_l" in state_dict_embedding else "text_encoders_0"],
-                    token=["<s0>", "<s1>"],
-                    text_encoder=pipe.text_encoder,
-                    tokenizer=pipe.tokenizer
-                )
-                pipe.load_textual_inversion(
-                    state_dict_embedding["clip_g" if "clip_g" in state_dict_embedding else "text_encoders_1"],
-                    token=["<s0>", "<s1>"],
-                    text_encoder=pipe.text_encoder_2,
-                    tokenizer=pipe.tokenizer_2
-                )
-        except Exception as e:
-            print(f"Error loading LoRA: {e}")
-            raise gr.Error(f"Failed to load LoRA: {str(e)}")
-
-    print("Processing prompt...")
-    conditioning, pooled = compel(prompt)
-    negative_conditioning, negative_pooled = compel(negative) if negative else (None, None)
-
-    # Enhanced generation parameters
-    num_inference_steps = 40  # Increased for better quality
-    
-    print("Generating image...")
-    image = pipe(
-        prompt_embeds=conditioning,
-        pooled_prompt_embeds=pooled,
-        negative_prompt_embeds=negative_conditioning,
-        negative_pooled_prompt_embeds=negative_pooled,
-        width=face_image.width,
-        height=face_image.height,
-        image_embeds=face_emb if face_detected else None,
-        image=face_image,
-        strength=1-image_strength,
-        control_image=control_images,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=guidance_scale,
-        controlnet_conditioning_scale=control_scales,
-    ).images[0]
-
-    last_lora = repo_name
-    return image
-
 def shuffle_gallery(sdxl_loras):
     random.shuffle(sdxl_loras)
     return [(item["image"], item["title"]) for item in sdxl_loras], sdxl_loras
@@ -505,75 +386,75 @@ def swap_gallery(order, sdxl_loras):
         return classify_gallery(sdxl_loras)
         
 def deselect():
-    return gr.Gallery(selected_index=None)
+  return gr.Gallery(selected_index=None)
 
 def get_huggingface_safetensors(link):
-    split_link = link.split("/")
-    if(len(split_link) == 2):
-        model_card = ModelCard.load(link)
-        image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
-        trigger_word = model_card.data.get("instance_prompt", "")
-        image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
-        fs = HfFileSystem()
-        try:
-            list_of_files = fs.ls(link, detail=False)
-            for file in list_of_files:
-                if(file.endswith(".safetensors")):
-                    safetensors_name = file.replace("/", "_")
-                    if(not os.path.exists(f"{lora_archive}/{safetensors_name}")):
-                        fs.get_file(file, lpath=f"{lora_archive}/{safetensors_name}")
-                if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))):
-                    image_elements = file.split("/")
-                    image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
-        except:
-            gr.Warning("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
-            raise Exception("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
-        return split_link[1], f"{lora_archive}/{safetensors_name}", trigger_word, image_url
+  split_link = link.split("/")
+  if(len(split_link) == 2):
+            model_card = ModelCard.load(link)
+            image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
+            trigger_word = model_card.data.get("instance_prompt", "")
+            image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
+            fs = HfFileSystem()
+            try:
+                list_of_files = fs.ls(link, detail=False)
+                for file in list_of_files:
+                    if(file.endswith(".safetensors")):
+                        safetensors_name = file.replace("/", "_")
+                        if(not os.path.exists(f"{lora_archive}/{safetensors_name}")):
+                          fs.get_file(file, lpath=f"{lora_archive}/{safetensors_name}")
+                    if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))):
+                      image_elements = file.split("/")
+                      image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
+            except:
+              gr.Warning("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+              raise Exception("You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
+            return split_link[1], f"{lora_archive}/{safetensors_name}", trigger_word, image_url
 
 def get_civitai_safetensors(link):
-    link_split = link.split("civitai.com/")
-    pattern = re.compile(r'models\/(\d+)')
-    regex_match = pattern.search(link_split[1])
-    if(regex_match):
-        civitai_model_id = regex_match.group(1)
-    else:
-        gr.Warning("No CivitAI model id found in your URL")
-        raise Exception("No CivitAI model id found in your URL")
-    model_request_url = f"https://civitai.com/api/v1/models/{civitai_model_id}?token={os.getenv('CIVITAI_TOKEN')}"
-    x = requests.get(model_request_url)
-    if(x.status_code != 200):
-        raise Exception("Invalid CivitAI URL")
-    model_data = x.json()
-    
-    if(model_data["type"] != "LORA"):
-        gr.Warning("The model isn't tagged at CivitAI as a LoRA")
-        raise Exception("The model isn't tagged at CivitAI as a LoRA")
-    
-    model_link_download = None
-    image_url = None
-    trigger_word = ""
-    for model in model_data["modelVersions"]:
-        if(model["baseModel"] == "SDXL 1.0"):
-            model_link_download = f"{model['downloadUrl']}/?token={os.getenv('CIVITAI_TOKEN')}"
-            safetensors_name = model["files"][0]["name"]
-            if(not os.path.exists(f"{lora_archive}/{safetensors_name}")):
-                safetensors_file_request = requests.get(model_link_download)
-                if(safetensors_file_request.status_code != 200):
-                    raise Exception("Invalid CivitAI download link")
-                with open(f"{lora_archive}/{safetensors_name}", 'wb') as file:
-                    file.write(safetensors_file_request.content)
-            trigger_word = model.get("trainedWords", [""])[0]
-            for image in model["images"]:
-                if(image["nsfwLevel"] == 1):
-                    image_url = image["url"]
-                    break
-            break
-    
-    if(not model_link_download):
-        gr.Warning("We couldn't find a SDXL LoRA on the model you've sent")
-        raise Exception("We couldn't find a SDXL LoRA on the model you've sent")
-    return model_data["name"], f"{lora_archive}/{safetensors_name}", trigger_word, image_url
-
+  link_split = link.split("civitai.com/")
+  pattern = re.compile(r'models\/(\d+)')
+  regex_match = pattern.search(link_split[1])
+  if(regex_match):
+    civitai_model_id = regex_match.group(1)
+  else:
+    gr.Warning("No CivitAI model id found in your URL")
+    raise Exception("No CivitAI model id found in your URL")
+  model_request_url = f"https://civitai.com/api/v1/models/{civitai_model_id}?token={os.getenv('CIVITAI_TOKEN')}"
+  x = requests.get(model_request_url)
+  if(x.status_code != 200):
+      raise Exception("Invalid CivitAI URL")
+  model_data = x.json()
+  #if(model_data["nsfw"] == True or model_data["nsfwLevel"] > 20):
+   # gr.Warning("The model is tagged by CivitAI as adult content and cannot be used in this shared environment.")
+   # raise Exception("The model is tagged by CivitAI as adult content and cannot be used in this shared environment.")
+  if(model_data["type"] != "LORA"):
+    gr.Warning("The model isn't tagged at CivitAI as a LoRA")
+    raise Exception("The model isn't tagged at CivitAI as a LoRA")
+  model_link_download = None
+  image_url = None
+  trigger_word = ""
+  for model in model_data["modelVersions"]:
+    if(model["baseModel"] == "SDXL 1.0"):
+      model_link_download = f"{model['downloadUrl']}/?token={os.getenv('CIVITAI_TOKEN')}"
+      safetensors_name = model["files"][0]["name"]
+      if(not os.path.exists(f"{lora_archive}/{safetensors_name}")):
+        safetensors_file_request = requests.get(model_link_download)
+        if(safetensors_file_request.status_code != 200):
+            raise Exception("Invalid CivitAI download link")
+        with open(f"{lora_archive}/{safetensors_name}", 'wb') as file:
+            file.write(safetensors_file_request.content)
+      trigger_word = model.get("trainedWords", [""])[0]
+      for image in model["images"]:
+        if(image["nsfwLevel"] == 1):
+          image_url = image["url"]
+          break
+      break
+  if(not model_link_download):
+    gr.Warning("We couldn't find a SDXL LoRA on the model you've sent")
+    raise Exception("We couldn't find a SDXL LoRA on the model you've sent")
+  return model_data["name"], f"{lora_archive}/{safetensors_name}", trigger_word, image_url
+  
 def check_custom_model(link):
     if(link.startswith("https://")):
         if(link.startswith("https://huggingface.co") or link.startswith("https://www.huggingface.co")):
@@ -584,6 +465,9 @@ def check_custom_model(link):
     else: 
         return get_huggingface_safetensors(link)
 
+def show_loading_widget():
+    return gr.update(visible=True)
+
 def load_custom_lora(link):
     if(link):
         try:
@@ -609,29 +493,33 @@ def load_custom_lora(link):
 
 def remove_custom_lora():
     return "", gr.update(visible=False), gr.update(visible=False), None
-
-# Build Gradio interface
 with gr.Blocks(css="custom.css") as demo:
     gr_sdxl_loras = gr.State(value=sdxl_loras_raw)
     title = gr.HTML(
         """<h1><img src="https://i.imgur.com/DVoGw04.png">
-<span>Face to All - Enhanced<br><small style="
+<span>Face to All<br><small style="
     font-size: 13px;
     display: block;
     font-weight: normal;
     opacity: 0.75;
-">🔥 Supports: No faces (landscape), Multiple faces, Improved quality, Custom LoRAs<br> diffusers InstantID + ControlNet</small></span></h1>""",
+">🧨 diffusers InstantID + ControlNet<br> inspired by fofr's <a href="https://github.com/fofr/cog-face-to-many" target="_blank">face-to-many</a></small></span></h1>""",
         elem_id="title",
     )
     selected_state = gr.State()
     custom_loaded_lora = gr.State()
-    
     with gr.Row(elem_id="main_app"):
         with gr.Column(scale=4, elem_id="box_column"):
             with gr.Group(elem_id="gallery_box"):
-                photo = gr.Image(label="Upload a picture (with or without faces)", interactive=True, type="pil", height=300)
-                selected_loras = gr.Gallery(label="Selected LoRAs", height=80, show_share_button=False, visible=False, elem_id="gallery_selected")
+                photo = gr.Image(label="Upload a picture of yourself", interactive=True, type="pil", height=300)
+                selected_loras = gr.Gallery(label="Selected LoRAs", height=80, show_share_button=False, visible=False, elem_id="gallery_selected", )
+                #order_gallery = gr.Radio(choices=["random", "likes"], value="random", label="Order by", elem_id="order_radio")
+                #new_gallery = gr.Gallery(
+                #    label="New LoRAs",
+                #    elem_id="gallery_new",
+                #    columns=3,
+                #    value=[(item["image"], item["title"]) for item in sdxl_loras_raw_new], allow_preview=False, show_share_button=False)
                 gallery = gr.Gallery(
+                    #value=[(item["image"], item["title"]) for item in sdxl_loras],
                     label="Pick a style from the gallery",
                     allow_preview=False,
                     columns=4,
@@ -642,82 +530,77 @@ with gr.Blocks(css="custom.css") as demo:
                 custom_model = gr.Textbox(label="or enter a custom Hugging Face or CivitAI SDXL LoRA", placeholder="Paste Hugging Face or CivitAI model path...")
                 custom_model_card = gr.HTML(visible=False)
                 custom_model_button = gr.Button("Remove custom LoRA", visible=False)
-        
         with gr.Column(scale=5):
             with gr.Row():
-                prompt = gr.Textbox(label="Prompt", show_label=False, lines=1, max_lines=1, 
-                                   info="Describe your subject or scene", value="a person", elem_id="prompt")
+                prompt = gr.Textbox(label="Prompt", show_label=False, lines=1, max_lines=1, info="Describe your subject (optional)", value="a person", elem_id="prompt")
                 button = gr.Button("Run", elem_id="run_button")
-            
             result = ImageSlider(
                 interactive=False, label="Generated Image", elem_id="result-image", position=0.1
             )
-            
             with gr.Group(elem_id="share-btn-container", visible=False) as share_group:
                 community_icon = gr.HTML(community_icon_html)
                 loading_icon = gr.HTML(loading_icon_html)
                 share_button = gr.Button("Share to community", elem_id="share-btn")
-            
             with gr.Accordion("Advanced options", open=False):
-                use_multiple_faces = gr.Checkbox(label="Use multiple faces (if detected)", value=False)
                 negative = gr.Textbox(label="Negative Prompt")
                 weight = gr.Slider(0, 10, value=0.9, step=0.1, label="LoRA weight")
-                face_strength = gr.Slider(0, 2, value=0.9, step=0.01, label="Face strength", 
-                                         info="Higher values increase face likeness (auto-adjusted for no-face images)")
-                image_strength = gr.Slider(0, 1, value=0.20, step=0.01, label="Image strength", 
-                                          info="Higher values increase similarity with original structure/colors")
-                guidance_scale = gr.Slider(0, 50, value=8, step=0.1, label="Guidance Scale")
-                depth_control_scale = gr.Slider(0, 1, value=0.8, step=0.01, label="Zoe Depth ControlNet strength")
-            
+                face_strength = gr.Slider(0, 2, value=0.85, step=0.01, label="Face strength", info="Higher values increase the face likeness but reduce the creative liberty of the models")
+                image_strength = gr.Slider(0, 1, value=0.15, step=0.01, label="Image strength", info="Higher values increase the similarity with the structure/colors of the original photo")
+                guidance_scale = gr.Slider(0, 50, value=7, step=0.1, label="Guidance Scale")
+                depth_control_scale = gr.Slider(0, 1, value=0.8, step=0.01, label="Zoe Depth ControlNet strenght")
             prompt_title = gr.Markdown(
                 value="### Click on a LoRA in the gallery to select it",
                 visible=True,
                 elem_id="selected_lora",
             )
-    
-    # Event handlers
+    #order_gallery.change(
+    #    fn=swap_gallery,
+    #    inputs=[order_gallery, gr_sdxl_loras],
+    #    outputs=[gallery, gr_sdxl_loras],
+    #    queue=False
+    #)
     custom_model.input(
         fn=load_custom_lora,
         inputs=[custom_model],
         outputs=[custom_model_card, custom_model_card, custom_model_button, custom_loaded_lora, gallery, prompt_title],
     )
-    
     custom_model_button.click(
         fn=remove_custom_lora,
         outputs=[custom_model, custom_model_button, custom_model_card, custom_loaded_lora]
     )
-    
     gallery.select(
         fn=update_selection,
         inputs=[gr_sdxl_loras, face_strength, image_strength, weight, depth_control_scale, negative],
         outputs=[prompt_title, prompt, face_strength, image_strength, weight, depth_control_scale, negative, selected_state],
         show_progress=False
     )
-    
+    #new_gallery.select(
+    #    fn=update_selection,
+    #    inputs=[gr_sdxl_loras_new, gr.State(True)],
+    #    outputs=[prompt_title, prompt, prompt, selected_state, gallery],
+    #    queue=False,
+    #    show_progress=False
+    #)
     prompt.submit(
         fn=check_selected,
         inputs=[selected_state, custom_loaded_lora],
         show_progress=False
     ).success(
         fn=run_lora,
-        inputs=[photo, prompt, negative, weight, selected_state, face_strength, image_strength, 
-               guidance_scale, depth_control_scale, gr_sdxl_loras, custom_loaded_lora, use_multiple_faces],
+        inputs=[photo, prompt, negative, weight, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, gr_sdxl_loras, custom_loaded_lora],
         outputs=[result, share_group],
     )
-    
     button.click(
         fn=check_selected,
         inputs=[selected_state, custom_loaded_lora],
         show_progress=False
     ).success(
         fn=run_lora,
-        inputs=[photo, prompt, negative, weight, selected_state, face_strength, image_strength, 
-               guidance_scale, depth_control_scale, gr_sdxl_loras, custom_loaded_lora, use_multiple_faces],
+        inputs=[photo, prompt, negative, weight, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, gr_sdxl_loras, custom_loaded_lora],
         outputs=[result, share_group],
     )
-    
     share_button.click(None, [], [], js=share_js)
-    demo.load(fn=classify_gallery, inputs=[gr_sdxl_loras], outputs=[gallery, gr_sdxl_loras])
+    demo.load(fn=classify_gallery, inputs=[gr_sdxl_loras], outputs=[gallery, gr_sdxl_loras], js=js)
 
 demo.queue(default_concurrency_limit=None, api_open=True)
 demo.launch(share=True)