allow custom size

app.py

@@ -238,7 +238,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, repo_name, loaded_state_dict, lora_scale, sdxl_loras, selected_state_index, size, st):
     print(loaded_state_dict)
     et = time.time()
     elapsed_time = et - st
@@ -303,8 +303,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=int(size),
+        height=int(size),
         image_embeds=face_emb,
         image=face_image,
         strength=1-image_strength,
@@ -319,7 +319,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, sdxl_loras, custom_lora, size, 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
@@ -373,7 +373,7 @@ 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)
+    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, size, st)
     return image, gr.update(visible=True)
 
 def shuffle_gallery(sdxl_loras):
@@ -539,6 +539,7 @@ with gr.Blocks(css="custom.css") as demo:
             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")
+            size = gr.Slider(minimum=512, maximum=1024, step=128, value=1024, interactive=True)
             result = gr.Image(
                 interactive=False, label="Generated Image", elem_id="result-image"
             )
@@ -595,7 +596,7 @@ with gr.Blocks(css="custom.css") as demo:
         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, gr_sdxl_loras, custom_loaded_lora, size],
         outputs=[result, share_group],
     )
     button.click(
@@ -605,7 +606,7 @@ with gr.Blocks(css="custom.css") as demo:
         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, gr_sdxl_loras, custom_loaded_lora, size],
         outputs=[result, share_group],
     )
     share_button.click(None, [], [], js=share_js)