Update app.py

app.py

@@ -5,9 +5,11 @@ import spaces
 import torch
 import uuid
 import os
+from transformers import pipeline
 
 from diffusers import StableDiffusionXLPipeline, ControlNetModel
 from diffusers.models import AutoencoderKL
+from PIL import Image
 
 # Pre-Initialize
 DEVICE = "auto"
@@ -38,19 +40,19 @@ repo_default = StableDiffusionXLPipeline.from_pretrained("sd-community/sdxl-flas
 repo_default.load_lora_weights("ehristoforu/dalle-3-xl-v2", adapter_name="base")
 repo_default.set_adapters(["base"], adapter_weights=[0.7])
 
-repo_pixel = StableDiffusionXLPipeline.from_pretrained("sd-community/sdxl-flash", vae=vae, controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False)
-repo_pixel.load_lora_weights("artificialguybr/PixelArtRedmond", adapter_name="base")
-repo_pixel.load_lora_weights("nerijs/pixel-art-xl", adapter_name="base2")
-repo_pixel.set_adapters(["base", "base2"], adapter_weights=[1, 1])
+#repo_pixel = StableDiffusionXLPipeline.from_pretrained("sd-community/sdxl-flash", vae=vae, controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False)
+#repo_pixel.load_lora_weights("artificialguybr/PixelArtRedmond", adapter_name="base")
+#repo_pixel.load_lora_weights("nerijs/pixel-art-xl", adapter_name="base2")
+#repo_pixel.set_adapters(["base", "base2"], adapter_weights=[1, 1])
 
-repo_large = StableDiffusionXLPipeline.from_pretrained("Corcelio/mobius", vae=vae, controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False)
+#repo_large = StableDiffusionXLPipeline.from_pretrained("Corcelio/mobius", vae=vae, controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False)
 
 repo_customs = {
     "Default": repo_default,
-    "Realistic": StableDiffusionXLPipeline.from_pretrained("stablediffusionapi/NightVision_XL", vae=vae, controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=False, add_watermarker=False),
-    "Anime": StableDiffusionXLPipeline.from_pretrained("cagliostrolab/animagine-xl-3.1", vae=vae, controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False),
-    "Pixel": repo_pixel,
-    "Large": repo_large,
+    "Realistic": None, #StableDiffusionXLPipeline.from_pretrained("stablediffusionapi/NightVision_XL", vae=vae, controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=False, add_watermarker=False),
+    "Anime": None, #StableDiffusionXLPipeline.from_pretrained("cagliostrolab/animagine-xl-3.1", vae=vae, controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False),
+    "Pixel": None, #repo_pixel,
+    "Large": None, #repo_large,
 }
 
 # Functions
@@ -120,8 +122,13 @@ def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATI
     
     images = repo(**parameters).images
     image_paths = [save_image(img, seed) for img in images]
+
+    classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")(Image.open(image_paths[0]))
+    
     print(image_paths)
-    return image_paths
+    print(classifier)
+    
+    return image_paths, classifier
 
 def cloud():
     print("[CLOUD] | Space maintained.")
@@ -148,8 +155,9 @@ with gr.Blocks(css=css) as main:
 
     with gr.Column():
         images = gr.Gallery(columns=1, label="Image")
+        classifier = gr.Label()
             
-    submit.click(generate, inputs=[input, filter_input, negative_input, model, height, width, steps, guidance, number, seed], outputs=[images], queue=False)
+    submit.click(generate, inputs=[input, filter_input, negative_input, model, height, width, steps, guidance, number, seed], outputs=[images, classifier], queue=False)
     maintain.click(cloud, inputs=[], outputs=[], queue=False)
 
 main.launch(show_api=True)