Update app.py

app.py

@@ -232,7 +232,7 @@ 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, width, height, repo_name, loaded_state_dict, lora_scale, sdxl_loras, selected_state_index, st):
     print(loaded_state_dict)
     et = time.time()
     elapsed_time = et - st
@@ -292,8 +292,8 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
         pooled_prompt_embeds=pooled,
         negative_prompt_embeds=negative_conditioning,
         negative_pooled_prompt_embeds=negative_pooled,
-        width=1024,
-        height=1024,
+        width=width,
+        height=height,
         image_embeds=face_emb,
         image=face_image,
         strength=1-image_strength,
@@ -308,7 +308,7 @@ def generate_image(prompt, negative, face_emb, face_image, face_kps, image_stren
     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)):
+def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, aspect_ratio, 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
@@ -362,7 +362,18 @@ def run_lora(face_image, prompt, negative, lora_scale, selected_state, face_stre
     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)
+
+    if(aspect_ratio == "portrait"):
+        width = 1024
+        height = 1536
+    elif(aspect_ratio == "landscape"):
+        width = 1536
+        height = 1024
+    else:
+        width = 1024
+        height = 1024
+        
+    image = generate_image(prompt, negative, face_emb, face_image, face_kps, image_strength, guidance_scale, face_strength, depth_control_scale, width, height, repo_name, full_path_lora, lora_scale, sdxl_loras, selected_state_index, st)
     return (face_image, image), gr.update(visible=True)
 
 def shuffle_gallery(sdxl_loras):
@@ -542,6 +553,7 @@ with gr.Blocks(css="custom.css") as demo:
                 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")
+                aspect_ratio = gr.CheckboxGroup(choices=["square", "portrait", "landscape"], value="square", label="Aspect Ratio")
             prompt_title = gr.Markdown(
                 value="### Click on a LoRA in the gallery to select it",
                 visible=True,
@@ -575,24 +587,17 @@ with gr.Blocks(css="custom.css") as demo:
     #    queue=False,
     #    show_progress=False
     #)
-    prompt.submit(
+    gr.on(
+        triggers=[prompt.submit, 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],
-    )
-    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],
+        inputs=[photo, prompt, negative, weight, selected_state, face_strength, image_strength, guidance_scale, depth_control_scale, aspect_ratio, 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)