Update app.py

app.py

@@ -24,6 +24,9 @@ import cv2
 import torch
 import numpy as np
 from PIL import Image
+from io import BytesIO
+import base64
+import re
 
 from insightface.app import FaceAnalysis
 from pipeline_stable_diffusion_xl_instantid_img2img import StableDiffusionXLInstantIDImg2ImgPipeline, draw_kps
@@ -34,6 +37,26 @@ from compel import Compel, ReturnedEmbeddingsType
 
 #from gradio_imageslider import ImageSlider
 
+
+# Regex pattern to match data URI scheme
+data_uri_pattern = re.compile(r'data:image/(png|jpeg|jpg|webp);base64,')
+
+def readb64(b64):
+    # Remove any data URI scheme prefix with regex
+    b64 = data_uri_pattern.sub("", b64)
+    # Decode and open the image with PIL
+    img = Image.open(BytesIO(base64.b64decode(b64)))
+    return img
+    
+# convert from PIL to base64
+def writeb64(image):
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    b64image = base64.b64encode(buffered.getvalue())
+    b64image_str = b64image.decode("utf-8")
+    return b64image_str
+
+
 with open("sdxl_loras.json", "r") as file:
     data = json.load(file)
     sdxl_loras_raw = [
@@ -56,7 +79,27 @@ with open("sdxl_loras.json", "r") as file:
 
 with open("defaults_data.json", "r") as file:
     lora_defaults = json.load(file)
-    
+
+
+def getLoraByRepoName(repo_name):
+    # Loop through each lora in sdxl_loras_raw
+    for lora in sdxl_loras_raw:
+        if lora["repo"] == repo_name:
+            # Return the lora if the repo name matches
+            return lora
+    # If no match is found, return the first lora in the array
+    return sdxl_loras_raw[0] if sdxl_loras_raw else None
+
+# Return the default values specific to this particular 
+def getLoraDefaultsByRepoName(repo_name):
+    # Loop through each lora in sdxl_loras_raw
+    for lora in lora_defaults:
+        if lora["repo"] == repo_name:
+            # Return the lora if the repo name matches
+            return lora
+    # If no match is found, return the first lora in the array
+    return lora_defaults[0] if lora_defaults else None
+
 
 device = "cuda" 
 
@@ -148,36 +191,6 @@ pipe.to(device)
 last_lora = ""
 last_fused = False
 
-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 '' }"
-
-    for lora_list in lora_defaults:
-        if lora_list["model"] == sdxl_loras[selected_state.index]["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
-    )
 
 def center_crop_image_as_square(img):
     square_size = min(img.size)
@@ -216,13 +229,13 @@ def merge_incompatible_lora(full_path_lora, lora_scale):
                 del weights_sd
                 del lora_model
 #@spaces.GPU
-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):
+def generate_image(prompt, negative, face_emb, face_image, face_kps, image_strength, guidance_scale, face_strength, depth_control_scale, lora, full_path_lora, lora_scale, st):
     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("Current LoRA: ", lora["repo"])
     print("Last fused: ", last_fused)
     #prepare face zoe
     st = time.time()
@@ -233,7 +246,7 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
     et = time.time()
     elapsed_time = et - st
     print('Zoe Depth calculations took: ', elapsed_time, 'seconds')
-    if last_lora != repo_name:
+    if last_lora != lora["repo"]:
         if(last_fused):
             st = time.time()
             pipe.unfuse_lora()
@@ -242,17 +255,16 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
             elapsed_time = et - st
             print('Unfuse and unload LoRA took: ', elapsed_time, 'seconds')
         st = time.time()
-        pipe.load_lora_weights(loaded_state_dict)
+        pipe.load_lora_weights(full_path_lora)
         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):
+        if(lora["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")
+            text_embedding_name = lora["text_embedding_weights"]
+            embedding_path = hf_hub_download(repo_id=lora["repo"], filename=text_embedding_name, repo_type="model")
             state_dict_embedding = load_file(embedding_path)
             try:
                 pipe.unload_textual_inversion()
@@ -293,12 +305,19 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
     et = time.time()
     elapsed_time = et - st
     print('Image processing took: ', elapsed_time, 'seconds')
-    last_lora = repo_name
+    last_lora = lora["repo"]
+        
     return image
     
-def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, sdxl_loras, progress=gr.Progress(track_tqdm=True)):
-    selected_state_index = selected_state.index
+def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, lora_repo_name):
+    # get the lora and its defaulrt values
+    lora = getLoraByRepoName(lora_repo_name)
+    default_values = getLoraDefaultsByRepoName(lora_repo_name)
+    
     st = time.time()
+
+    face_image = readb64(face_image)
+    
     face_image = center_crop_image_as_square(face_image)
     try:
         face_info = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
@@ -312,160 +331,92 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
     print('Cropping and calculating face embeds took: ', elapsed_time, 'seconds')
 
     st = time.time()
-    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)
+  
+    if default_values:
+        prompt_full = default_values.get("prompt", None)
+        if(prompt_full):
+            prompt = prompt_full.replace("<subject>", prompt)
 
     
     print("Prompt:", prompt) 
     if(prompt == ""):
         prompt = "a person"
