Update app.py

app.py

@@ -1,606 +1,166 @@
-import gradio as gr
-import torch
-import spaces
-torch.jit.script = lambda f: f
-import timm
+import os
+import re
 import time
-
-from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_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 copy
 import random
-from urllib.parse import quote
-import gdown
-import os
-import re
 import requests
+import torch
+import cv2
+import numpy as np
+import gradio as gr
+import spaces
+from PIL import Image
+from urllib.parse import quote
+
+# Disable Torch JIT compilation for compatibility
+torch.jit.script = lambda f: f  
 
+# Model & Utilities
+import timm
 import diffusers
 from diffusers.utils import load_image
 from diffusers.models import ControlNetModel
 from diffusers import AutoencoderKL, DPMSolverMultistepScheduler, UNet2DConditionModel
-import cv2
-import torch
-import numpy as np
-from PIL import Image
-
+from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
 from insightface.app import FaceAnalysis
-from pipeline_stable_diffusion_xl_instantid_img2img import StableDiffusionXLInstantIDImg2ImgPipeline, draw_kps
 from controlnet_aux import ZoeDetector
-
 from compel import Compel, ReturnedEmbeddingsType
-
 from gradio_imageslider import ImageSlider
 
+# Custom imports
+try:
+    from pipeline_stable_diffusion_xl_instantid_img2img import StableDiffusionXLInstantIDImg2ImgPipeline, draw_kps
+    from cog_sdxl_dataset_and_utils import TokenEmbeddingsHandler
+except ImportError as e:
+    print(f"Import Error: {e}. Check if modules exist or paths are correct.")
+    exit()
 
-#from gradio_imageslider import ImageSlider
+# Device setup
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# Load LoRA configuration
 with open("sdxl_loras.json", "r") as file:
-    data = json.load(file)
-    sdxl_loras_raw = [
-        {
-            "image": item["image"],
-            "title": item["title"],
-            "repo": item["repo"],
-            "trigger_word": item["trigger_word"],
-            "weights": item["weights"],
-            "is_compatible": item["is_compatible"],
-            "is_pivotal": item.get("is_pivotal", False),
-            "text_embedding_weights": item.get("text_embedding_weights", None),
-            "likes": item.get("likes", 0),
-            "downloads": item.get("downloads", 0),
-            "is_nc": item.get("is_nc", False),
-            "new": item.get("new", False),
-        }
-        for item in data
-    ]
+    sdxl_loras_raw = json.load(file)
 
 with open("defaults_data.json", "r") as file:
     lora_defaults = json.load(file)
-    
-
-device = "cuda" 
-
-state_dicts = {}
-
-for item in sdxl_loras_raw:
-    saved_name = hf_hub_download(item["repo"], item["weights"])
-    
-    if not saved_name.endswith('.safetensors'):
-        state_dict = torch.load(saved_name)
-    else:
-        state_dict = load_file(saved_name)
-    
-    state_dicts[item["repo"]] = {
-        "saved_name": saved_name,
-        "state_dict": state_dict
-    }
 
-sdxl_loras_raw = [item for item in sdxl_loras_raw if item.get("new") != True]
-    
-# 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
-#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 required models
+CHECKPOINT_DIR = "/data/checkpoints"
+hf_hub_download(repo_id="InstantX/InstantID", filename="ControlNetModel/config.json", local_dir=CHECKPOINT_DIR)
+hf_hub_download(repo_id="InstantX/InstantID", filename="ControlNetModel/diffusion_pytorch_model.safetensors", local_dir=CHECKPOINT_DIR)
+hf_hub_download(repo_id="InstantX/InstantID", filename="ip-adapter.bin", local_dir=CHECKPOINT_DIR)
+hf_hub_download(repo_id="latent-consistency/lcm-lora-sdxl", filename="pytorch_lora_weights.safetensors", local_dir=CHECKPOINT_DIR)
 
+# Download Antelopev2 Face Recognition model
 antelope_download = snapshot_download(repo_id="DIAMONIK7777/antelopev2", local_dir="/data/models/antelopev2")
