Update app.py

app.py

@@ -1,39 +1,40 @@
-from flask import Flask, request, send_file
-from flask_cors import CORS
+import gradio as gr
 from huggingface_hub import InferenceClient
 from PIL import Image
 import io
 import base64
 
-app = Flask(__name__)
-CORS(app)  # Enable CORS for all routes
+# Initialize the Inference Client
 client = InferenceClient()
 
-@app.route('/')
-def home():
-    return "Welcome to the Image Background Remover!"
-
-@app.route('/generate-image', methods=['POST'])
-def generate_image():
-    data = request.json  # Get the JSON data from the request
-    base64_image = data['image']  # Get the base64 image string
-    prompt = data['prompt']  # Get the prompt
-
+# Function to handle the image generation
+def generate_image(base64_image, prompt):
     # Decode the base64 image
-    image_data = base64.b64decode(base64_image)
+    image_data = base64.b64decode(base64_image.split(",")[1])
     image = Image.open(io.BytesIO(image_data))
 
     # Generate image using the InferenceClient
     generated_image = client.image_to_image(image, prompt=prompt)
 
-    # Save the generated image to a BytesIO object
+    # Save generated image to a BytesIO object
     img_byte_arr = io.BytesIO()
     generated_image.save(img_byte_arr, format='PNG')
-    img_byte_arr.seek(0)  # Move the cursor to the beginning of the BytesIO object
-
-    # Send the generated image back to the client
-    return send_file(img_byte_arr, mimetype='image/png')
-
-
-if __name__ == "__main__":
-    app.run(host='0.0.0.0', port=7860)  # Run directly if needed for testing
+    img_byte_arr.seek(0)  # Move the cursor to the beginning
+
+    # Return the generated image
+    return img_byte_arr
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate_image,
+    inputs=[
+        gr.inputs.Image(type="filepath", label="Input Image"),  # Filepath for uploaded image
+        gr.inputs.Textbox(label="Prompt")  # Textbox for the prompt
+    ],
+    outputs=gr.outputs.Image(type="numpy", label="Generated Image"),  # Output as numpy array
+    title="Image Generation with Hugging Face",
+    description="Upload an image and provide a prompt to generate a new image."
+)
+
+# Launch the Gradio interface
+iface.launch()