Update app.py

app.py

@@ -1,138 +1,32 @@
-from flask import Flask, request, jsonify, send_file
+from flask import Flask, request, send_file
 from flask_cors import CORS
-import os
 from huggingface_hub import InferenceClient
-from io import BytesIO
 from PIL import Image
-import base64
+import io
 
-# Initialize the Flask app
 app = Flask(__name__)
 CORS(app)  # Enable CORS for all routes
+client = InferenceClient()
 
-# 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)
+@app.route('/generate-image', methods=['POST'])
+def generate_image():
+    # Get the image file from the request
+    file = request.files['image']
+    prompt = request.form.get('prompt', '')
 
-@app.route('/')
-def home():
-    return "Welcome to the Image Generation Service!"
+    # Read the image directly from the file stream
+    image = Image.open(file.stream)
 
-# Function to generate an image from a text prompt
-def generate_text_to_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:
-        # Generate the image using Hugging Face's inference API with additional parameters
-        image = client.text_to_image(
-            prompt=prompt, 
-            negative_prompt=negative_prompt, 
-            height=height, 
-            width=width, 
-            model=model,
-            num_inference_steps=num_inference_steps,  # Control the number of inference steps
-            guidance_scale=guidance_scale,  # Control the guidance scale
-            seed=seed  # Control the seed for reproducibility
-        )
-        return image  # Return the generated image
-    except Exception as e:
-        print(f"Error generating image: {str(e)}")
-        return None
+    # Generate image using the InferenceClient
+    generated_image = client.image_to_image(image, prompt=prompt)
 
-# Function to modify an existing image using image-to-image processing
-def generate_image_to_image(input_image, prompt, height, width, num_inference_steps, guidance_scale, seed=None):
-    try:
-        # Hard-coded model for image-to-image processing
-        hardcoded_model = "stabilityai/sdxl-1.0"
-        
-        # Convert the base64-encoded input image to a PIL image
-        decoded_image = Image.open(BytesIO(base64.b64decode(input_image)))
-        
-        # Modify the image using image-to-image transformation
-        modified_image = client.image_to_image(
-            image=decoded_image,
-            prompt=prompt,
-            model=hardcoded_model,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            seed=seed
-        )
-        return modified_image
-    except Exception as e:
-        print(f"Error generating image-to-image: {str(e)}")
-        return None
+    # Save the generated image to a BytesIO object
+    img_byte_arr = io.BytesIO()
+    generated_image.save(img_byte_arr, format='JPEG')
+    img_byte_arr.seek(0)  # Move the cursor to the beginning of the BytesIO object
 
-# Flask route for the API endpoint to generate an image
-@app.route('/generate_image', methods=['POST'])
-def generate_api():
-    data = request.get_json()
+    # Send the generated image back to the client
+    return send_file(img_byte_arr, mimetype='image/jpeg')
 
-    # Extract common fields
-    prompt = data.get('prompt', '')
-    negative_prompt = data.get('negative_prompt', None)
-    num_inference_steps = data.get('num_inference_steps', 50)
-    guidance_scale = data.get('guidance_scale', 7.5)
-    model_name = data.get('model', 'stabilityai/stable-diffusion-2-1')  # Model specified by the user for text-to-image
-    seed = data.get('seed', None)
-
-    if not prompt:
-        return jsonify({"error": "Prompt is required"}), 400
-
-    # Check if the request contains an image for image-to-image processing
-    input_image = data.get('image', None)  # Expecting a base64-encoded image string
-    if input_image:
-        # Extract parameters specific to image-to-image
-        height = data.get('height', 1024)
-        width = data.get('width', 720)
-
-        # Call the image-to-image function
-        final_image = generate_image_to_image(
-            input_image, prompt, height, width, num_inference_steps, guidance_scale, seed
-        )
-
-        if final_image:
-            # Save the modified image to a BytesIO object
-            img_byte_arr = BytesIO()
-            final_image.save(img_byte_arr, format='PNG')
-            img_byte_arr.seek(0)
-
-            # Send the modified image as a response
-            return send_file(
-                img_byte_arr, 
-                mimetype='image/png', 
-                as_attachment=False, 
-                download_name='final_image.png'
-            )
-        else:
-            return jsonify({"error": "Failed to generate image-to-image."}), 500
-
-    # If no image is provided, proceed with text-to-image generation
-    else:
-        # Default height and width for text-to-image
-        height = data.get('height', 512)
-        width = data.get('width', 512)
-
-        # Call the text-to-image function
-        generated_image = generate_text_to_image(
-            prompt, negative_prompt, height, width, model_name, num_inference_steps, guidance_scale, seed
-        )
-
-        if generated_image:
-            # Save the generated image to a BytesIO object
-            img_byte_arr = BytesIO()
-            generated_image.save(img_byte_arr, format='PNG')
-            img_byte_arr.seek(0)
-
-            # Send the generated image as a response
-            return send_file(
-                img_byte_arr, 
-                mimetype='image/png', 
-                as_attachment=False, 
-                download_name='generated_image.png'
-            )
-        else:
-            return jsonify({"error": "Failed to generate text-to-image."}), 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
+    app.run(debug=True)