Update app.py
Browse files
app.py
CHANGED
@@ -0,0 +1,123 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from flask import Flask, request, jsonify, send_file
|
2 |
+
from flask_cors import CORS
|
3 |
+
import os
|
4 |
+
from huggingface_hub import InferenceClient
|
5 |
+
from io import BytesIO
|
6 |
+
from PIL import Image
|
7 |
+
|
8 |
+
# Initialize the Flask app
|
9 |
+
app = Flask(__name__)
|
10 |
+
CORS(app) # Enable CORS for all routes
|
11 |
+
|
12 |
+
# Initialize the InferenceClient with your Hugging Face token
|
13 |
+
HF_TOKEN = os.environ.get("HF_TOKEN") # Ensure to set your Hugging Face token in the environment
|
14 |
+
client = InferenceClient(token=HF_TOKEN)
|
15 |
+
|
16 |
+
@app.route('/')
|
17 |
+
def home():
|
18 |
+
return "Welcome to the Image Background Remover!"
|
19 |
+
|
20 |
+
# Simple content moderation function
|
21 |
+
def is_prompt_explicit(prompt):
|
22 |
+
explicit_keywords = ["sexual", "nudity", "erotic", "explicit", "porn", "pornographic", "xxx", "hentai", "fetish", "sex", "sensual", "nude", "strip", "stripping", "adult", "lewd", "provocative", "obscene", "vulgar", "intimacy", "intimate", "lust", "arouse", "seductive", "seduction", "kinky", "bdsm", "dominatrix", "bondage", "hardcore", "softcore", "topless", "bottomless", "threesome", "orgy", "incest", "taboo", "masturbation", "genital", "penis", "vagina", "breast", "boob", "nipple", "butt", "anal", "oral", "ejaculation", "climax", "moan", "foreplay", "intercourse", "naked", "exposed", "suicide", "self-harm", "overdose", "poison", "hang", "end life", "kill myself", "noose", "depression", "hopeless", "worthless", "die", "death", "harm myself"] # Add more keywords as needed
|
23 |
+
for keyword in explicit_keywords:
|
24 |
+
if keyword.lower() in prompt.lower():
|
25 |
+
return True
|
26 |
+
return False
|
27 |
+
|
28 |
+
# Function to generate an image from a text prompt
|
29 |
+
def generate_image(prompt, negative_prompt=None, height=512, width=512, model="stabilityai/sd-3.5", num_inference_steps=50, guidance_scale=7.5, seed=None):
|
30 |
+
try:
|
31 |
+
# Generate the image using Hugging Face's inference API with additional parameters
|
32 |
+
image = client.text_to_image(
|
33 |
+
prompt=prompt,
|
34 |
+
negative_prompt=negative_prompt,
|
35 |
+
height=height,
|
36 |
+
width=width,
|
37 |
+
model=model,
|
38 |
+
num_inference_steps=num_inference_steps, # Control the number of inference steps
|
39 |
+
guidance_scale=guidance_scale, # Control the guidance scale
|
40 |
+
seed=seed # Control the seed for reproducibility
|
41 |
+
)
|
42 |
+
return image # Return the generated image
|
43 |
+
except Exception as e:
|
44 |
+
print(f"Error generating image: {str(e)}")
|
45 |
+
return None
|
46 |
+
|
47 |
+
# Function to refine an image using the refiner model
|
48 |
+
def refine_image(image, prompt, negative_prompt=None, model="stabilityai/stable-diffusion-xl-refiner-1.0", num_inference_steps=50, guidance_scale=7.5):
|
49 |
+
try:
|
50 |
+
# Use Hugging Face's image-to-image API to refine the image
|
51 |
+
refined_image = client.image_to_image(
|
52 |
+
prompt=prompt,
|
53 |
+
negative_prompt=negative_prompt,
|
54 |
+
image=image,
|
55 |
+
model=model,
|
56 |
+
num_inference_steps=num_inference_steps,
|
57 |
+
guidance_scale=guidance_scale
|
58 |
+
)
|
59 |
+
return refined_image
|
60 |
+
except Exception as e:
|
61 |
+
print(f"Error refining image: {str(e)}")
|
62 |
+
return None
|
63 |
+
|
64 |
+
@app.route('/generate_image', methods=['POST'])
|
65 |
+
def generate_api():
|
66 |
+
data = request.get_json()
|
67 |
+
|
68 |
+
# Extract required fields from the request
|
69 |
+
prompt = data.get('prompt', '')
|
70 |
+
negative_prompt = data.get('negative_prompt', None)
|
71 |
+
height = data.get('height', 1024) # Default height
|
72 |
+
width = data.get('width', 720) # Default width
|
73 |
+
num_inference_steps = data.get('num_inference_steps', 50) # Default number of inference steps
|
74 |
+
guidance_scale = data.get('guidance_scale', 7.5) # Default guidance scale
|
75 |
+
model_name = data.get('model', 'stabilityai/sd-3.5') # Base model
|
76 |
+
refiner_model_name = 'stabilityai/sd-xl-refiner-1.0' # Refiner model
|
77 |
+
seed = data.get('seed', None) # Seed for reproducibility, default is None
|
78 |
+
|
79 |
+
if not prompt:
|
80 |
+
return jsonify({"error": "Prompt is required"}), 400
|
81 |
+
|
82 |
+
try:
|
83 |
+
# Check for explicit content
|
84 |
+
if is_prompt_explicit(prompt):
|
85 |
+
# Return the pre-defined "thinkgood.png" image
|
86 |
+
return send_file(
|
87 |
+
"thinkgood.jpeg",
|
88 |
+
mimetype='image/png',
|
89 |
+
as_attachment=False,
|
90 |
+
download_name='thinkgood.png'
|
91 |
+
)
|
92 |
+
|
93 |
+
# Step 1: Generate the base image
|
94 |
+
base_image = generate_image(prompt, negative_prompt, height, width, model_name, num_inference_steps, guidance_scale, seed)
|
95 |
+
|
96 |
+
if not base_image:
|
97 |
+
return jsonify({"error": "Failed to generate base image"}), 500
|
98 |
+
|
99 |
+
# Step 2: Refine the image with the refiner model
|
100 |
+
refined_image = refine_image(base_image, prompt, negative_prompt, refiner_model_name, num_inference_steps, guidance_scale)
|
101 |
+
|
102 |
+
if not refined_image:
|
103 |
+
return jsonify({"error": "Failed to refine image"}), 500
|
104 |
+
|
105 |
+
# Save the refined image to a BytesIO object
|
106 |
+
img_byte_arr = BytesIO()
|
107 |
+
refined_image.save(img_byte_arr, format='PNG') # Convert the image to PNG
|
108 |
+
img_byte_arr.seek(0) # Move to the start of the byte stream
|
109 |
+
|
110 |
+
# Send the refined image as a response
|
111 |
+
return send_file(
|
112 |
+
img_byte_arr,
|
113 |
+
mimetype='image/png',
|
114 |
+
as_attachment=False, # Send the file inline
|
115 |
+
download_name='refined_image.png' # File name for download
|
116 |
+
)
|
117 |
+
except Exception as e:
|
118 |
+
print(f"Error in generate_api: {str(e)}") # Log the error
|
119 |
+
return jsonify({"error": str(e)}), 500
|
120 |
+
|
121 |
+
# Add this block to make sure your app runs when called
|
122 |
+
if __name__ == "__main__":
|
123 |
+
app.run(host='0.0.0.0', port=7860) # Run directly if needed for testing
|