add NSFW checker and GPU mode

app.py

@@ -61,7 +61,13 @@ class GlobalText:
         self.pipeline              = None
         self.torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
         self.lora_model_state_dict = {}
-        self.device = torch.device("cpu")
+        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+        self.nsfw_image = Image.open('./data/nsfw.jpg') # to float in [0,1]
+        
+        
+        
+
 
     def init_source_image_path(self, source_path):
         self.source_paths = sorted(glob(os.path.join(source_path, '*')))
@@ -83,9 +89,9 @@ class GlobalText:
 
         self.scheduler = 'LCM'
         scheduler = LCMScheduler.from_pretrained(model_path, subfolder="scheduler")
-        self.pipeline = ZePoPipeline.from_pretrained(model_path,scheduler=scheduler,torch_dtype=torch.float16,)
-        # if is_xformers:
-        #     self.pipeline.enable_xformers_memory_efficient_attention()
+        self.pipeline = ZePoPipeline.from_pretrained(model_path,scheduler=scheduler,torch_dtype=torch.float16,).to('cuda')
+        if is_xformers:
+            self.pipeline.enable_xformers_memory_efficient_attention()
         time_end = datetime.now()
         print(f'Load {model_path} successful in {time_end-time_start}')
         return gr.Dropdown()
@@ -171,7 +177,7 @@ class GlobalText:
                                                  de_bug=de_bug,)
 
             time_begin = datetime.now()
-            generate_image = model(prompt=prompts,
+            results = model(prompt=prompts,
                                 negative_prompt=negative_prompt_textbox,
                                 image=source,
                                 style=style,
@@ -183,7 +189,16 @@ class GlobalText:
                                 fix_step_index=co_feat_step,
                                 de_bug = de_bug,
                                 callback = None
-                   ).images
+                   )
+            generate_image = results.images
+
+        
+            for idx, has_nsfw_concept in enumerate(results.nsfw_content_detected):
+                if has_nsfw_concept:
+                    generate_image[idx] = np.array(self.nsfw_image.resize((height_slider,width_slider))).astype(np.float32) / 255.0
+
+                
+
             time_end = datetime.now()
             print('generate one image with time {}'.format(time_end-time_begin))
 
@@ -191,18 +206,19 @@ class GlobalText:
 
             
             save_file_path = os.path.join(self.savedir, save_file_name)
-                    
+                   
             save_image(torch.tensor(generate_image).permute(0, 3, 1, 2), save_file_path, nrow=3, padding=0)
             save_image(torch.tensor(generate_image[2:]).permute(0, 3, 1, 2), os.path.join(self.savedir_sample, save_file_name), nrow=3, padding=0)
             self.init_results_image_path()
-        return [
-            generate_image[0],
-            generate_image[1],
-            generate_image[2],
-            self.init_results_image_path()
-            ]
-    
 
+            return [
+                generate_image[0],
+                generate_image[1],
+                generate_image[2],
+                self.init_results_image_path()
+                ]
+        
+        
 global_text = GlobalText()
 
 
@@ -309,23 +325,23 @@ def ui():
                                 
                                 style_gallery_index.change(fn=update_style_list, inputs=[style_gallery_index], outputs=[style_image_gallery])
                             
-                            with gr.Tab("Results Gallery"):
-                                with gr.Row():
-                                    refresh_results_list_button = gr.Button(value="\U0001F503", elem_classes="toolbutton")
-                                    results_gallery_index = gr.Slider(label="Index", value=0, minimum=0, maximum=50, step=1)   
-                                num_gallery_images = 12   
-                                results_image_gallery = gr.Gallery(value=[], columns=4, label="style Image List")
-                                refresh_results_list_button.click(fn=global_text.init_results_image_path, inputs=[], outputs=[results_image_gallery]) 
+                            # with gr.Tab("Results Gallery"):
+                            #     with gr.Row():
+                            #         refresh_results_list_button = gr.Button(value="\U0001F503", elem_classes="toolbutton")
+                            #         results_gallery_index = gr.Slider(label="Index", value=0, minimum=0, maximum=50, step=1)   
+                            #     num_gallery_images = 12   
+                            #     results_image_gallery = gr.Gallery(value=[], columns=4, label="style Image List")
+                            #     refresh_results_list_button.click(fn=global_text.init_results_image_path, inputs=[], outputs=[results_image_gallery]) 
                                 
 
-                                def update_results_list(index):
-                                    if int(index) < 0:
-                                        index = 0
-                                    if int(index) > global_text.max_results_index:
-                                        index = global_text.max_results_index
-                                    return global_text.results_paths[int(index)*num_gallery_images:(int(index)+1)*num_gallery_images]
+                            #     def update_results_list(index):
+                            #         if int(index) < 0:
+                            #             index = 0
+                            #         if int(index) > global_text.max_results_index:
+                            #             index = global_text.max_results_index
+                            #         return global_text.results_paths[int(index)*num_gallery_images:(int(index)+1)*num_gallery_images]
                                 
-                                results_gallery_index.change(fn=update_results_list, inputs=[results_gallery_index], outputs=[style_image_gallery])
+                            #     results_gallery_index.change(fn=update_results_list, inputs=[results_gallery_index], outputs=[style_image_gallery])
                             
                             
 

utils/pipeline.py

@@ -157,6 +157,20 @@ class ZePoPipeline(DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMix
             extra_step_kwargs["generator"] = generator
         return extra_step_kwargs
 
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker
+    def run_safety_checker(self, image, device, dtype):
+        if self.safety_checker is None:
+            has_nsfw_concept = None
+        else:
+            if torch.is_tensor(image):
+                feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")
+            else:
+                feature_extractor_input = self.image_processor.numpy_to_pil(image)
+            safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)
+            image, has_nsfw_concept = self.safety_checker(
+                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
+            )
+        return image, has_nsfw_concept
    
     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents
     def decode_latents(self, latents):
@@ -416,7 +430,7 @@ class ZePoPipeline(DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMix
         # 9. Post-processing
         if not output_type == "latent":
             image = self.vae.decode(pred_x0 / self.vae.config.scaling_factor, return_dict=False)[0]
-            has_nsfw_concept = None
+            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
         else:
             image = pred_x0
             has_nsfw_concept = None