resolution options

app.py

@@ -72,9 +72,9 @@ class Ours:
             "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16
         ).to(device)
 
-    def inference(self, text, negative_text, steps, guidance_scale):
+    def inference(self, text, negative_text, steps, guidance_scale, width, height):
         gen_images = []
-        for i in range(4):
+        for i in range(2):
             image_emb, negative_image_emb = self.pipe_prior(
                 text, negative_prompt=negative_text
             ).to_tuple()
@@ -83,6 +83,8 @@ class Ours:
                 negative_image_embeds=negative_image_emb,
                 num_inference_steps=steps,
                 guidance_scale=guidance_scale,
+                width=width,
+                height=height,
             ).images
             gen_images.append(image[0])
         return gen_images
@@ -91,8 +93,11 @@ class Ours:
 selected_model = Ours(device=__DEVICE__)
 
 @spaces.GPU
-def get_images(text, negative_text, steps, guidance_scale):
-    images = selected_model.inference(text, negative_text, steps, guidance_scale)
+def get_images(text, negative_text, steps, guidance_scale, width, height, fixed_res):
+    if fixed_res!="manual":
+        print(f"Using {fixed_res} resolution")
+        width, height = fixed_res.split("x")
+    images = selected_model.inference(text, negative_text, steps, guidance_scale, width=int(width), height=int(height))
     new_images = []
     for img in images:
         new_images.append(img)
@@ -134,12 +139,37 @@ with gr.Blocks() as demo:
                 label="Guidance Scale", minimum=0, maximum=10, value=7.5, step=0.1
             )
 
+        with gr.Row():
+            with gr.Group():
+                width_inp = gr.Textbox(
+                    label="Please provide the width",
+                    value="512",
+                    max_lines=1,
+                ).style(
+                    border=(True, False, True, True),
+                    rounded=(True, False, False, True),
+                    container=False,
+                )
+                height_inp = gr.Textbox(
+                    label="Please provide the height",
+                    max_lines=1,
+                    value="512",
+                ).style(
+                    border=(True, False, True, True),
+                    rounded=(True, False, False, True),
+                    container=False,
+                )
+
+            fixed_res = gr.Dropdown(
+                ["manual", "512x512", "1024x1024", "1920x1080", "1280x720"], value="manual", label="Prefined Resolution", info="Either select one or manually define one!"
+            )
+
         with gr.Row():
             btn = gr.Button(value="Generate Image")
 
     gallery = gr.Gallery(
         label="Generated images", show_label=False, elem_id="gallery"
-    , columns=[4], rows=[1], object_fit="contain", height="auto")
+    , columns=[2], rows=[1], object_fit="contain", height="auto")
 
     btn.click(
         get_images,
@@ -148,6 +178,9 @@ with gr.Blocks() as demo:
             negative_text,
             steps,
             guidance_scale,
+            width_inp,
+            height_inp,
+            fixed_res[0],
         ],
         outputs=gallery,
     )
@@ -158,6 +191,9 @@ with gr.Blocks() as demo:
             negative_text,
             steps,
             guidance_scale,
+            width_inp,
+            height_inp,
+            fixed_res[0],
         ],
         outputs=gallery,
     )
@@ -168,6 +204,9 @@ with gr.Blocks() as demo:
             negative_text,
             steps,
             guidance_scale,
+            width_inp,
+            height_inp,
+            fixed_res[0],
         ],
         outputs=gallery,
     )