Update app.py
Browse files
app.py
CHANGED
@@ -1,123 +0,0 @@
|
|
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=15, 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
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|