Update myapp.py

myapp.py

@@ -1,9 +1,10 @@
-from flask import Flask, request, jsonify, send_file, render_template
+from flask import Flask, request, jsonify
 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
 
 # Initialize the Flask app
 myapp = Flask(__name__)
@@ -23,46 +24,30 @@ def generate_image(prompt, seed=1, model="prompthero/openjourney-v4"):
         print(f"Error generating image: {str(e)}")
         return None
 
-# 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()
-
-    # Extract required fields from the request
-    prompt = data.get('prompt', '')
-    seed = data.get('seed', 1)
-    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 custom model name
-        image = generate_image(prompt, seed, 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 an attachment
-            return send_file(
-                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
-            )
-        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
-
-# Flask route for the UI
-@myapp.route('/')
-def index():
-    return render_template('index.html')  # Serve the HTML page
+# Gradio interface function
+def gradio_interface(prompt, seed, model_name):
+    image = generate_image(prompt, seed, model_name)
+
+    if image:
+        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
+        return img_byte_arr  # Return the image as bytes
+    else:
+        return "Failed to generate image"
+
+# Set up the Gradio interface
+gr.Interface(
+    fn=gradio_interface,
+    inputs=[
+        gr.Textbox(label="Prompt", placeholder="Enter a text prompt", lines=2),
+        gr.Number(label="Seed", value=1, precision=0),
+        gr.Textbox(label="Model Name", value="prompthero/openjourney-v4", placeholder="Enter model name"),
+    ],
+    outputs="image",
+    title="Image Generation with Hugging Face",
+    description="Enter a prompt, seed, and model name to generate an image."
+).launch()  # Launch the Gradio interface
 
 # Add this block to make sure your app runs when called
 if __name__ == "__main__":