Update app.py

app.py

@@ -1,14 +1,20 @@
+from flask import Flask, request, jsonify, send_file
+from flask_cors import CORS
 import os
 from huggingface_hub import InferenceClient
+from io import BytesIO
 from PIL import Image
-import gradio as gr  # Import Gradio for the UI
+
+# Initialize the Flask app
+app = Flask(__name__)
+CORS(app)  # Enable CORS for all routes
 
 # Initialize the InferenceClient with your Hugging Face token
 HF_TOKEN = os.environ.get("HF_TOKEN")  # Ensure to set your Hugging Face token in the environment
-client = InferenceClient(token=HF_TOKEN)  # Initialize the client
+client = InferenceClient(token=HF_TOKEN)
 
 # Function to generate an image from a text prompt
-def generate_image(prompt, negative_prompt=None, model=None):
+def generate_image(prompt, negative_prompt=None, model="stabilityai/stable-diffusion-2-1"):
     try:
         # Generate the image using Hugging Face's inference API
         image = client.text_to_image(prompt=prompt, negative_prompt=negative_prompt, model=model)
@@ -17,15 +23,42 @@ def generate_image(prompt, negative_prompt=None, model=None):
         print(f"Error generating image: {str(e)}")
         return None
 
-# Set up the Gradio interface
-gr.Interface(
-    fn=generate_image,
-    inputs=[
-        gr.Textbox(label="Prompt", placeholder="Enter a text prompt", lines=2),
-        gr.Textbox(label="Negative Prompt (Optional)", placeholder="Enter negative prompt", lines=2),
-        gr.Textbox(label="Model Name", placeholder="Enter model name", value="stabilityai/stable-diffusion-2-1"),
-    ],
-    outputs="image",
-    title="Image Generation with Hugging Face",
-    description="Enter a prompt, optional negative prompt, and model name to generate an image.",
-).launch()  # Launch the Gradio interface
+# Flask route for the API endpoint to generate an image
+@app.route('/generate_image', methods=['POST'])
+def generate_api():
+    data = request.get_json()
+
+    # Extract required fields from the request
+    prompt = data.get('prompt', '')
+    negative_prompt = data.get('negative_prompt', None)
+    model_name = data.get('model', 'stabilityai/stable-diffusion-2-1')  # Default model
+
+    if not prompt:
+        return jsonify({"error": "Prompt is required"}), 400
+
+    try:
+        # Call the generate_image function with the provided parameters
+        image = generate_image(prompt, negative_prompt, model_name)
+
+        if image:
+            # Save the image to a BytesIO object
+            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 a response
+            return send_file(
+                img_byte_arr, 
+                mimetype='image/png', 
+                as_attachment=False,  # Send the file as an attachment
+                download_name='generated_image.png'  # The file name for download
+            )
+        else:
+            return jsonify({"error": "Failed to generate image"}), 500
+    except Exception as e:
+        print(f"Error in generate_api: {str(e)}")  # Log the error
+        return jsonify({"error": str(e)}), 500
+
+# 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