Added safety checker back

app.py

@@ -8,19 +8,27 @@ from diffusers import DPMSolverMultistepScheduler
 
 deviceStr = "cuda" if torch.cuda.is_available() else "cpu"
 device = torch.device(deviceStr)
+latents = None
+
+def GenerateNewLatentsForInference():
+    global latents
+    if deviceStr == "cuda":
+        latents = torch.randn((1, 4, 64, 64), device=device, dtype=torch.float16)
+    else:
+        latents = torch.randn((1, 4, 64, 64), device=device)
 
 if deviceStr == "cuda":
     pipeline = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting",
                                                               revision="fp16",
-                                                              torch_dtype=torch.float16,
-                                                              safety_checker=lambda images, **kwargs: (images, False))
+                                                              torch_dtype=torch.float16)
+                                                              #safety_checker=lambda images, **kwargs: (images, False))
     pipeline.to(device)
     pipeline.enable_xformers_memory_efficient_attention()
-    latents = torch.randn((1, 4, 64, 64), device=device, dtype=torch.float16)
 else:
-    pipeline = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting",
-                                                              safety_checker=lambda images, **kwargs: (images, False))                                                       
-    latents = torch.randn((1, 4, 64, 64), device=device)
+    pipeline = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting")
+                                                              #safety_checker=lambda images, **kwargs: (images, False))
+
+GenerateNewLatentsForInference()
 
 imageSize = (512, 512)
 lastImage = Image.new(mode="RGB", size=imageSize)
@@ -35,10 +43,7 @@ def diffuse(staticLatents, inputImage, mask, pauseInference, prompt, negativePro
         return lastImage
     
     if staticLatents is False:
-        if deviceStr == "cuda":
-            latents = torch.randn((1, 4, 64, 64), device=device, dtype=torch.float16)
-        else:
-            latents = torch.randn((1, 4, 64, 64), device=device)
+        GenerateNewLatentsForInference()
     
     if lastSeed != seed:
         generator = torch.Generator(device).manual_seed(seed)
@@ -69,6 +74,8 @@ numInferenceSteps = gradio.Slider(label="Number of Inference Steps", maximum=100
 seed = gradio.Slider(label="Generator Seed", maximum=10000, value=4096)
 staticLatents =gradio.Checkbox(label="Static Latents", value=True)
 pauseInference = gradio.Checkbox(label="Pause Inference", value=False)
+#generateNewLatents = gradio.Button(label="Generate New Latents")
+#generateNewLatents.click(GenerateNewLatentsForInference)
 
 inputs=[staticLatents, inputImage, mask, pauseInference, prompt, negativePrompt, guidanceScale, numInferenceSteps, seed]
 ux = gradio.Interface(fn=diffuse, title="View Diffusion", inputs=inputs, outputs=outputImage, live=True)