Update app.py

app.py

@@ -7,7 +7,6 @@ import uuid
 import os
 
 from diffusers import StableDiffusionXLPipeline, StableDiffusion3Pipeline
-from transformers import AutoModelForImageClassification, ViTImageProcessor
 from PIL import Image
 
 # Pre-Initialize
@@ -17,6 +16,8 @@ if DEVICE == "auto":
 print(f"[SYSTEM] | Using {DEVICE} type compute device.")
 
 # Variables
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
 MAX_SEED = 9007199254740991
 DEFAULT_INPUT = ""
 DEFAULT_NEGATIVE_INPUT = "(bad, ugly, amputation, abstract, blur, blurry, deformed, distorted, disfigured, disconnected, mutation, mutated, low quality, lowres), unfinished, title, text, signature, watermark, (limbs, legs, feet, arms, hands), (porn, nude, naked, nsfw)"
@@ -24,6 +25,8 @@ DEFAULT_MODEL = "Default"
 DEFAULT_HEIGHT = 1024
 DEFAULT_WIDTH = 1024
 
+headers = {"Content-Type": "application/json", "Authorization": f"Bearer {HF_TOKEN}" }
+
 css = '''
 .gradio-container{max-width: 560px !important}
 h1{text-align:center}
@@ -32,9 +35,6 @@ footer {
 }
 '''
 
-repo_nsfw_classifier = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
-processor_nsfw_classifier = ViTImageProcessor.from_pretrained("Falconsai/nsfw_image_detection")
-
 repo_default = StableDiffusionXLPipeline.from_pretrained("fluently/Fluently-XL-Final", torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False)
 #repo_default.load_lora_weights("ehristoforu/dalle-3-xl-v2", adapter_name="base")
 #repo_default.set_adapters(["base"], adapter_weights=[0.7])
@@ -65,6 +65,12 @@ def get_seed(seed):
     else:
         return random.randint(0, MAX_SEED)
 
+def api_classification_request(url, filename, headers):
+    with open(filename, "rb") as file:
+        data = file.read()
+    response = requests.request("POST", url, headers=headers or {}, data=data)
+    return json.loads(response.content.decode("utf-8"))
+
 @spaces.GPU(duration=60)
 def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATIVE_INPUT, model=DEFAULT_MODEL, height=DEFAULT_HEIGHT, width=DEFAULT_WIDTH, steps=1, guidance=0, number=1, seed=None):
 
@@ -100,7 +106,6 @@ def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATI
     
     print(steps, guidance)
 
-    repo_nsfw_classifier.to(DEVICE)
     repo.to(DEVICE)
     
     parameters  = {
@@ -120,11 +125,9 @@ def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATI
 
     print(image_paths)
     
-    nsfw_prediction = repo_nsfw_classifier(**processor_nsfw_classifier(images=Image.open(image_paths[0]), return_tensors="pt")).logits
+    nsfw_prediction = api_classification_request("https://api-inference.huggingface.co/models/Falconsai/nsfw_image_detection", image_paths[0], headers)
 
-    print(nsfw_prediction.argmax(-1))
-    print(nsfw_prediction.argmax(-1).item())
-    print(repo_nsfw_classifier.config.id2label[nsfw_prediction])
+    print(nsfw_prediction)
 
     return image_paths, {item['label']: round(item['score'], 3) for item in nsfw_prediction}