-print(antelope_download)
-app = FaceAnalysis(name='antelopev2', root='/data', providers=['CPUExecutionProvider'])
+print("Antelopev2 Download Path:", antelope_download)
+
+# Initialize FaceAnalysis
+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 identity & depth models
+face_adapter = os.path.join(CHECKPOINT_DIR, "ip-adapter.bin")
+controlnet_path = os.path.join(CHECKPOINT_DIR, "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)
-et = time.time()
-elapsed_time = et - st
-print('Loading ControlNet took: ', elapsed_time, 'seconds')
-st = time.time()
+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)
-et = time.time()
-elapsed_time = et - st
-print('Loading VAE took: ', elapsed_time, 'seconds')
-st = time.time()
 
-#pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained("stablediffusionapi/albedobase-xl-v21",
-pipe = StableDiffusionXLInstantIDImg2ImgPipeline.from_pretrained("frankjoshua/albedobaseXL_v21",
-                                                                 vae=vae,
-                                                                 controlnet=[identitynet, zoedepthnet],
-                                                                 torch_dtype=torch.float16)
+# Load main pipeline
+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.8)
-et = time.time()
-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])
-et = time.time()
-elapsed_time = et - st
-print('Loading Compel took: ', elapsed_time, 'seconds')
 
-st = time.time()
+# Initialize Compel for text conditioning
+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]
+)
+
+# Load ZoeDetector for depth estimation
 zoe = ZoeDetector.from_pretrained("lllyasviel/Annotators")
-et = time.time()
-elapsed_time = et - st
-print('Loading Zoe took: ', elapsed_time, 'seconds')
 zoe.to(device)
 pipe.to(device)
 
+# LoRA Management
 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"
 
-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"
-    weight_name = sdxl_loras[selected_state.index]["weights"]
-    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨ {'(non-commercial LoRA, `cc-by-nc`)' if sdxl_loras[selected_state.index]['is_nc'] else '' }"
+# --- Utility Functions ---
+def update_selection(selected_state, sdxl_loras, face_strength, image_strength, weight, depth_control_scale, negative):
+    index = selected_state.index
+    lora_repo = sdxl_loras[index]["repo"]
+    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo})"
 
     for lora_list in lora_defaults:
-        if lora_list["model"] == sdxl_loras[selected_state.index]["repo"]:
+        if lora_list["model"] == lora_repo:
             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", "")
-    
-    if(is_new):
-        if(selected_state.index == 0):
-            selected_state.index = -9999
-        else:
-            selected_state.index *= -1
-    
+
     return (
-        updated_text,
-        gr.update(placeholder=new_placeholder),
-        face_strength,
-        image_strength,
-        weight,
-        depth_control_scale,
-        negative,
-        selected_state
+        updated_text, gr.update(placeholder="Type a prompt"), face_strength, 
+        image_strength, weight, depth_control_scale, negative, selected_state
     )
 
