curretn progress

mcp_server.py

@@ -1,11 +1,9 @@
 import gradio as gr
 from src.utils.change_format import change_format
-from src.utils.remove_background import remove_background_from_url
-from src.utils.visualize_image import visualize_base64_image
+from src.utils.remove_background import remove_background
 from src.utils.generate_image import generate_image
 from src.utils.apply_filter import apply_filter
 from src.utils.add_text import add_text_to_image
-from src.utils.resize_image import resize_image
 from src.utils.watermark import add_watermark, remove_watermark
 from src.utils.describe import describe_image
 from src.utils.compress import compress_image
@@ -24,9 +22,55 @@ def image_to_base64(image):
 def base64_to_image(base64_str):
     if not base64_str:
         return None
-    image_data = base64.b64decode(base64_str)
-    return Image.open(io.BytesIO(image_data))
-
+    
+    # Remove data URI prefix if present (e.g., "data:image/png;base64,")
+    if isinstance(base64_str, str) and "base64," in base64_str:
+        base64_str = base64_str.split("base64,", 1)[1]
+    
+    try:
+        # Strip any whitespace that might be in the base64 string
+        if isinstance(base64_str, str):
+            base64_str = base64_str.strip()
+        
+        # Decode the base64 data
+        image_data = base64.b64decode(base64_str)
+        
+        # Check if we have data
+        if not image_data:
+            print("Decoded base64 data is empty")
+            return None
+        
+        # Attempt to open the image
+        image = Image.open(io.BytesIO(image_data))
+        
+        # Convert the image to ensure it's valid
+        return image.copy()
+        
+    except base64.binascii.Error as e:
+        print(f"Base64 decoding error: {str(e)}")
+        if isinstance(base64_str, str):
+            preview = base64_str[:30] + "..." if len(base64_str) > 30 else base64_str
+            print(f"Base64 preview: {preview}")
+        return None
+        
+    except Exception as e:
+        print(f"Error converting base64 to image: {str(e)}")
+        
+        # Print preview of the base64 string for debugging
+        if isinstance(base64_str, str):
+            preview = base64_str[:30] + "..." if len(base64_str) > 30 else base64_str
+            print(f"Base64 preview: {preview}")
+        
+        # Additional debug information
+        if 'image_data' in locals() and image_data:
+            try:
+                magic_bytes = image_data[:12].hex()
+                print(f"First 12 bytes: {magic_bytes}")
+            except:
+                pass
+        
+        return None
+    
 def url_to_base64(url):
     response = requests.get(url)
     return base64.b64encode(response.content).decode()
@@ -35,8 +79,26 @@ def gradio_remove_background(image):
     if image is None:
         return None
     base64_img = image_to_base64(image)
-    result = remove_background_from_url(f"data:image/png;base64,{base64_img}")
-    return base64_to_image(result)
+    result = remove_background(f"data:image/png;base64,{base64_img}")
+    
+    # Check if the result is directly a base64 string or has an image_data key
+    if isinstance(result, str):
+        return base64_to_image(result)
+    elif isinstance(result, dict) and "image_data" in result:
+        # If image_data contains a data URI prefix
+        if isinstance(result["image_data"], str) and result["image_data"].startswith("data:"):
+            # The response already contains the full data URI
+            return base64_to_image(result["image_data"])
+        else:
+            # Try to process it as a regular base64 string
+            try:
+                return base64_to_image(result["image_data"])
+            except Exception as e:
+                print(f"Error processing image data: {e}")
+                return None
+    else:
+        print(f"Unexpected response format from remove_background: {type(result)}")
+        return None
 
 def gradio_describe_image(image):
     if image is None:
@@ -53,7 +115,7 @@ def gradio_change_format(image, format_type):
 
 def gradio_generate_image(prompt, width=512, height=512):
     result = generate_image(prompt, width, height)
-    return base64_to_image(result)
+    return base64_to_image(result["b64"])
 
 def gradio_apply_filter(image, filter_type):
     if image is None:

src/utils/describe.py

@@ -4,6 +4,8 @@ import requests
 from pathlib import Path
 from openai import OpenAI
 from urllib.parse import urlparse
+from dotenv import load_dotenv
+
 
 def describe_image(image_path: str) -> str:
     """
@@ -14,6 +16,7 @@ def describe_image(image_path: str) -> str:
     
     Returns:
         A string description of the image """
+    load_dotenv()
     
     # Check if API key is available
     api_key = os.getenv("NEBIUS_API_KEY")

src/utils/generate_image.py

