feat: working magic replace

agent.py

@@ -17,6 +17,18 @@ from pydantic_ai.messages import (
 )
 import asyncio
 from src.agents.mask_generation_agent import mask_generation_agent, ImageEditDeps
+from src.hopter.client import Hopter, Environment
+import os
+from src.services.generate_mask import GenerateMaskService
+from dotenv import load_dotenv
+
+load_dotenv()
+
+hopter = Hopter(
+    api_key=os.getenv("HOPTER_API_KEY"),
+    environment=Environment.STAGING
+)
+mask_service = GenerateMaskService(hopter=hopter)
 
 class ChatMessage(TypedDict):
     """Format of messages sent to the browser/API."""
@@ -60,7 +72,9 @@ async def run_agent(user_input: str, image_b64: str):
     ]
     deps = ImageEditDeps(
         edit_instruction=user_input,
-        image_url=image_b64
+        image_url=image_b64,
+        hopter_client=hopter,
+        mask_service=mask_service
     )
     async with mask_generation_agent.run_stream(
         messages,

src/agents/mask_generation_agent.py

@@ -1,5 +1,4 @@
 from pydantic_ai import Agent, RunContext
-from pydantic_ai.settings import ModelSettings
 from pydantic_ai.models.openai import OpenAIModel
 from dotenv import load_dotenv
 import os
@@ -8,6 +7,7 @@ import base64
 from dataclasses import dataclass
 import logfire
 from src.services.generate_mask import GenerateMaskService
+from src.hopter.client import Hopter, Environment, MagicReplaceInput
 
 load_dotenv()
 
@@ -23,6 +23,8 @@ if the edit instruction involved modifying parts of the image, please generate a
 class ImageEditDeps:
     edit_instruction: str
     image_url: str
+    hopter_client: Hopter
+    mask_service: GenerateMaskService
 
 model = OpenAIModel(
     "gpt-4o",
@@ -33,26 +35,53 @@ model = OpenAIModel(
 class MaskGenerationResult:
     mask_image_base64: str
 
+
+@dataclass
+class EditImageResult:
+    edited_image_base64: str
+
 mask_generation_agent = Agent(
     model,
     system_prompt=system_prompt,
     deps_type=ImageEditDeps
 )
 
+# @mask_generation_agent.tool
+# async def generate_mask(ctx: RunContext[ImageEditDeps]) -> MaskGenerationResult:
+#     """
+#     Generate a mask for the image editing instruction.
+#     """
+#     print("Invoking generate_mask tool")
+#     service = GenerateMaskService()
+#     mask_instruction = await service.get_mask_generation_instruction(ctx.deps.edit_instruction, ctx.deps.image_url)
+#     response = mask_instruction.model_dump_json(indent=4)
+#     print(f"generate_mask tool response: {response}")
+
+#     mask = await service.generate_mask(mask_instruction, ctx.deps.image_url)
+#     print("Exiting generate_mask tool")
+#     return MaskGenerationResult(mask_image_base64=mask)
+
 @mask_generation_agent.tool
-async def generate_mask(ctx: RunContext[ImageEditDeps]) -> MaskGenerationResult:
+async def edit_object(ctx: RunContext[ImageEditDeps]) -> EditImageResult:
     """
-    Generate a mask for the image editing instruction. 
+    Edit an object in the image.
     """
-    print("Invoking generate_mask tool")
-    service = GenerateMaskService()
-    mask_instruction = await service.get_mask_generation_instruction(ctx.deps.edit_instruction, ctx.deps.image_url)
-    response = mask_instruction.model_dump_json(indent=4)
-    print(f"generate_mask tool response: {response}")
+    print("Invoking edit_object tool")
+    edit_instruction = ctx.deps.edit_instruction
+    image_url = ctx.deps.image_url
+    mask_service = ctx.deps.mask_service
+    hopter_client = ctx.deps.hopter_client
+
+    # Generate mask
+    print("Generating mask")
+    mask_instruction = await mask_service.get_mask_generation_instruction(edit_instruction, image_url)
+    mask = await mask_service.generate_mask(mask_instruction, image_url)
 
-    mask = await service.generate_mask(mask_instruction, ctx.deps.image_url)
-    print("Exiting generate_mask tool")
-    return MaskGenerationResult(mask_image_base64=mask)
+    # Magic replace
+    input = MagicReplaceInput(image=image_url, mask=mask, prompt=mask_instruction.target_caption)
+    result = await hopter_client.magic_replace(input)
+    print("Exiting edit_object tool: ", result)
+    return EditImageResult(edited_image_base64=result.base64_image)
 
 async def main():
     image_file_path = "./assets/lakeview.jpg"
@@ -75,12 +104,23 @@ async def main():
             }
         ]
 
+        # Initialize services
+        hopter = Hopter(api_key=os.environ.get("HOPTER_API_KEY"), environment=Environment.STAGING)
+        mask_service = GenerateMaskService(hopter=hopter)
+
+        # Initialize dependencies
         deps = ImageEditDeps(
             edit_instruction=prompt,
-            image_url=image_url
+            image_url=image_url,
+            hopter_client=hopter,
+            mask_service=mask_service
         )
-        r = await mask_generation_agent.run(messages, deps=deps)
-        print(r.all_messages())
+        async with mask_generation_agent.run_stream(
+            messages, 
+            deps=deps
+        ) as result:
+            async for message in result.stream():
+                print(message)
 
 
 if __name__ == "__main__":

src/hopter/client.py

@@ -32,6 +32,14 @@ class RamGroundedSamResult(BaseModel):
     confidence: float = Field(..., description="The confidence score of the mask.")
     bbox: List[float] = Field(..., description="The bounding box of the mask in the format [x1, y1, x2, y2].")
 
+class MagicReplaceInput(BaseModel):
+    image: str = Field(..., description="The image in base64 format.")
+    mask: str = Field(..., description="The mask in base64 format.")
+    prompt: str = Field(..., description="The prompt for the magic replace.")
+
+class MagicReplaceResult(BaseModel):
+    base64_image: str = Field(..., description="The edited image in base64 format.")
+
 class Hopter:
     def __init__(
         self, 
@@ -64,7 +72,45 @@ class Hopter:
         except Exception as exc:
             print(f"An unexpected error occurred: {exc}")
     
-    # async def _ram_grounded_sam(self, prompt: str, image_base64: str) -> str:
+    async def magic_replace(self, input: MagicReplaceInput) -> MagicReplaceResult:
+        print(f"Magic replacing with input: {input.prompt}")
+        try:
+            response = await self.client.post(
+                f"{self.base_url}/api/v1/services/sdxl-magic-replace/predictions",
+                headers={
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json"
+                },
+                json={
+                    "input": input.model_dump()
+                }
+            )
+            print(response)
+            response.raise_for_status()  # Raise an error for bad responses
+            instance = response.json().get("output")
+            print("Magic replaced.")
+            return MagicReplaceResult(**instance)
+        except httpx.HTTPStatusError as exc:
+            print(f"HTTP error occurred: {exc.response.status_code} - {exc.response.text}")
+        except Exception as exc:
+            print(f"An unexpected error occurred: {exc}")
+
+async def test_generate_mask(hopter: Hopter, image_url: str) -> str:
+    input = RamGroundedSamInput(
+        text_prompt="pole",
+        image_b64=image_url
+    )
+    mask = await hopter.generate_mask(input)
+    return mask.mask_b64
+    
+async def test_magic_replace(hopter: Hopter, image_url: str, mask: str, prompt: str) -> str:
+    input = MagicReplaceInput(
+        image=image_url,
+        mask=mask,
+        prompt=prompt
+    )
+    result = await hopter.magic_replace(input)
+    return result.base64_image
 
 async def main():
     hopter = Hopter(
@@ -77,12 +123,8 @@ async def main():
         image_base64 = base64.b64encode(image_bytes).decode("utf-8")
         image_url = f"data:image/jpeg;base64,{image_base64}"
 
-        input = RamGroundedSamInput(
-            text_prompt="pole",
-            image_b64=image_url
-        )
-        mask = await hopter.generate_mask(input)
-        print(mask)
-
+        mask = await test_generate_mask(hopter, image_url)
+        magic_replace_result = await test_magic_replace(hopter, image_url, mask, "remove the pole")
+        print(magic_replace_result)
 if __name__ == "__main__":
     asyncio.run(main())

src/services/generate_mask.py

@@ -6,6 +6,8 @@ import base64
 import asyncio
 from src.hopter.client import Hopter, RamGroundedSamInput, Environment
 from src.models.generate_mask_instruction import GenerateMaskInstruction
+from src.services.openai_file_upload import OpenAIFileUpload
+
 load_dotenv()
 
 system_prompt = """
@@ -38,13 +40,13 @@ Do not output 'sorry, xxx', even if it's a guess, directly output the answer you
 """
 
 class GenerateMaskService:
-    def __init__(self):
+    def __init__(self, hopter: Hopter):
         self.llm = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
-        self.hopter = Hopter(api_key=os.environ.get("HOPTER_API_KEY"), environment=Environment.STAGING)
         self.model = "gpt-4o"
+        self.openai_file_upload = OpenAIFileUpload()
+        self.hopter = hopter
 
     async def get_mask_generation_instruction(self, edit_instruction: str, image_url: str) -> GenerateMaskInstruction:
-        
         messages = [
             {
                 "role": "system",
@@ -93,13 +95,11 @@ class GenerateMaskService:
         return generate_mask_result.mask_b64
     
 async def main():
-    service = GenerateMaskService()
+    service = GenerateMaskService(Hopter(api_key=os.environ.get("HOPTER_API_KEY"), environment=Environment.STAGING))
     edit_instruction = "remove the light post"
     image_file_path = "./assets/lakeview.jpg"
     with open(image_file_path, "rb") as image_file:
-        image_bytes = image_file.read()
-        image_base64 = base64.b64encode(image_bytes).decode("utf-8")
-        image_url = f"data:image/jpeg;base64,{image_base64}"
+        image_url = service.openai_file_upload.upload_image(image_file.read(), "vision")
 
     instruction = await service.get_mask_generation_instruction(edit_instruction, image_url)
     print(instruction)

src/services/image_uploader.py

@@ -0,0 +1,161 @@
+import requests
+from typing import Optional, Union
+import base64
+from pathlib import Path
+import os
+from pydantic import BaseModel
+
+class ImageInfo(BaseModel):
+    filename: str
+    name: str
+    mime: str
+    extension: str
+    url: str
+
+class ImgBBData(BaseModel):
+    id: str
+    title: str
+    url_viewer: str
+    url: str
+    display_url: str
+    width: int
+    height: int
+    size: int
+    time: int
+    expiration: int
+    image: ImageInfo
+    thumb: ImageInfo
+    medium: ImageInfo
+    delete_url: str
+
+class ImgBBResponse(BaseModel):
+    data: ImgBBData
+    success: bool
+    status: int
+
+class ImageUploader:
+    """A class to handle image uploads to ImgBB service."""
+    
+    def __init__(self, api_key: str):
+        """
+        Initialize the ImageUploader with an API key.
+        
+        Args:
+            api_key (str): The ImgBB API key
+        """
+        self.api_key = api_key
+        self.base_url = "https://api.imgbb.com/1/upload"
+        
+    def upload(
+        self,
+        image: Union[str, bytes, Path],
+        name: Optional[str] = None,
+        expiration: Optional[int] = None
+    ) -> ImgBBResponse:
+        """
+        Upload an image to ImgBB.
+        
+        Args:
+            image: Can be:
+                - A file path (str or Path)
+                - Base64 encoded string
+                - Base64 data URI (e.g., data:image/jpeg;base64,...)
+                - URL to an image
+                - Bytes of an image
+            name: Optional name for the uploaded file
+            expiration: Optional expiration time in seconds (60-15552000)
+            
+        Returns:
+            ImgBBResponse containing the parsed upload response from ImgBB
+            
+        Raises:
+            ValueError: If the image format is invalid or upload fails
+            requests.RequestException: If the API request fails
+        """
+        # Prepare the parameters
+        params = {'key': self.api_key}
+        if expiration:
+            if not 60 <= expiration <= 15552000:
+                raise ValueError("Expiration must be between 60 and 15552000 seconds")
+            params['expiration'] = expiration
+
+        # Handle different image input types
+        if isinstance(image, (str, Path)):
+            image_str = str(image)
+            files = {}
+            if os.path.isfile(image_str):
+                # It's a file path
+                with open(image_str, 'rb') as file:
+                    files['image'] = file
+            elif image_str.startswith(('http://', 'https://')):
+                # It's a URL
+                files['image'] = (None, image_str)
+            elif image_str.startswith('data:image/'):
+                # It's a data URI
+                # Extract the base64 part after the comma
+                base64_data = image_str.split(',', 1)[1]
+                files['image'] = (None, base64_data)
+            else:
+                # Assume it's base64 data
+                files['image'] = (None, image_str)
+
+            if name:
+                files['name'] = (None, name)
+            response = requests.post(self.base_url, params=params, files=files)
+        elif isinstance(image, bytes):
+            # Convert bytes to base64
+            base64_image = base64.b64encode(image).decode('utf-8')
+            files = {
+                'image': (None, base64_image)
+            }
+            if name:
+                files['name'] = (None, name)
+            response = requests.post(self.base_url, params=params, files=files)
+        else:
+            raise ValueError("Invalid image format. Must be file path, URL, base64 string, or bytes")
+
+        # Check the response
+        if response.status_code != 200:
+            raise ValueError(f"Upload failed with status {response.status_code}: {response.text}")
+
+        # Parse the response using Pydantic model
+        response_json = response.json()
+        return ImgBBResponse.parse_obj(response_json)
+
+    def upload_file(
+        self,
+        file_path: Union[str, Path],
+        name: Optional[str] = None,
+        expiration: Optional[int] = None
+    ) -> ImgBBResponse:
+        """
+        Convenience method to upload an image file.
+        
+        Args:
+            file_path: Path to the image file
+            name: Optional name for the uploaded file
+            expiration: Optional expiration time in seconds (60-15552000)
+            
+        Returns:
+            ImgBBResponse containing the parsed upload response from ImgBB
+        """
+        return self.upload(file_path, name=name, expiration=expiration)
+
+    def upload_url(
+        self,
+        image_url: str,
+        name: Optional[str] = None,
+        expiration: Optional[int] = None
+    ) -> ImgBBResponse:
+        """
+        Convenience method to upload an image from a URL.
+        
+        Args:
+            image_url: URL of the image to upload
+            name: Optional name for the uploaded file
+            expiration: Optional expiration time in seconds (60-15552000)
+            
+        Returns:
+            ImgBBResponse containing the parsed upload response from ImgBB
+        """
+        return self.upload(image_url, name=name, expiration=expiration)

src/services/openai_file_upload.py

@@ -0,0 +1,13 @@
+from openai import OpenAI
+from dotenv import load_dotenv
+import os
+
+load_dotenv()
+
+class OpenAIFileUpload:
+    def __init__(self):
+        self.client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+
+    def upload_image(self, image, purpose) -> str:
+        file = self.client.files.create(file=image, purpose=purpose)
+        return file.id