Update app.py

app.py

@@ -69,7 +69,7 @@ def custom_model_changed(custom_model, previous_model):
     return status_message
 
 @spaces.GPU
-def infer (custom_model, weight_name, prompt, inf_steps, guidance_scale, seed, lora_weight, progress=gr.Progress(track_tqdm=True)):
+def infer (custom_model, weight_name, prompt, inf_steps, guidance_scale, width, height, seed, lora_weight, progress=gr.Progress(track_tqdm=True)):
     
     vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
     
@@ -106,6 +106,8 @@ def infer (custom_model, weight_name, prompt, inf_steps, guidance_scale, seed, l
     image = pipe(
         prompt=prompt, 
         num_inference_steps=inf_steps,
+        width=width,
+        height=height,
         guidance_scale = guidance_scale,
         generator=generator,
         cross_attention_kwargs={"scale": lora_weight}
@@ -235,6 +237,22 @@ with gr.Blocks(css=css) as demo:
                     step=0.1,
                     value=7.5
                 )
+
+            with gr.Row():
+                width = gr.Slider(
+                    label="Width",
+                    minimum=256,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=1024,
+                )
+                height = gr.Slider(
+                    label="Height",
+                    minimum=256,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=1024,
+                )
     
             with gr.Row():
                 seed = gr.Slider(
@@ -273,7 +291,7 @@ with gr.Blocks(css=css) as demo:
     )
     submit_btn.click(
         fn = infer,
-        inputs = [custom_model, weight_name, prompt_in, inf_steps, guidance_scale, seed, lora_weight],
+        inputs = [custom_model, weight_name, prompt_in, inf_steps, guidance_scale, width, height, seed, lora_weight],
         outputs = [image_out, last_used_seed]
     )