Update app.py

app.py

@@ -4,7 +4,6 @@ import os
 from huggingface_hub import InferenceClient
 from io import BytesIO
 from PIL import Image
-import threading
 
 # Initialize the Flask app
 app = Flask(__name__)
@@ -14,20 +13,14 @@ 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, 
@@ -47,9 +40,6 @@ 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
@@ -69,10 +59,6 @@ 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()
@@ -92,13 +78,6 @@ 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