Update app.py

app.py

@@ -291,30 +291,21 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
     print('Prompt processing took: ', elapsed_time, 'seconds')
     print("Processing image...")
     st = time.time()
-    
-    inputs = {
-        "prompt_embeds": conditioning,
-        "pooled_prompt_embeds": pooled,
-        "negative_prompt_embeds": negative_conditioning,
-        "negative_pooled_prompt_embeds": negative_pooled,
-        "width": 1024,
-        "height": 1024,
-        "image": face_image,
-        "strength": 1 - image_strength,
-        "guidance_scale": guidance_scale,
-        "num_inference_steps": 37,
-    }
-    
-    if face_emb is not None:
-        inputs["image_embeds"] = face_emb
-        inputs["control_image"] = [face_kps, image_zoe.resize((height, width))]
-        inputs["controlnet_conditioning_scale"] = [face_strength, depth_control_scale]
-    else:
-        inputs["control_image"] = [image_zoe.resize((height, width))]
-        inputs["controlnet_conditioning_scale"] = [depth_control_scale]
-    
-    image = pipe(**inputs).images[0]
-
+    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')
@@ -326,20 +317,14 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
     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)
+    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] if face_info else None
-    except:
-        face_info = None
-    
-    if face_info:
+        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'])
-    else:
-        face_emb = None
-        face_kps = Image.new("RGB", face_image.size, (0, 0, 0))  # empty placeholder for consistency
-
+    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')