:tada: add external model

app.py

@@ -3,6 +3,7 @@ import spaces
 import torch
 from PIL import Image
 from diffusers import DiffusionPipeline
+from diffusers.utils import make_image_grid
 
 DIFFUSERS_MODEL_IDS = [
     # SD Models
@@ -32,17 +33,24 @@ if device == 'cuda':
     ).to(device)
 
 
-
 @spaces.GPU()
 @torch.inference_mode()
 def inference(
         model_id: str,
         prompt: str,
         negative_prompt: str = "",
+        width: int = 512,
+        height: int = 512,
+        guidance_scale: float = 7.5,
+        num_inference_steps: int = 50,
+        num_images: int = 4,
         progress=gr.Progress(track_tqdm=True),
 ) -> Image.Image:
     progress(0, "Starting inference...")
 
+    if device != 'cuda':
+        raise gr.Error("This model requires a GPU to run. Please switch to a GPU runtime.")
+
     global current_model_id, pipe
 
     if model_id != current_model_id:
@@ -58,13 +66,21 @@ def inference(
         except Exception as e:
             raise gr.Error(str(e))
 
-    if device != 'cuda':
-        raise gr.Error("This model requires a GPU to run. Please switch to a GPU runtime.")
-
-    image = pipe(
+    # Generation
+    images = pipe(
         prompt,
         negative_prompt=negative_prompt,
-    ).images[0]
+        width=width,
+        height=height,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        num_images_per_prompt=num_images,
+    ).images
+
+    if num_images % 2 == 1:
+        image = make_image_grid(images, rows=num_images, cols=1)
+    else:
+        image = make_image_grid(images, rows=2, cols=num_images // 2)
 
     return image
 
@@ -75,41 +91,54 @@ if __name__ == "__main__":
 
         with gr.Row():
             with gr.Column():
-                inputs = [
-                    gr.Dropdown(
-                        label="Model ID",
-                        choices=MODEL_CHOICES,
-                        value="stabilityai/stable-diffusion-3-medium-diffusers",
-                    ),
-                    gr.Text(label="Prompt", value=""),
-                    gr.Text(label="Negative Prompt", value=""),
-                ]
+                model_id = gr.Dropdown(
+                    label="Model ID",
+                    choices=MODEL_CHOICES,
+                    value="stabilityai/stable-diffusion-3-medium-diffusers",
+                )
+                prompt = gr.Text(label="Prompt", value="")
 
                 with gr.Accordion("Additional Settings (W.I.P)", open=False):
+                    negative_prompt = gr.Text(label="Negative Prompt", value="")
+
                     with gr.Row():
-                        width_component = gr.Number(label="Width", value=512, step=64, minimum=64, maximum=1024)
-                        height_component = gr.Number(label="Height", value=512, step=64, minimum=64, maximum=1024)
+                        width = gr.Number(label="Width", value=512, step=64, minimum=64, maximum=2048)
+                        height = gr.Number(label="Height", value=512, step=64, minimum=64, maximum=2048)
+                        num_images = gr.Number(label="Num Images", value=4, minimum=1, maximum=10, step=1)
 
-                    additional_inputs = [
-                        width_component,
-                        height_component,
-                        gr.Number(label="Guidance Scale", value=7.5, step=0.5, minimum=0, maximum=10),
-                        gr.Slider(label="Num Inference Steps", value=None, minimum=1, maximum=1000, step=1)
-                    ]
+                    guidance_scale = gr.Slider(label="Guidance Scale", value=7.5, step=0.5, minimum=0, maximum=10)
+                    num_inference_step = gr.Slider(label="Num Inference Steps", value=50, minimum=1, maximum=100,
+                                                   step=1)
 
             with gr.Column():
-                outputs = [
-                    gr.Image(label="Image", type="pil"),
-                ]
+                output_image = gr.Image(label="Image", type="pil")
+
+        inputs = [
+            model_id,
+            prompt,
+            negative_prompt,
+            width,
+            height,
+            guidance_scale,
+            num_inference_step,
+            num_images,
+        ]
 
         btn = gr.Button("Generate")
-        btn.click(fn=inference, inputs=inputs, outputs=outputs)
+        btn.click(
+            fn=inference,
+            inputs=inputs,
+            outputs=output_image
+        )
 
         gr.Examples(
             examples=[
                 ['stabilityai/stable-diffusion-3-medium-diffusers', 'A cat holding a sign that says Hello world', ""],
-                ['stabilityai/stable-diffusion-3-medium-diffusers', 'Beautiful pixel art of a Wizard with hovering text "Achivement unlocked: Diffusion models can spell now"', ''],
-                ['stabilityai/stable-diffusion-3-medium-diffusers', 'A corgi wearing sunglasses says "U-Net is OVER!!"', ''],
+                ['stabilityai/stable-diffusion-3-medium-diffusers',
+                 'Beautiful pixel art of a Wizard with hovering text "Achivement unlocked: Diffusion models can spell now"',
+                 ''],
+                ['stabilityai/stable-diffusion-3-medium-diffusers', 'A corgi wearing sunglasses says "U-Net is OVER!!"',
+                 ''],
             ],
             inputs=inputs,
         )

example.csv

@@ -0,0 +1,4 @@
+model_uri,description,additional_info
+stabilityai/stable-diffusion-3-medium-diffusers,"A cat holding a sign that says Hello world",
+stabilityai/stable-diffusion-3-medium-diffusers,"Beautiful pixel art of a Wizard with hovering text "Achivement unlocked: Diffusion models can spell now"",
+stabilityai/stable-diffusion-3-medium-diffusers,"A corgi wearing sunglasses says ""U-Net is OVER!!"",

requirements.txt

@@ -7,3 +7,5 @@ transformers
 ftfy
 accelerate
 sentencepiece
+spaces
+Pillow