Update main.py

main.py

@@ -288,20 +288,21 @@ def get_random_model(models):
     last_two_models = last_two_models[-5:]        
     return model
    
-def nsfw_check(attempts=1):
+def nsfw_check(item, attempts=1):
     try:
         API_URL = "https://api-inference.huggingface.co/models/Falconsai/nsfw_image_detection"
         with open('response.png', 'rb') as f:
             data = f.read()
         response = requests.request("POST", API_URL, headers=headers, data=data)
         decoded_response = response.content.decode("utf-8")
+        print(item.prompt)
         print(decoded_response)
         
         json_response = json.loads(decoded_response)
         
         if "error" in json_response:
             time.sleep(json_response["estimated_time"])
-            return nsfw_check(attempts+1)
+            return nsfw_check(item, attempts+1)
         
         scores = {item['label']: item['score'] for item in json_response}
         error_msg = scores.get('nsfw', 0) > scores.get('normal', 0)
@@ -313,7 +314,7 @@ def nsfw_check(attempts=1):
         print(f'NSFW Check Error: {e}')
         if attempts > 30:
             return True
-        return nsfw_check(attempts+1)
+        return nsfw_check(item, attempts+1)
     
     
 @app.post("/api")
@@ -350,7 +351,7 @@ async def inference(item: Item):
             base64_img, model = inferenceAPI(item.modelID, item)
         if 'error' in base64_img:
             return {"output": base64_img, "model": model}
-        NSFW = nsfw_check()
+        NSFW = nsfw_check(item)
             
         save_image(base64_img, item, model, NSFW)
     except Exception as e: