Geek7 commited on
Commit
fed8daa
·
verified ·
1 Parent(s): 3b57a34

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +123 -0
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