-    
-    print("Selected State: ", selected_state_index)
-    print(sdxl_loras[selected_state_index]["repo"])
+
     if negative == "":
         negative = None
 
     if not selected_state:
         raise gr.Error("You must select a LoRA")
-    repo_name = sdxl_loras[selected_state_index]["repo"]
-    weight_name = sdxl_loras[selected_state_index]["weights"]
+
+    weight_name = lora["weights"]
     
-    full_path_lora = state_dicts[repo_name]["saved_name"]
-    #loaded_state_dict = copy.deepcopy(state_dicts[repo_name]["state_dict"])
+    full_path_lora = state_dicts[lora["repo"]]["saved_name"]
+    #loaded_state_dict = copy.deepcopy(state_dicts[lora_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 image
 
-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)
+    image = generate_image(prompt, negative, face_emb, face_image, face_kps, image_strength, guidance_scale, face_strength, depth_control_scale, lora, full_path_lora, lora_scale, st)
         
-def deselect():
-  return gr.Gallery(selected_index=None)
+    image_base64 = writeb64(image)
+        
+    return image_base64
+
 
 with gr.Blocks() 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()
-    with gr.Row(elem_id="main_app"):
-        with gr.Column(scale=4):
-            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="Style gallery",
-                    allow_preview=False,
-                    columns=4,
-                    elem_id="gallery",
-                    show_share_button=False,
-                    height=550
-                )
-                custom_model = gr.Textbox(label="Enter a custom Hugging Face or CivitAI SDXL LoRA", interactive=False, placeholder="Coming soon...")
-        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 = gr.Image(
-                interactive=False, label="Generated Image", elem_id="result-image"
-            )
-
-            with gr.Accordion("Advanced options", open=False):
-                negative = gr.Textbox(label="Negative Prompt")
-
-                # initial value was 0.9
-                weight = gr.Slider(0, 10, value=6, step=0.1, label="LoRA weight")
-
-                # initial value was 0.85
-                face_strength = gr.Slider(0, 1, value=0.75, step=0.01, label="Face strength", info="Higher values increase the face likeness but reduce the creative liberty of the models")
-
-                # initial value was 0.15
-                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")
-
-                # initial value was 7
-                guidance_scale = gr.Slider(0, 50, value=7, step=0.1, label="Guidance Scale")
-
-                # initial value was 1
-                depth_control_scale = gr.Slider(0, 4, 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
-    #)
-    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],
-        queue=False,
-        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],
-        queue=False,
-        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],
-        outputs=[result],
-    )
-    button.click(
-        fn=check_selected,
-        inputs=[selected_state],
-        queue=False,
-        show_progress=False
-    ).success(
+    gr.HTML("""
+        <div style="z-index: 100; position: fixed; top: 0px; right: 0px; left: 0px; bottom: 0px; width: 100%; height: 100%; background: white; display: flex; align-items: center; justify-content: center; color: black;">
+        <div style="text-align: center; color: black;">
+        <p style="color: black;">This space is a REST API to programmatically generate an image from a face.</p>
+        <p style="color: black;">Interested in using it through an UI? Please use the <a href="https://huggingface.co/spaces/multimodalart/face-to-all" target="_blank">original space</a>, thank you!</p>
+        </div>
+        </div>""")
+    
+    input_image_base64 = gr.Text()
+
+    lora_repo_name = gr.Text(label="name of the LoRA repo nape on HF")
+
+    prompt = gr.Textbox(label="Prompt", show_label=False, lines=1, max_lines=1, info="Describe your subject (optional)", value="a person", elem_id="prompt")
+
+    negative = gr.Textbox(label="Negative Prompt")
+    
+    # initial value was 0.9
+    weight = gr.Slider(0, 10, value=6, step=0.1, label="LoRA weight")
+    
+    # initial value was 0.85
+    face_strength = gr.Slider(0, 1, value=0.75, step=0.01, label="Face strength", info="Higher values increase the face likeness but reduce the creative liberty of the models")
+    
+    # initial value was 0.15
+    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")
+    
+    # initial value was 7
+    guidance_scale = gr.Slider(0, 50, value=7, step=0.1, label="Guidance Scale")
+    
+    # initial value was 1
+    depth_control_scale = gr.Slider(0, 4, value=0.8, step=0.01, label="Zoe Depth ControlNet strenght")
+
+    button = gr.Button(value="Generate")
+    output_image_base64 = gr.Text()
+    text2image_predict.click(
         fn=run_lora,
-        inputs=[photo, prompt, negative, weight, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, gr_sdxl_loras],
-        outputs=[result],
+        inputs=[
+            input_image_base64,
+            prompt,
+            negative,
+            weight,
+            face_strength,
+            image_strength,
+            guidance_scale,
+            depth_control_scale,
+            lora_repo_name
+        ],
+        outputs=output_image_base64,
+        api_name='run',
     )
 
-    demo.load(fn=classify_gallery, inputs=[gr_sdxl_loras], outputs=[gallery, gr_sdxl_loras], queue=False)
+
 demo.queue(max_size=20)
 demo.launch()