src/hopter/client.py

import asyncio
import httpx
from enum import Enum
from src.utils import image_path_to_uri
from dotenv import load_dotenv
import os
from pydantic import BaseModel, Field
from typing import List

load_dotenv()


class Environment(Enum):
    STAGING = "staging"
    PRODUCTION = "production"

    @property
    def base_url(self) -> str:
        match self:
            case Environment.STAGING:
                return "https://serving.hopter.staging.picc.co"
            case Environment.PRODUCTION:
                return "https://serving.hopter.picc.co"


class RamGroundedSamInput(BaseModel):
    text_prompt: str = Field(
        ..., description="The text prompt for the mask generation."
    )
    image_b64: str = Field(..., description="The image in base64 format.")


class RamGroundedSamResult(BaseModel):
    mask_b64: str = Field(..., description="The mask image in base64 format.")
    class_label: str = Field(..., description="The class label of the mask.")
    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 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, api_key: str, environment: Environment = Environment.PRODUCTION):
        self.api_key = api_key
        self.base_url = environment.base_url
        self.client = httpx.Client()

    def generate_mask(self, input: RamGroundedSamInput) -> RamGroundedSamResult:
        print(f"Generating mask with input: {input.text_prompt}")
        try:
            response = self.client.post(
                f"{self.base_url}/api/v1/services/ram-grounded-sam-api/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").get("instances")[0]
            print("Generated mask.")
            return RamGroundedSamResult(**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}")

    def magic_replace(self, input: MagicReplaceInput) -> MagicReplaceResult:
        print(f"Magic replacing with input: {input.prompt}")
        try:
            response = 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()},
                timeout=None,
            )
            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}")

    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", image_b64=image_url)
    mask = 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 = hopter.magic_replace(input)
    return result.base64_image


async def main():
    hopter = Hopter(
        api_key=os.getenv("HOPTER_API_KEY"), environment=Environment.STAGING
    )
    image_file_path = "./assets/lakeview.jpg"
    image_url = image_path_to_uri(image_file_path)

    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())