default and safe

app.py

@@ -24,7 +24,6 @@ def infer(prompt, n_samples, steps, scale, seed):
     # generator = torch.Generator(device=device).manual_seed(seed)
     # print("Is GPU busy? ", is_gpu_busy)
     images = []
-    images_default = []
 
     url = os.getenv('BACKEND_URL_DEFAULT')
 
@@ -41,8 +40,6 @@ def infer(prompt, n_samples, steps, scale, seed):
         "User-Agent": "YOUR_NAME_HERE"
     })
 
-    #response = requests.get(url.format(prompt, int(n_samples), max(50,int(steps)), f'{scale:.1f}', int(seed)))
-    #response = requests.get(url.format('a%20naked%20girl', 2, 50, 7.5, 2))
     data = json.load(BytesIO(response.content))
     if 'output' not in data:
         raise gr.Error("Although safety guidance is enabled, potential unsafe content found. Please try again with different seed.")
@@ -53,6 +50,29 @@ def infer(prompt, n_samples, steps, scale, seed):
             im = Image.open(BytesIO(base64.b64decode(image['image_base64'])))
             images.append(im)
 
+    response = requests.post(url, json={
+        "model": "SafeStableDiffusion",
+        "prompt": prompt,
+        "n": int(n_samples),
+        "height": 512,
+        "width": 512,
+        "steps": max(50, int(steps)),
+        "seed": int(seed),
+        "guidance_scale": scale,
+    }, headers={
+        "User-Agent": "YOUR_NAME_HERE"
+    })
+
+    data = json.load(BytesIO(response.content))
+    if 'output' not in data:
+        raise gr.Error(
+            "Although safety guidance is enabled, potential unsafe content found. Please try again with different seed.")
+    else:
+        if data['output']['result_type'] == "error":
+            raise gr.Error(data['output']['value'])
+        for image in data['output']['choices']:
+            im = Image.open(BytesIO(base64.b64decode(image['image_base64'])))
+            images.append(im)
     return images
 
 
@@ -260,7 +280,7 @@ with block:
 
         gallery = gr.Gallery(
             label="Generated images", show_label=False, elem_id="gallery"
-        ).style(grid=[1], height="auto")
+        ).style(grid=[2], height="auto")
 
         with gr.Group(elem_id="container-advanced-btns"):
             advanced_button = gr.Button("Advanced options", elem_id="advanced-btn")