Update app.py

app.py

@@ -29,7 +29,7 @@ MODEL_ID = os.getenv("MODEL_VAL_PATH", "SG161222/RealVisXL_V4.0_Lightning")
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
-BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))  # Allow generating multiple images at once
+BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1")) 
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 pipe = StableDiffusionXLPipeline.from_pretrained(
@@ -59,24 +59,37 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
+def set_wallpaper_size(size):
+    if size == "Mobile (1080x1920)":
+        return 1080, 1920
+    elif size == "Desktop (1920x1080)":
+        return 1920, 1080
+    elif size == "Extended (1920x512)":
+        return 1920, 512
+    elif size == "Headers (1080x512)":
+        return 1080, 512
+    else:
+        return 1024, 1024  # Default return if none of the conditions are met
+
 @spaces.GPU(duration=60, enable_queue=True)
 def generate(
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
     seed: int = 1,
-    width: int = 1024,
-    height: int = 1024,
     guidance_scale: float = 3,
     num_inference_steps: int = 25,
     randomize_seed: bool = False,
     use_resolution_binning: bool = True, 
     num_images: int = 1,  # Number of images to generate
+    wallpaper_size: str = "Default (1024x1024)",
     progress=gr.Progress(track_tqdm=True),
 ):
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device=device).manual_seed(seed)
 
+    width, height = set_wallpaper_size(wallpaper_size)
+
     options = {
         "prompt": [prompt] * num_images,
         "negative_prompt": [negative_prompt] * num_images if use_negative_prompt else None,
@@ -92,7 +105,7 @@ def generate(
         options["use_resolution_binning"] = True
 
     images = []
-    for i in range(0, num_images, BATCH_SIZE):
+    for i in range 0, num_images, BATCH_SIZE):
         batch_options = options.copy()
         batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE]
         if "negative_prompt" in batch_options:
@@ -103,7 +116,6 @@ def generate(
     return image_paths, seed
 
 with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
-
     gr.Markdown(DESCRIPTIONx)  
     with gr.Row():
         prompt = gr.Text(
@@ -115,6 +127,13 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         )
         run_button = gr.Button("Run ⚡", scale=0)
     result = gr.Gallery(label="Result", columns=1, show_label=False) 
+
+    with gr.Row(visible=True):
+        wallpaper_size = gr.Radio(
+            choices=["Mobile (1080x1920)", "Desktop (1920x1080)", "Extended (1920x512)", "Headers (1080x512)", "Default (1024x1024)"],
+            label="Pixel Size(x*y)",
+            value="Default (1024x1024)"
+        )
         
     with gr.Accordion("Advanced options", open=False, visible=True):
         num_images = gr.Slider(
@@ -142,21 +161,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             value=0,
         )
         randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-        with gr.Row(visible=True):
-            width = gr.Slider(
-                label="Width",
-                minimum=512,
-                maximum=MAX_IMAGE_SIZE,
-                step=64,
-                value=1024,
-            )
-            height = gr.Slider(
-                label="Height",
-                minimum=512,
-                maximum=MAX_IMAGE_SIZE,
-                step=64,
-                value=1024,
-            )
+
         with gr.Row():
             guidance_scale = gr.Slider(
                 label="Guidance Scale",
@@ -197,12 +202,11 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             negative_prompt,
             use_negative_prompt,
             seed,
-            width,
-            height,
             guidance_scale,
             num_inference_steps,
             randomize_seed,
-            num_images
+            num_images,
+            wallpaper_size,
         ],
         outputs=[result, seed],
         api_name="run",