Update myapp.py

myapp.py

@@ -1,9 +1,9 @@
-from flask import Flask, request, jsonify, send_file
+from flask import Flask, request, jsonify, send_file, render_template
 from flask_cors import CORS
 import os
 from huggingface_hub import InferenceClient
 from io import BytesIO
-from PIL import Image  # Importing Pillow for image processing
+from PIL import Image
 
 # Initialize the Flask app
 myapp = Flask(__name__)
@@ -13,12 +13,8 @@ CORS(myapp)  # 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)
 
-@myapp.route('/')
-def home():
-    return "Welcome to the Image Background Remover!"
-    
-# Function to generate an image from a prompt using the specified model
-def generate_image(prompt, seed=1, model="prompthero/openjourney-v4"):  # Default model if none provided
+# Function to generate an image from a text prompt
+def generate_image(prompt, seed=1, model="prompthero/openjourney-v4"):
     try:
         # Generate the image using Hugging Face's inference API
         result_image = client.text_to_image(prompt=prompt, seed=seed, model=model)
@@ -27,7 +23,7 @@ def generate_image(prompt, seed=1, model="prompthero/openjourney-v4"):  # Defaul
         print(f"Error generating image: {str(e)}")
         return None
 
-# Flask route for the API endpoint
+# Flask route for the API endpoint to generate an image based on a text prompt
 @myapp.route('/generate_image', methods=['POST'])
 def generate_api():
     data = request.get_json()
@@ -35,24 +31,24 @@ def generate_api():
     # Extract required fields from the request
     prompt = data.get('prompt', '')
     seed = data.get('seed', 1)
-    model_name = data.get('model', 'prompthero/openjourney-v4')  # Use the provided model name or a default one
+    model_name = data.get('model', 'prompthero/openjourney-v4')  # Default model
 
     if not prompt:
         return jsonify({"error": "Prompt is required"}), 400
 
     try:
-        # Call the generate_image function with the dynamically provided model name
+        # Call the generate_image function with the custom model name
         image = generate_image(prompt, seed, model_name)
 
         if image:
             # Save the image to a BytesIO object
-            image_bytes = BytesIO()
-            image.save(image_bytes, format='PNG')
-            image_bytes.seek(0)  # Go to the start of the byte stream
+            img_byte_arr = BytesIO()
+            image.save(img_byte_arr, format='PNG')  # Convert the image to PNG
+            img_byte_arr.seek(0)  # Move to the start of the byte stream
 
             # Send the generated image as an attachment
             return send_file(
-                image_bytes, 
+                img_byte_arr, 
                 mimetype='image/png', 
                 as_attachment=True,  # Send the file as an attachment
                 download_name='generated_image.png'  # The file name for download
@@ -63,6 +59,11 @@ def generate_api():
         print(f"Error in generate_api: {str(e)}")  # Log the error
         return jsonify({"error": str(e)}), 500
 
+# Flask route for the UI
+@myapp.route('/')
+def index():
+    return render_template('index.html')  # Serve the HTML page
+
 # Add this block to make sure your app runs when called
 if __name__ == "__main__":
     myapp.run(host='0.0.0.0', port=7860)  # Run directly if needed for testing