From b873c4b70837dce28a2c7d14930a0b42d86e4cf7 Mon Sep 17 00:00:00 2001
From: Pedro Cuenca <pedro@huggingface.co>
Date: Wed, 1 Feb 2023 13:19:38 +0100
Subject: [PATCH] Allow safety checker to be disabled.

---
 Diffusion-macOS/ControlsView.swift | 33 ++++++++++++++++++++++++++++++
 Diffusion/State.swift              | 13 +++++++++++-
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/Diffusion-macOS/ControlsView.swift b/Diffusion-macOS/ControlsView.swift
index c3a967e..3619820 100644
--- a/Diffusion-macOS/ControlsView.swift
+++ b/Diffusion-macOS/ControlsView.swift
@@ -62,6 +62,13 @@ struct ControlsView: View {
     @State private var stateSubscriber: Cancellable?
     @State private var pipelineState: PipelineState = .downloading(0)
 
+    // TODO: make this computed, and observable, and easy to read
+    @State private var mustShowSafetyCheckerDisclaimer = false
+    
+    func updateSafetyCheckerState() {
+        mustShowSafetyCheckerDisclaimer = generation.disableSafety && !Settings.shared.safetyCheckerDisclaimerShown
+    }
+    
     func modelDidChange(model: ModelInfo) {
         print("Loading model \(model)")
         Settings.shared.currentModel = model
@@ -182,6 +189,32 @@ struct ControlsView: View {
             }
             .disclosureGroupStyle(LabelToggleDisclosureGroupStyle())
             
+            Toggle("Disable Safety Checker", isOn: $generation.disableSafety).onChange(of: generation.disableSafety) { value in
+                updateSafetyCheckerState()
+            }
+                .popover(isPresented: $mustShowSafetyCheckerDisclaimer) {
+                        VStack {
+                            Text("You have disabled the safety checker").font(.title).padding(.top)
+                            Text("""
+                                 Please, ensure that you abide \
+                                 by the conditions of the Stable Diffusion license and do not expose \
+                                 unfiltered results to the public.
+                                 """)
+                            .lineLimit(nil)
+                            .padding(.all, 5)
+                            Button {
+                                Settings.shared.safetyCheckerDisclaimerShown = true
+                                updateSafetyCheckerState()
+                            } label: {
+                                Text("I Accept").frame(maxWidth: 200)
+                            }
+                            .padding(.bottom)
+                        }
+                        .frame(minWidth: 400, idealWidth: 400, maxWidth: 400)
+                        .fixedSize()
+                    }
+            Divider()
+            
             StatusView(pipelineState: $pipelineState)
         }
         .padding()
diff --git a/Diffusion/State.swift b/Diffusion/State.swift
index 13c9c49..167851d 100644
--- a/Diffusion/State.swift
+++ b/Diffusion/State.swift
@@ -71,11 +71,13 @@ class Settings {
     
     enum Keys: String {
         case model
+        case safetyCheckerDisclaimer
     }
     
     private init() {
         defaults.register(defaults: [
-            Keys.model.rawValue: ModelInfo.v2Base.modelId
+            Keys.model.rawValue: ModelInfo.v2Base.modelId,
+            Keys.safetyCheckerDisclaimer.rawValue: false
         ])
     }
     
@@ -88,4 +90,13 @@ class Settings {
             return ModelInfo.from(modelId: modelId) ?? DEFAULT_MODEL
         }
     }
+    
+    var safetyCheckerDisclaimerShown: Bool {
+        set {
+            defaults.set(newValue, forKey: Keys.safetyCheckerDisclaimer.rawValue)
+        }
+        get {
+            return defaults.bool(forKey: Keys.safetyCheckerDisclaimer.rawValue)
+        }
+    }
 }