Update app.py

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 import numpy as np
 import random
+import uuid
+from PIL import Image
 
 import spaces
 from diffusers import DiffusionPipeline
@@ -47,6 +49,15 @@ style_list = [
 STYLE_NAMES = [style["name"] for style in style_list]
 DEFAULT_STYLE_NAME = STYLE_NAMES[0]
 
+grid_sizes = {
+    "2x1": (2, 1),
+    "1x2": (1, 2),
+    "2x2": (2, 2),
+    "2x3": (2, 3),
+    "3x2": (3, 2),
+    "1x1": (1, 1)
+}
+
 @spaces.GPU
 def infer(
     prompt,
@@ -58,6 +69,7 @@ def infer(
     guidance_scale=0.0,
     num_inference_steps=4,
     style="Style Zero",
+    grid_size="1x1",
     progress=gr.Progress(track_tqdm=True),
 ):
     
@@ -70,17 +82,33 @@ 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]
-
-    return image, seed
+    grid_size_x, grid_size_y = grid_sizes.get(grid_size, (2, 2))
+    num_images = grid_size_x * grid_size_y
+
+    images = []
+    for _ in range(num_images):
+        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]
+        images.append(image)
+
+    # Create a 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))
+
+    # Save the grid image
+    unique_name = str(uuid.uuid4()) + ".png"
+    grid_img.save(unique_name)
+    
+    return unique_name, seed
 
 examples = [
     "A capybara wearing a suit holding a sign that reads Hello World",
@@ -96,7 +124,7 @@ css = """
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(" # [Stable Diffusion 3.5 Large Turbo (8B)](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-turbo)")
-        gr.Markdown("[Learn more](https://stability.ai/news/introducing-stable-diffusion-3-5) about the Stable Diffusion 3.5 series. Try on [Stability AI API](https://platform.stability.ai/docs/api-reference#tag/Generate/paths/~1v2beta~1stable-image~1generate~1sd3/post), or [download model](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-turbo) to run locally with ComfyUI or diffusers.")
+        gr.Markdown("[Learn more](https://stability.ai/news/introducing-stable-diffusion-3-5) about the Stable Diffusion 3.5 series.")
         
         with gr.Row():
             prompt = gr.Text(
@@ -112,6 +140,13 @@ with gr.Blocks(css=css) as demo:
         result = gr.Image(label="Result", show_label=False)
         
         with gr.Row(visible=True):
+            grid_size_selection = gr.Dropdown(
+                choices=["2x1", "1x2", "2x2", "2x3", "3x2", "1x1"],
+                value="1x1",
+                label="Grid Size"
+            )
+
+        with gr.Row(visible=True):     
             style_selection = gr.Radio(
                 show_label=True,
                 container=True,
@@ -138,7 +173,7 @@ with gr.Blocks(css=css) as demo:
             )
 
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
+            
             with gr.Row():
                 width = gr.Slider(
                     label="Width",
@@ -173,6 +208,8 @@ with gr.Blocks(css=css) as demo:
                     value=4,
                 )
 
+
+
         gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=True, cache_mode="lazy")
         
     gr.on(
@@ -188,6 +225,7 @@ with gr.Blocks(css=css) as demo:
             guidance_scale,
             num_inference_steps,
             style_selection,
+            grid_size_selection,
         ],
         outputs=[result, seed],
     )