Geek7 commited on
Commit
7468f56
·
verified ·
1 Parent(s): aabdb2a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +25 -26
app.py CHANGED
@@ -14,38 +14,38 @@ CORS(app) # Enable CORS for all routes
14
  HF_TOKEN = os.environ.get("HF_TOKEN") # Ensure to set your Hugging Face token in the environment
15
  client = InferenceClient(token=HF_TOKEN)
16
 
17
- # Hardcoded negative prompt
18
- NEGATIVE_PROMPT_FINGERS = """2D,missing fingers, extra fingers, elongated fingers, fused fingers,
19
- mutated fingers, poorly drawn fingers, disfigured fingers,
20
- too many fingers, deformed hands, extra hands, malformed hands,
21
- blurry hands, disproportionate fingers"""
22
 
23
- @app.route('/')
24
- def home():
25
- return "Welcome to the Image Background Remover!"
26
 
27
- # Simple content moderation function
28
  def is_prompt_explicit(prompt):
29
- # Streamlined keyword list to avoid unnecessary restrictions
30
- explicit_keywords = [
31
- "sexual","boobs","boob","penis","porn", "hentai", "fetish", "nude", "provocative", "obscene", "vulgar", "intimate", "kinky", "hardcore",
32
- "threesome", "orgy", "masturbation", "genital", "suicide",
33
- "self-harm", "depression", "kill myself", "worthless","vagina","intercourse","fuck","fucking","fuckers","fucker"
34
- ]
35
- for keyword in explicit_keywords:
36
- if keyword.lower() in prompt.lower():
37
- return True
38
- return False
 
 
 
39
 
40
  # Function to generate an image from a text prompt
41
  def generate_image(prompt, negative_prompt=None, height=512, width=512, model="stabilityai/stable-diffusion-2-1", num_inference_steps=50, guidance_scale=7.5, seed=None):
42
  try:
43
  # Generate the image using Hugging Face's inference API with additional parameters
44
  image = client.text_to_image(
45
- prompt=prompt,
46
  negative_prompt=NEGATIVE_PROMPT_FINGERS,
47
- height=height,
48
- width=width,
49
  model=model,
50
  num_inference_steps=num_inference_steps, # Control the number of inference steps
51
  guidance_scale=guidance_scale, # Control the guidance scale
@@ -77,7 +77,7 @@ def generate_api():
77
  try:
78
  # Check for explicit content
79
  if is_prompt_explicit(prompt):
80
- # Return the pre-defined "thinkgood.png" image
81
  return send_file(
82
  "nsfw.jpg",
83
  mimetype='image/png',
@@ -96,8 +96,8 @@ def generate_api():
96
 
97
  # Send the generated image as a response
98
  return send_file(
99
- img_byte_arr,
100
- mimetype='image/png',
101
  as_attachment=False, # Send the file as an attachment
102
  download_name='generated_image.png' # The file name for download
103
  )
@@ -110,5 +110,4 @@ def generate_api():
110
  # Add this block to make sure your app runs when called
111
  if __name__ == "__main__":
112
  subprocess.Popen(["python", "wk.py"]) # Start awake.py
113
-
114
  app.run(host='0.0.0.0', port=7860) # Run directly if needed for testing
 
14
  HF_TOKEN = os.environ.get("HF_TOKEN") # Ensure to set your Hugging Face token in the environment
15
  client = InferenceClient(token=HF_TOKEN)
16
 
17
+ # Initialize NSFW model
18
+ NSFW_MODEL = "MichalMlodawski/nsfw-text-detection-large"
19
+ nsfw_client = InferenceClient(model=NSFW_MODEL, token=HF_TOKEN)
 
 
20
 
21
+ # Hardcoded negative prompt
22
+ NEGATIVE_PROMPT_FINGERS = """2D,missing fingers, extra fingers, elongated fingers, fused fingers, mutated fingers, poorly drawn fingers, disfigured fingers, too many fingers, deformed hands, extra hands, malformed hands, blurry hands, disproportionate fingers"""
 
23
 
24
+ # NSFW detection function
25
  def is_prompt_explicit(prompt):
26
+ try:
27
+ response = nsfw_client(prompt, task="text-classification")
28
+ if "error" in response:
29
+ print(f"Error in NSFW detection: {response['error']}")
30
+ return False
31
+
32
+ # Parse the classification result
33
+ predicted_class = response[0]["label"] # E.g., "LABEL_2"
34
+ class_id = int(predicted_class.split("_")[-1]) # Extract numerical label
35
+ return class_id == 2 # Class 2 indicates UNSAFE
36
+ except Exception as e:
37
+ print(f"Error in NSFW detection: {str(e)}")
38
+ return False
39
 
40
  # Function to generate an image from a text prompt
41
  def generate_image(prompt, negative_prompt=None, height=512, width=512, model="stabilityai/stable-diffusion-2-1", num_inference_steps=50, guidance_scale=7.5, seed=None):
42
  try:
43
  # Generate the image using Hugging Face's inference API with additional parameters
44
  image = client.text_to_image(
45
+ prompt=prompt,
46
  negative_prompt=NEGATIVE_PROMPT_FINGERS,
47
+ height=height,
48
+ width=width,
49
  model=model,
50
  num_inference_steps=num_inference_steps, # Control the number of inference steps
51
  guidance_scale=guidance_scale, # Control the guidance scale
 
77
  try:
78
  # Check for explicit content
79
  if is_prompt_explicit(prompt):
80
+ # Return the pre-defined "nsfw.jpg" image
81
  return send_file(
82
  "nsfw.jpg",
83
  mimetype='image/png',
 
96
 
97
  # Send the generated image as a response
98
  return send_file(
99
+ img_byte_arr,
100
+ mimetype='image/png',
101
  as_attachment=False, # Send the file as an attachment
102
  download_name='generated_image.png' # The file name for download
103
  )
 
110
  # Add this block to make sure your app runs when called
111
  if __name__ == "__main__":
112
  subprocess.Popen(["python", "wk.py"]) # Start awake.py
 
113
  app.run(host='0.0.0.0', port=7860) # Run directly if needed for testing