:tada: change

app.py

@@ -4,20 +4,28 @@ import torch
 from PIL import Image
 from diffusers import DiffusionPipeline
 
-MODEL_CHOICES = [
+DIFFUSERS_MODEL_IDS = [
     "stabilityai/stable-diffusion-3-medium-diffusers",
     "stabilityai/stable-diffusion-xl-base-1.0",
     "stabilityai/stable-diffusion-2-1",
     "runwayml/stable-diffusion-v1-5",
 ]
 
+EXTERNAL_MODEL_URL_MAPPING = {
+    "Beautiful Realistic Asians": "https://civitai.com/api/download/models/177164?type=Model&format=SafeTensor&size=full&fp=fp16",
+}
+
+MODEL_CHOICES = DIFFUSERS_MODEL_IDS + list(EXTERNAL_MODEL_URL_MAPPING.keys())
+
 # Global Variables
 current_model_id = "stabilityai/stable-diffusion-3-medium-diffusers"
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-pipe = DiffusionPipeline.from_pretrained(
-    current_model_id,
-    torch_dtype=torch.float16,
-).to(device)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+if device == 'cuda':
+    pipe = DiffusionPipeline.from_pretrained(
+        current_model_id,
+        torch_dtype=torch.float16,
+    ).to(device)
 
 
 @spaces.GPU()
@@ -26,6 +34,10 @@ 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 | None = None,
         progress=gr.Progress(track_tqdm=True),
 ) -> Image.Image:
     global current_model_id, pipe
@@ -43,6 +55,10 @@ def inference(
     image = pipe(
         prompt,
         negative_prompt=negative_prompt,
+        width=width,
+        height=height,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
     ).images[0]
 
     return image
@@ -65,13 +81,15 @@ if __name__ == "__main__":
                 ]
 
                 with gr.Accordion("Additional Settings (W.I.P)", open=False):
+                    with gr.Row():
+                        width = gr.Number(label="Width", value=512, step=64, minimum=64, maximum=1024)
+                        height = gr.Number(label="Height", value=512, step=64, minimum=64, maximum=1024)
+
                     additional_inputs = [
-                        gr.Text(
-                            label="Model URL",
-                            lines=2,
-                            placeholder="e.g. ) https://civitai.com/api/download/models/177164?type=Model&format=SafeTensor&size=full&fp=fp16"
-                        ),
-                        gr.Number(label="Num Inference Steps", value=None, minimum=1, maximum=1000, step=1)
+                        width,
+                        height,
+                        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)
                     ]
 
             with gr.Column():
@@ -87,6 +105,6 @@ if __name__ == "__main__":
         )
 
         btn = gr.Button("Generate")
-        btn.click(fn=inference, inputs=inputs, outputs=outputs)
+        btn.click(fn=inference, inputs=inputs + additional_inputs, outputs=outputs)
 
     demo.queue().launch()