feat: integrated super resolution

src/agents/mask_generation_agent.py

@@ -6,7 +6,7 @@ import asyncio
 from dataclasses import dataclass
 import logfire
 from src.services.generate_mask import GenerateMaskService
-from src.hopter.client import Hopter, Environment, MagicReplaceInput
+from src.hopter.client import Hopter, Environment, MagicReplaceInput, SuperResolutionInput
 from src.services.image_uploader import ImageUploader
 from src.utils import image_path_to_uri
 
@@ -67,25 +67,36 @@ async def edit_object(ctx: RunContext[ImageEditDeps]) -> EditImageResult:
     """
     Edit an object in the image.
     """
-    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 = mask_service.get_mask_generation_instruction(edit_instruction, image_url)
     mask = mask_service.generate_mask(mask_instruction, image_url)
 
     # Magic replace
     input = MagicReplaceInput(image=image_url, mask=mask, prompt=mask_instruction.target_caption)
     result = hopter_client.magic_replace(input)
-    print("Exiting edit_object tool: ", result)
     uploader = ImageUploader(os.environ.get("IMG_BB_API_KEY"))
     uploaded_image = uploader.upload_url(result.base64_image)
     return EditImageResult(edited_image_url=uploaded_image.data.url)
 
+@mask_generation_agent.tool
+async def super_resolution(ctx: RunContext[ImageEditDeps]) -> EditImageResult:
+    """
+    run super resolution, upscale, or enhance the image
+    """
+    image_url = ctx.deps.image_url
+    hopter_client = ctx.deps.hopter_client
+
+    input = SuperResolutionInput(image_b64=image_url, scale=4, use_face_enhancement=False)
+    result = hopter_client.super_resolution(input)
+    uploader = ImageUploader(os.environ.get("IMG_BB_API_KEY"))
+    uploaded_image = uploader.upload_url(result.scaled_image)
+    return EditImageResult(edited_image_url=uploaded_image.data.url)
+
 async def main():
     image_file_path = "./assets/lakeview.jpg"
     image_url = image_path_to_uri(image_file_path)

src/hopter/client.py

@@ -39,6 +39,14 @@ class MagicReplaceInput(BaseModel):
 class MagicReplaceResult(BaseModel):
     base64_image: str = Field(..., description="The edited image in base64 format.")
 
+class SuperResolutionInput(BaseModel):
+    image_b64: str = Field(..., description="The image in base64 format.")
+    scale: int = Field(4, description="The scale of the image to upscale to.")
+    use_face_enhancement: bool = Field(False, description="Whether to use face enhancement.")
+
+class SuperResolutionResult(BaseModel):
+    scaled_image: str = Field(..., description="The super-resolved image in base64 format.")
+
 class Hopter:
     def __init__(
         self, 
@@ -95,6 +103,29 @@ class Hopter:
         except Exception as exc:
             print(f"An unexpected error occurred: {exc}")
 
+    def super_resolution(self, input: SuperResolutionInput) -> SuperResolutionResult:
+        try:
+            response = self.client.post(
+                f"{self.base_url}/api/v1/services/super-resolution-esrgan/predictions",
+                headers={
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json"
+                },
+                json={
+                    "input": input.model_dump()
+                },
+                timeout=None
+            )
+            response.raise_for_status()  # Raise an error for bad responses
+            instance = response.json().get("output")
+            print("Super-resolutin done")
+            return SuperResolutionResult(**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",