Update app.py

app.py

@@ -4,6 +4,7 @@ import os
 from huggingface_hub import InferenceClient
 from io import BytesIO
 from PIL import Image
+import base64
 
 # Initialize the Flask app
 app = Flask(__name__)
@@ -15,11 +16,10 @@ client = InferenceClient(token=HF_TOKEN)
 
 @app.route('/')
 def home():
-    return "Welcome to the Image Background Remover!"
-
+    return "Welcome to the Image Generation Service!"
 
 # Function to generate an image from a text prompt
-def generate_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):
+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(
@@ -37,46 +37,101 @@ def generate_image(prompt, negative_prompt=None, height=512, width=512, model="s
         print(f"Error generating image: {str(e)}")
         return None
 
+# 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
+
 # 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
+    # Extract common fields
     prompt = data.get('prompt', '')
     negative_prompt = data.get('negative_prompt', None)
-    height = data.get('height', 1024)  # Default height
-    width = data.get('width', 720)  # Default width
-    num_inference_steps = data.get('num_inference_steps', 50)  # Default number of inference steps
-    guidance_scale = data.get('guidance_scale', 7.5)  # Default guidance scale
-    model_name = data.get('model', 'stabilityai/stable-diffusion-2-1')  # Default model
-    seed = data.get('seed', None)  # Seed for reproducibility, default is 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
 
-    try:
-        # Call the generate_image function with the provided parameters
-        image = generate_image(prompt, negative_prompt, height, width, model_name, num_inference_steps, guidance_scale, seed)
+    # 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 image:
-            # Save the image to a BytesIO object
+        if final_image:
+            # Save the modified 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
+            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,  # Send the file as an attachment
-                download_name='generated_image.png'  # The file name for download
+                as_attachment=False, 
+                download_name='generated_image.png'
             )
         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
+            return jsonify({"error": "Failed to generate text-to-image."}), 500
 
 # Add this block to make sure your app runs when called
 if __name__ == "__main__":