Update app.py

app.py

@@ -173,6 +173,23 @@ button.addEventListener('click', function() {
 '''
 lora_archive = "/data"
 
+def resize_image_aspect_ratio(img, max_dim=512):
+    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 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"
@@ -204,17 +221,6 @@ 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")
@@ -241,120 +247,84 @@ def merge_incompatible_lora(full_path_lora, lora_scale):
                 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')
+@spaces.GPU(duration=85)
+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):
     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')
+
+    control_images = [face_kps, zoe(face_image)] if face_detected else [zoe(face_image)]
+    control_scales = [face_strength, depth_control_scale] if face_detected else [depth_control_scale]
+
     if last_lora != repo_name:
-        if(last_fused):
-            st = time.time()
+        if last_fused:
             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
+        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)
-            
-    print("Processing prompt...")
-    st = time.time()
+
     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()
+    negative_conditioning, negative_pooled = compel(negative) if negative else (None, None)
+
     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,
+        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=images,
+        control_image=control_images,
         num_inference_steps=36,
-        guidance_scale = guidance_scale,
-        controlnet_conditioning_scale=[face_strength, depth_control_scale],
+        guidance_scale=guidance_scale,
+        controlnet_conditioning_scale=control_scales,
     ).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)
+
+    original_width, original_height = face_image.size
+
+    face_detected = True
     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_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(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()
+    except:
+        face_detected = False
+        face_emb = None
+        face_kps = face_image
 
-    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 == ""):
+    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:
@@ -362,18 +332,17 @@ 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"]
-    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)
+
+    resized_image = resize_with_aspect_ratio(face_image)
+
+    image = generate_image(
+        prompt, negative, face_emb, resized_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
+    )
+
+    return (resized_image, image), gr.update(visible=True)
+
 
 run_lora.zerogpu = True