enable better output

app.py

@@ -13,6 +13,8 @@ current_steps = 15
 
 pipe = DiffusionPipeline.from_pretrained("timbrooks/instruct-pix2pix", torch_dtype=torch.float16, safety_checker=None)
 pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+pipe.enable_xformers_memory_efficient_attention()
+pipe.unet.to(memory_format=torch.channels_last)
 
 device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
 
@@ -62,9 +64,9 @@ def inference(
         )
 
         # return replace_nsfw_images(result)
-        return result.images, f"Done. Seed: {seed}"
+        return result.images, result.nsfw_content_detected, seed
     except Exception as e:
-        return None, error_str(e)
+        return None, None, error_str(e)
 
 
 def replace_nsfw_images(results):
@@ -119,7 +121,8 @@ with gr.Blocks(css="style.css") as demo:
             state_info = gr.Textbox(label="State", show_label=False, max_lines=2).style(
                 container=False
             )
-            error_output = gr.Markdown()
+            nsfw_output = gr.JSON()
+            error_output = gr.JSON()
 
         with gr.Column(scale=45):
             with gr.Tab("Options"):
@@ -180,7 +183,7 @@ with gr.Blocks(css="style.css") as demo:
         height,
         seed,
     ]
-    outputs = [gallery, error_output]
+    outputs = [gallery, nsfw_output, error_output]
     prompt.submit(inference, inputs=inputs, outputs=outputs)
     generate.click(inference, inputs=inputs, outputs=outputs)
 
@@ -197,4 +200,4 @@ with gr.Blocks(css="style.css") as demo:
 print(f"Space built in {time.time() - start_time:.2f} seconds")
 
 demo.queue(concurrency_count=1)
-demo.launch()
+demo.launch(debug=True, show_api=False)

requirements.txt

@@ -7,4 +7,8 @@ scipy
 ftfy
 psutil
 accelerate
-safetensors
+safetensors
+transformers
+safetensors
+--pre
+xformers