@@ -3,9 +3,8 @@ import base64
 from typing import Dict, Any
 from openai import OpenAI
 
-async def generate_image(
+def generate_image(
     prompt: str,
-    output_path: str = "generated_image.png",
     width: int = 1024,
     height: int = 1024,
     num_inference_steps: int = 28,
@@ -50,17 +49,12 @@ async def generate_image(
         
         image_data = base64.b64decode(response.data[0].b64_json)
         
-        with open(output_path, 'wb') as output_file:
-            output_file.write(image_data)
-        
-        output_size = os.path.getsize(output_path)
         
         return {
             "success": True,
             "message": "Image generated successfully",
             "prompt": prompt,
-            "output_path": output_path,
-            "output_size_bytes": output_size,
+            "b64": image_data,
             "generation_params": {
                 "width": width,
                 "height": height,

src/utils/remove_background.py

@@ -1,46 +1,88 @@
 import requests
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, Union
 import os
 import rembg
+import numpy as np
+from PIL import Image
+import io
+import base64
+import re
 
-async def remove_background_from_url(
-    image_url: str,
-    output_path: str,
+def remove_background(
+    image_input: Union[str, bytes, np.ndarray, Image.Image],
     model_name: str = "u2net"
 ) -> Dict[str, Any]:
     """
-    Remove background from an image downloaded from a URL.
+    Remove background from an image.
     
     Args:
-        image_url: URL of the image to process
-        output_path: Path where to save the processed image
+        image_input: Can be one of:
+            - URL string
+            - Data URL string (base64 encoded)
+            - Image bytes
+            - NumPy array
+            - PIL Image
         model_name: Background removal model to use
     
     Returns:
-        Dictionary with result information
+        Dictionary with result information and processed image data
     """
     
     try:
-        # Download image from URL
-        response = requests.get(image_url, timeout=30)
-        response.raise_for_status()
-        
-        # Remove background
+        # Initialize session
         session = rembg.new_session(model_name=model_name)
-        output_data = rembg.remove(response.content, session=session)
         
-        # Save the result
-        with open(output_path, 'wb') as output_file:
-            output_file.write(output_data)
+        # Handle different input types
+        if isinstance(image_input, str):
+            if image_input.startswith('http://') or image_input.startswith('https://'):
+                # If input is a URL, download the image
+                response = requests.get(image_input, timeout=30)
+                response.raise_for_status()
+                input_data = response.content
+                source_info = f"URL: {image_input}"
+            elif image_input.startswith('data:'):
+                # If input is a data URL (base64 encoded string)
+                # Extract the base64 part after the comma
+                base64_data = re.sub('^data:image/.+;base64,', '', image_input)
+                input_data = base64.b64decode(base64_data)
+                source_info = "data URL"
+            else:
+                return {
+                    "success": False,
+                    "error": f"Unsupported string input format: {image_input[:30]}...",
+                    "image_data": None
+                }
+        elif isinstance(image_input, bytes):
+            # If input is bytes, use directly
+            input_data = image_input
+            source_info = "image bytes"
+        elif isinstance(image_input, np.ndarray):
+            # If input is numpy array, convert to bytes
+            pil_img = Image.fromarray(image_input)
+            buffer = io.BytesIO()
+            pil_img.save(buffer, format="PNG")
+            input_data = buffer.getvalue()
+            source_info = "numpy array"
+        elif isinstance(image_input, Image.Image):
+            # If input is PIL Image, convert to bytes
+            buffer = io.BytesIO()
+            image_input.save(buffer, format="PNG")
+            input_data = buffer.getvalue()
+            source_info = "PIL Image"
+        else:
+            return {
+                "success": False,
+                "error": f"Unsupported input type: {type(image_input)}",
+                "image_data": None
+            }
         
-        output_size = os.path.getsize(output_path)
+        # Remove background
+        output_data = rembg.remove(input_data, session=session)
         
         return {
             "success": True,
-            "message": f"Background removed from URL image using {model_name} model",
-            "source_url": image_url,
-            "output_path": output_path,
-            "output_size_bytes": output_size,
+            "message": f"Background removed from {source_info} using {model_name} model",
+            "image_data": output_data,
             "model_used": model_name
         }
         
@@ -48,11 +90,11 @@ async def remove_background_from_url(
         return {
             "success": False,
             "error": f"Failed to download image: {str(e)}",
-            "output_path": None
+            "image_data": None
         }
     except Exception as e:
         return {
             "success": False,
             "error": f"Failed to process image: {str(e)}",
-            "output_path": None
+            "image_data": None
         }