Update app.py

app.py

@@ -31,9 +31,18 @@ def save_to_s3(image_data, payload, file_name):
 
 
 # Function to handle the NovelAI API request
-def generate_novelai_image(input_text, quality_tags, seed, negative_prompt, scale, sampler):
+def generate_novelai_image(input_text, quality_tags, seed, negative_prompt, scale, ratio, sampler):
     jwt_token = os.environ.get('NAI_API_KEY')
-
+    if ratio == "Landscape (1216x832)":
+        width = 1216
+        height = 832
+    elif ratio == "Square (1024x1024)":
+        width = 1024
+        height = 1024
+    elif ratio == "Portrait (832x1216)":
+        width = 832
+        height = 1216
+        
     # Check if quality tags are provided and append to input
     final_input = input_text
     if quality_tags:
@@ -60,8 +69,8 @@ def generate_novelai_image(input_text, quality_tags, seed, negative_prompt, scal
         "input": final_input,
         "model": "nai-diffusion-3",
         "parameters": {
-            "width": 832,
-            "height": 1216,
+            "width": width,
+            "height": height,
             "scale": scale,
             "sampler": sampler,
             "steps": 28,
@@ -120,6 +129,7 @@ iface = gr.Interface(
         gr.Slider(minimum=-1, maximum=2**32 - 1, step=1, value=-1, label="Seed"),
         gr.Textbox(label="Negative Prompt", value="nsfw, lowres, {bad}, error, fewer, extra, missing, worst quality, jpeg artifacts, bad quality, watermark, unfinished, displeasing, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]"),
         gr.Slider(minimum=1, maximum=20, step=1, value=5, label="Scale"),
+        gr.Radio(choices=["Landscape (1216x832)", "Square (1024x1024)", "Portrait (832x1216)"]),
         gr.Dropdown(
             choices=[
                 "k_euler", "k_euler_ancestral", "k_dpmpp_2s_ancestral",