Update app.py
Browse files
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 |
-
#
|
18 |
-
|
19 |
-
|
20 |
-
too many fingers, deformed hands, extra hands, malformed hands,
|
21 |
-
blurry hands, disproportionate fingers"""
|
22 |
|
23 |
-
|
24 |
-
|
25 |
-
return "Welcome to the Image Background Remover!"
|
26 |
|
27 |
-
#
|
28 |
def is_prompt_explicit(prompt):
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
|
36 |
-
|
37 |
-
|
38 |
-
|
|
|
|
|
|
|
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 "
|
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
|