Update app.py

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
 import numpy as np
 import random
-import spaces
 from diffusers import DiffusionPipeline
 import torch
+from PIL import Image
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "stabilityai/stable-diffusion-3.5-large-turbo"
@@ -46,7 +46,16 @@ style_list = [
 STYLE_NAMES = [style["name"] for style in style_list]
 DEFAULT_STYLE_NAME = STYLE_NAMES[0]
 
-@spaces.GPU(duration=60)
+grid_sizes = {
+    "2x1": (2, 1),
+    "1x2": (1, 2),
+    "2x2": (2, 2),
+    "2x3": (2, 3),
+    "3x2": (3, 2),
+    "1x1": (1, 1)
+}
+
+@gr.GPU(duration=60)
 def infer(
     prompt,
     negative_prompt="",
@@ -57,6 +66,7 @@ def infer(
     guidance_scale=0.0,
     num_inference_steps=4,
     style="Style Zero",
+    grid_size="1x1",
     progress=gr.Progress(track_tqdm=True),
 ):
     # Apply selected style
@@ -69,23 +79,37 @@ def infer(
 
     generator = torch.Generator().manual_seed(seed)
 
-    image = pipe(
-        prompt=styled_prompt,
-        negative_prompt=styled_negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
+    # Get grid size
+    grid_size_x, grid_size_y = grid_sizes.get(grid_size, (1, 1))
+    num_images = grid_size_x * grid_size_y
+
+    # Generate multiple images
+    options = {
+        "prompt": styled_prompt,
+        "negative_prompt": styled_negative_prompt,
+        "guidance_scale": guidance_scale,
+        "num_inference_steps": num_inference_steps,
+        "width": width,
+        "height": height,
+        "generator": generator,
+    }
+
+    torch.cuda.empty_cache()  # Clear GPU memory
+    images = pipe(**options).images
+
+    # Create grid image
+    grid_img = Image.new('RGB', (width * grid_size_x, height * grid_size_y))
+
+    for i, img in enumerate(images[:num_images]):
+        grid_img.paste(img, (i % grid_size_x * width, i // grid_size_x * height))
 
-    return image, seed
+    return grid_img, seed
 
 examples = [
     "Chocolate dripping from a donut against a yellow background, in the style of brocore, hyper-realistic oil --ar 2:3 --q 2 --s 750 --v 5  --ar 2:3 --q 2 --s 750 --v 5",
     "3d image, cute girl, in the style of Pixar --ar 1:2 --stylize 750, 4K resolution highlights, Sharp focus, octane render, ray tracing, Ultra-High-Definition, 8k, UHD, HDR, (Masterpiece:1.5), (best quality:1.5)",
     "Cold coffee in a cup bokeh --ar 85:128 --v 6.0 --style raw5, 4K, Photo-Realistic",
-    "Food photography of a milk shake with flying strawberrys against a pink background, professionally studio shot with cinematic lighting. The image is in the style of a professional studio shot --ar 85:128 --v 6.0 --style raw" 
+    "Food photography of a milk shake with flying strawberrys against a pink background, professionally studio shot with cinematic lighting. The image is in the style of a professional studio shot --ar 85:128 --v 6.0 --style raw"
 ]
 
 css = '''
@@ -99,7 +123,7 @@ footer {
 with gr.Blocks(css=css, theme="prithivMLmods/Minecraft-Theme") as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown("## SD3.5 TURBO")
-        
+
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
@@ -112,7 +136,7 @@ with gr.Blocks(css=css, theme="prithivMLmods/Minecraft-Theme") as demo:
             run_button = gr.Button("Run", scale=0, variant="primary")
 
         result = gr.Image(label="Result", show_label=False)
-        
+
         with gr.Row(visible=True):
             style_selection = gr.Radio(
                 show_label=True,
@@ -123,6 +147,13 @@ with gr.Blocks(css=css, theme="prithivMLmods/Minecraft-Theme") as demo:
                 label="Quality Style",
             )
 
+        with gr.Row(visible=True):
+            grid_size_selection = gr.Dropdown(
+                choices=["2x1", "1x2", "2x2", "2x3", "3x2", "1x1"],
+                value="1x1",
+                label="Grid Size"
+            )
+
         with gr.Accordion("Advanced Settings", open=False, visible=False):
             negative_prompt = gr.Text(
                 label="Negative prompt",
@@ -180,7 +211,7 @@ with gr.Blocks(css=css, theme="prithivMLmods/Minecraft-Theme") as demo:
                     outputs=[result, seed], 
                     fn=infer, 
                     cache_examples=True)
-        
+
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
@@ -194,9 +225,10 @@ with gr.Blocks(css=css, theme="prithivMLmods/Minecraft-Theme") as demo:
             guidance_scale,
             num_inference_steps,
             style_selection,
+            grid_size_selection,
         ],
         outputs=[result, seed],
     )
 
 if __name__ == "__main__":
-    demo.launch() 
+    demo.launch()