Update app.py

app.py

@@ -4,6 +4,7 @@ import os
 from huggingface_hub import InferenceClient
 from io import BytesIO
 from PIL import Image
+import threading
 
 # Initialize the Flask app
 app = Flask(__name__)
@@ -13,14 +14,20 @@ CORS(app)  # Enable CORS for all routes
 HF_TOKEN = os.environ.get("HF_TOKEN")  # Ensure to set your Hugging Face token in the environment
 client = InferenceClient(token=HF_TOKEN)
 
+# Global variable to manage generation rejection state
+generation_rejected = threading.Event()
+
 @app.route('/')
 def home():
     return "Welcome to the Image Background Remover!"
 
-
 # Function to generate an image from a text prompt
 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):
     try:
+        # Check if generation has been marked for rejection
+        if generation_rejected.is_set():
+            return "rejected"
+
         # Generate the image using Hugging Face's inference API with additional parameters
         image = client.text_to_image(
             prompt=prompt, 
@@ -40,6 +47,9 @@ def generate_image(prompt, negative_prompt=None, height=512, width=512, model="s
 # Flask route for the API endpoint to generate an image
 @app.route('/generate_image', methods=['POST'])
 def generate_api():
+    global generation_rejected
+    generation_rejected.clear()  # Reset rejection state at the start of a new request
+
     data = request.get_json()
 
     # Extract required fields from the request
@@ -59,6 +69,10 @@ def generate_api():
         # Call the generate_image function with the provided parameters
         image = generate_image(prompt, negative_prompt, height, width, model_name, num_inference_steps, guidance_scale, seed)
 
+        # Check if the response was rejected
+        if image == "rejected":
+            return jsonify({"error": "Image generation was rejected."}), 400
+
         if image:
             # Save the image to a BytesIO object
             img_byte_arr = BytesIO()
@@ -78,6 +92,13 @@ def generate_api():
         print(f"Error in generate_api: {str(e)}")  # Log the error
         return jsonify({"error": str(e)}), 500
 
+# Add a new endpoint to reject the image generation
+@app.route('/reject', methods=['POST'])
+def reject_generation():
+    global generation_rejected
+    generation_rejected.set()  # Set the flag to indicate the request should be rejected
+    return jsonify({"message": "Image generation has been marked for rejection."})
+
 # Add this block to make sure your app runs when called
 if __name__ == "__main__":
     app.run(host='0.0.0.0', port=7860)  # Run directly if needed for testing