-def center_crop_image_as_square(img):
+def center_crop_image(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 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=80)
-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')
+    left = (img.width - square_size) // 2
+    top = (img.height - square_size) // 2
+    return img.crop((left, top, left + square_size, top + square_size))
+
+def process_face(image):
+    face_info = app.get(cv2.cvtColor(np.array(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]
+    face_emb = face_info['embedding']
+    face_kps = draw_kps(image, face_info['kps'])
+    return face_emb, face_kps
+
+def generate_image(prompt, negative, face_emb, face_image, face_kps, image_strength, guidance_scale, face_strength, depth_control_scale, repo_name, lora_scale):
     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()
-    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()
+    if last_lora != repo_name and last_fused:
+        pipe.unfuse_lora()
+        pipe.unload_lora_weights()
+    pipe.load_lora_weights(repo_name)
+    pipe.fuse_lora(lora_scale)
+    last_lora, last_fused = repo_name, True
+
     conditioning, pooled = compel(prompt)
-    if(negative):
-        negative_conditioning, negative_pooled = compel(negative)
-    else:
-        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=20,
-        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
+    negative_conditioning, negative_pooled = compel(negative) if negative else (None, None)
     
-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()
-    elapsed_time = et - st
-    print('Cropping and calculating face embeds took: ', elapsed_time, 'seconds')
-
-    st = time.time()
-
-    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):
-                    prompt = prompt_full.replace("<subject>", prompt)
-
-    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 == "":
-        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:
-        repo_name = custom_lora_path
-        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"]
-    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()
-    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, st)
-    return (face_image, image), gr.update(visible=True)
-
-run_lora.zerogpu = True
-
-def shuffle_gallery(sdxl_loras):
-    random.shuffle(sdxl_loras)
-    return [(item["image"], item["title"]) for item in sdxl_loras], sdxl_loras
-
-def classify_gallery(sdxl_loras):
-    sorted_gallery = sorted(sdxl_loras, key=lambda x: x.get("likes", 0), reverse=True)
-    return [(item["image"], item["title"]) for item in sorted_gallery], sorted_gallery
-
-def swap_gallery(order, sdxl_loras):
-    if(order == "random"):
-        return shuffle_gallery(sdxl_loras)
-    else:
-        return classify_gallery(sdxl_loras)
-        
-def deselect():
-  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
-
-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["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")):
-            link_split = link.split("huggingface.co/")
-            return get_huggingface_safetensors(link_split[1])
-        elif(link.startswith("https://civitai.com") or link.startswith("https://www.civitai.com")):
-            return get_civitai_safetensors(link)
-    else: 
-        return get_huggingface_safetensors(link)
-
-def show_loading_widget():
-    return gr.update(visible=True)
+    images = [face_kps, zoe(face_image).resize(face_kps.size)]
+    return 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=20,
+        guidance_scale=guidance_scale, controlnet_conditioning_scale=[face_strength, depth_control_scale]
+    ).images[0]
 
-def load_custom_lora(link):
-    if(link):
-        try:
-            title, path, trigger_word, image = check_custom_model(link)
-            card = f'''
-            <div class="custom_lora_card">
-              <span>Loaded custom LoRA:</span>
-              <div class="card_internal">
-                <img src="{image}" />
-                <div>
-                    <h3>{title}</h3>
-                    <small>{"Using: <code><b>"+trigger_word+"</code></b> as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}<br></small>
-                </div>
-              </div>
-            </div>
-            '''
-            return gr.update(visible=True), card, gr.update(visible=True), [path, trigger_word], gr.Gallery(selected_index=None), f"Custom: {path}"
-        except Exception as e:
-            gr.Warning("Invalid LoRA: either you entered an invalid link, a non-SDXL LoRA or a LoRA with mature content")
-            return gr.update(visible=True), "Invalid LoRA: either you entered an invalid link, a non-SDXL LoRA or a LoRA with mature content", gr.update(visible=False), None, gr.update(visible=True), gr.update(visible=True)
-    else:
-        return gr.update(visible=False), "", gr.update(visible=False), None, gr.update(visible=True), gr.update(visible=True)
+# --- UI Setup ---
+with gr.Blocks() as demo:
+    photo = gr.Image(label="Upload a picture", interactive=True, type="pil", height=300)
+    gallery = gr.Gallery(label="Pick a style", allow_preview=False, columns=4, height=550)
+    prompt = gr.Textbox(label="Prompt", placeholder="Enter prompt...")
+    button = gr.Button("Run")
+    result = ImageSlider(interactive=False, label="Generated Image")
 
-def remove_custom_lora():
-    return "", gr.update(visible=False), gr.update(visible=False), None
-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<br><small style="
-    font-size: 13px;
-    display: block;
-    font-weight: normal;
-    opacity: 0.75;
-">🧨 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 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,
-                    elem_id="gallery",
-                    show_share_button=False,
-                    height=550
-                )
-                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 (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):
-                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.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",
-            )
-    #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],
-        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],
-        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], js=js)
+    button.click(fn=generate_image, inputs=[prompt, gr.State(), gr.State()], outputs=result)
 
-demo.queue(default_concurrency_limit=None, api_open=True)
-demo.launch(share=True)
+demo.queue()
+demo.launch(share=True)