isa-model 0.3.9__py3-none-any.whl → 0.4.0__py3-none-any.whl

isa_model/inference/services/img/services/replicate_sticker_maker.py

@@ -0,0 +1,249 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+"""
+Replicate Sticker Maker Service
+Specialized service for generating stickers using the fofr/sticker-maker model
+"""
+
+import os
+import logging
+from typing import Dict, Any, Optional
+import replicate
+
+from ..base_image_gen_service import BaseImageGenService
+
+logger = logging.getLogger(__name__)
+
+class ReplicateStickerMakerService(BaseImageGenService):
+    """
+    Replicate Sticker Maker Service - $0.0024 per generation
+    Specialized for creating cute stickers from text prompts
+    """
+    
+    def __init__(self, provider_name: str, model_name: str, **kwargs):
+        super().__init__(provider_name, model_name, **kwargs)
+        
+        # Get configuration from centralized config manager
+        provider_config = self.get_provider_config()
+        
+        try:
+            self.api_token = provider_config.get("api_key") or provider_config.get("replicate_api_token")
+            
+            if not self.api_token:
+                raise ValueError("Replicate API token not found in provider configuration")
+            
+            # Set API token
+            os.environ["REPLICATE_API_TOKEN"] = self.api_token
+            
+            # Model path
+            self.model_path = "fofr/sticker-maker:4acb778eb059772225ec213948f0660867b2e03f277448f18cf1800b96a65a1a"
+            
+            # Statistics
+            self.total_generation_count = 0
+            
+            logger.info(f"Initialized ReplicateStickerMakerService with model '{self.model_name}'")
+            
+        except Exception as e:
+            logger.error(f"Failed to initialize Replicate Sticker Maker client: {e}")
+            raise ValueError(f"Failed to initialize Replicate Sticker Maker client: {e}") from e
+
+    async def generate_sticker(
+        self, 
+        prompt: str,
+        steps: int = 17,
+        width: int = 1152,
+        height: int = 1152,
+        output_format: str = "webp",
+        output_quality: int = 100,
+        negative_prompt: str = "",
+        number_of_images: int = 1
+    ) -> Dict[str, Any]:
+        """Generate sticker from text prompt"""
+        
+        input_data = {
+            "steps": steps,
+            "width": width,
+            "height": height,
+            "prompt": prompt,
+            "output_format": output_format,
+            "output_quality": output_quality,
+            "negative_prompt": negative_prompt,
+            "number_of_images": number_of_images
+        }
+        
+        return await self._generate_internal(input_data)
+
+    async def _generate_internal(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+        """Internal generation method"""
+        try:
+            logger.info(f"Starting sticker generation with prompt: {input_data.get('prompt', '')[:50]}...")
+            
+            # Call Replicate API
+            output = await replicate.async_run(self.model_path, input=input_data)
+            
+            # Process output - convert FileOutput objects to URL strings
+            if isinstance(output, list):
+                raw_urls = output
+            else:
+                raw_urls = [output]
+            
+            # Convert to string URLs
+            urls = []
+            for url in raw_urls:
+                if hasattr(url, 'url'):
+                    urls.append(str(url.url))
+                else:
+                    urls.append(str(url))
+
+            # Update statistics
+            self.total_generation_count += len(urls)
+            
+            # Calculate cost
+            cost = self._calculate_cost(len(urls))
+            
+            # Track billing information
+            await self._track_usage(
+                service_type="image_generation",
+                operation="sticker_generation",
+                input_tokens=0,
+                output_tokens=0,
+                input_units=1,  # Input prompt
+                output_units=len(urls),  # Generated stickers count
+                metadata={
+                    "model": self.model_name,
+                    "prompt": input_data.get("prompt", "")[:100],
+                    "generation_type": "sticker",
+                    "image_count": len(urls),
+                    "cost_usd": cost
+                }
+            )
+            
+            # Return URLs
+            result = {
+                "urls": urls,
+                "url": urls[0] if urls else None,
+                "format": input_data.get("output_format", "webp"),
+                "width": input_data.get("width", 1152),
+                "height": input_data.get("height", 1152),
+                "count": len(urls),
+                "cost_usd": cost,
+                "metadata": {
+                    "model": self.model_name,
+                    "input": input_data,
+                    "generation_count": len(urls)
+                }
+            }
+            
+            logger.info(f"Sticker generation completed: {len(urls)} stickers, cost: ${cost:.6f}")
+            return result
+            
+        except Exception as e:
+            logger.error(f"Sticker generation failed: {e}")
+            raise
+
+    def _calculate_cost(self, image_count: int) -> float:
+        """Calculate generation cost - $0.0024 per generation"""
+        return image_count * 0.0024
+
+    def get_generation_stats(self) -> Dict[str, Any]:
+        """Get generation statistics"""
+        total_cost = self.total_generation_count * 0.0024
+        
+        return {
+            "total_generation_count": self.total_generation_count,
+            "total_cost_usd": total_cost,
+            "cost_per_generation": 0.0024,
+            "model": self.model_name
+        }
+
+    def get_model_info(self) -> Dict[str, Any]:
+        """Get model information"""
+        return {
+            "name": self.model_name,
+            "type": "sticker_generation",
+            "cost_per_generation": 0.0024,
+            "supports_negative_prompt": True,
+            "max_width": 1152,
+            "max_height": 1152,
+            "output_formats": ["webp", "jpg", "png"]
+        }
+
+    async def load(self) -> None:
+        """Load service"""
+        if not self.api_token:
+            raise ValueError("Missing Replicate API token")
+        logger.info(f"Replicate Sticker Maker service ready with model: {self.model_name}")
+
+    async def unload(self) -> None:
+        """Unload service"""
+        logger.info(f"Unloading Replicate Sticker Maker service: {self.model_name}")
+
+    async def close(self):
+        """Close service"""
+        await self.unload()
+
+    async def generate_image(
+        self, 
+        prompt: str,
+        negative_prompt: Optional[str] = None,
+        width: int = 512,
+        height: int = 512,
+        num_inference_steps: int = 17,
+        guidance_scale: float = 7.5,
+        seed: Optional[int] = None
+    ) -> Dict[str, Any]:
+        """Generate single sticker"""
+        return await self.generate_sticker(
+            prompt=prompt,
+            steps=num_inference_steps,
+            width=width,
+            height=height,
+            negative_prompt=negative_prompt or ""
+        )
+
+    async def generate_images(
+        self, 
+        prompt: str,
+        num_images: int = 1,
+        negative_prompt: Optional[str] = None,
+        width: int = 512,
+        height: int = 512,
+        num_inference_steps: int = 17,
+        guidance_scale: float = 7.5,
+        seed: Optional[int] = None
+    ) -> list[Dict[str, Any]]:
+        """Generate multiple stickers"""
+        results = []
+        for i in range(num_images):
+            result = await self.generate_sticker(
+                prompt=prompt,
+                steps=num_inference_steps,
+                width=width,
+                height=height,
+                negative_prompt=negative_prompt or "",
+                number_of_images=1
+            )
+            results.append(result)
+        return results
+
+    async def image_to_image(
+        self,
+        prompt: str,
+        init_image,
+        strength: float = 0.8,
+        negative_prompt: Optional[str] = None,
+        num_inference_steps: int = 17,
+        guidance_scale: float = 7.5,
+        seed: Optional[int] = None
+    ) -> Dict[str, Any]:
+        """Not supported by sticker maker"""
+        raise NotImplementedError("Sticker maker does not support image-to-image generation")
+
+    def get_supported_sizes(self) -> list[Dict[str, int]]:
+        """Get supported image sizes"""
+        return [
+            {"width": 1152, "height": 1152},
+            {"width": 1024, "height": 1024},
+            {"width": 768, "height": 768},
+        ]

isa_model/inference/services/img/tests/test_img_client.py

@@ -0,0 +1,297 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+"""
+Test Image Generation Client using ISA Model Client
+Tests the four specialized image generation services through the unified client:
+1. FLUX Schnell (text-to-image)
+2. FLUX Kontext Pro (image-to-image)  
+3. Sticker Maker
+4. Face Swap
+"""
+
+import asyncio
+import logging
+from typing import Dict, Any
+
+from isa_model.client import ISAModelClient
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class ImageGenerationTester:
+    """Test client for image generation services using ISA Model Client"""
+    
+    def __init__(self):
+        self.client = ISAModelClient()
+        
+        # Test configurations for each service
+        self.test_configs = {
+            "flux_schnell": {
+                "model": "flux-schnell",
+                "provider": "replicate",
+                "task": "generate",
+                "prompt": "a cute cat sitting in a garden"
+            },
+            "flux_kontext": {
+                "model": "flux-kontext-pro", 
+                "provider": "replicate",
+                "task": "img2img",
+                "prompt": "transform this into a futuristic cityscape",
+                "init_image": "https://replicate.delivery/pbxt/Mb44XIUHkUrmyyH1OP5K1WmFN7SNN0eUSU16A8rBtuXe7eYV/cyberpunk_80s_example.png"
+            },
+            "sticker_maker": {
+                "model": "sticker-maker",
+                "provider": "replicate", 
+                "task": "generate",
+                "prompt": "a cute cat"
+            },
+            "face_swap": {
+                "model": "face-swap",
+                "provider": "replicate",
+                "task": "face_swap",
+                "swap_image": "https://replicate.delivery/pbxt/Mb44Wp0W7Xfa1Pp91zcxDzSSQQz8GusUmXQXi3GGzRxDvoCI/0_1.webp",
+                "target_image": "https://replicate.delivery/pbxt/Mb44XIUHkUrmyyH1OP5K1WmFN7SNN0eUSU16A8rBtuXe7eYV/cyberpunk_80s_example.png"
+            }
+        }
+    
+    async def test_flux_text_to_image(self) -> Dict[str, Any]:
+        """Test FLUX Schnell text-to-image generation"""
+        logger.info("Testing FLUX Schnell text-to-image...")
+        
+        try:
+            config = self.test_configs["flux_schnell"]
+            
+            result = await self.client.invoke(
+                input_data=config["prompt"],
+                task=config["task"],
+                service_type="image",
+                model=config["model"],
+                provider=config["provider"],
+                width=1024,
+                height=1024,
+                num_inference_steps=4
+            )
+            
+            if result.get("success"):
+                response = result["result"]
+                logger.info(f"FLUX generation successful: {response.get('count', 0)} images")
+                logger.info(f"Cost: ${response.get('cost_usd', 0):.6f}")
+                logger.info(f"URL: {response.get('url', 'N/A')}")
+                
+                return {
+                    "status": "success",
+                    "result": response,
+                    "metadata": result.get("metadata", {})
+                }
+            else:
+                error_msg = result.get("error", "Unknown error")
+                logger.error(f"FLUX generation failed: {error_msg}")
+                return {"status": "error", "error": error_msg}
+                
+        except Exception as e:
+            logger.error(f"FLUX generation failed with exception: {e}")
+            return {"status": "error", "error": str(e)}
+    
+    async def test_flux_kontext_image_to_image(self) -> Dict[str, Any]:
+        """Test FLUX Kontext Pro image-to-image generation"""
+        logger.info("Testing FLUX Kontext Pro image-to-image...")
+        
+        try:
+            config = self.test_configs["flux_kontext"]
+            
+            result = await self.client.invoke(
+                input_data=config["prompt"],
+                task=config["task"],
+                service_type="image",
+                model=config["model"],
+                provider=config["provider"],
+                init_image=config["init_image"],
+                strength=0.8
+            )
+            
+            if result.get("success"):
+                response = result["result"]
+                logger.info(f"FLUX Kontext generation successful: {response.get('count', 0)} images")
+                logger.info(f"Cost: ${response.get('cost_usd', 0):.6f}")
+                logger.info(f"URL: {response.get('url', 'N/A')}")
+                
+                return {
+                    "status": "success",
+                    "result": response,
+                    "metadata": result.get("metadata", {})
+                }
+            else:
+                error_msg = result.get("error", "Unknown error")
+                logger.error(f"FLUX Kontext generation failed: {error_msg}")
+                return {"status": "error", "error": error_msg}
+                
+        except Exception as e:
+            logger.error(f"FLUX Kontext generation failed with exception: {e}")
+            return {"status": "error", "error": str(e)}
+    
+    async def test_sticker_generation(self) -> Dict[str, Any]:
+        """Test sticker generation"""
+        logger.info("Testing Sticker Maker...")
+        
+        try:
+            config = self.test_configs["sticker_maker"]
+            
+            result = await self.client.invoke(
+                input_data=config["prompt"],
+                task=config["task"],
+                service_type="image",
+                model=config["model"],
+                provider=config["provider"],
+                steps=17,
+                width=1152,
+                height=1152,
+                output_format="webp",
+                output_quality=100
+            )
+            
+            if result.get("success"):
+                response = result["result"]
+                logger.info(f"Sticker generation successful: {response.get('count', 0)} stickers")
+                logger.info(f"Cost: ${response.get('cost_usd', 0):.6f}")
+                logger.info(f"URL: {response.get('url', 'N/A')}")
+                
+                return {
+                    "status": "success",
+                    "result": response,
+                    "metadata": result.get("metadata", {})
+                }
+            else:
+                error_msg = result.get("error", "Unknown error")
+                logger.error(f"Sticker generation failed: {error_msg}")
+                return {"status": "error", "error": error_msg}
+                
+        except Exception as e:
+            logger.error(f"Sticker generation failed with exception: {e}")
+            return {"status": "error", "error": str(e)}
+    
+    async def test_face_swap(self) -> Dict[str, Any]:
+        """Test face swap"""
+        logger.info("Testing Face Swap...")
+        
+        try:
+            config = self.test_configs["face_swap"]
+            
+            result = await self.client.invoke(
+                input_data=config["swap_image"],  # Use swap_image as input_data
+                task=config["task"],
+                service_type="image",
+                model=config["model"],
+                provider=config["provider"],
+                target_image=config["target_image"],
+                hair_source="target",
+                user_gender="default",
+                user_b_gender="default"
+            )
+            
+            if result.get("success"):
+                response = result["result"]
+                logger.info(f"Face swap successful: {response.get('count', 0)} images")
+                logger.info(f"Cost: ${response.get('cost_usd', 0):.6f}")
+                logger.info(f"URL: {response.get('url', 'N/A')}")
+                
+                return {
+                    "status": "success",
+                    "result": response,
+                    "metadata": result.get("metadata", {})
+                }
+            else:
+                error_msg = result.get("error", "Unknown error")
+                logger.error(f"Face swap failed: {error_msg}")
+                return {"status": "error", "error": error_msg}
+                
+        except Exception as e:
+            logger.error(f"Face swap failed with exception: {e}")
+            return {"status": "error", "error": str(e)}
+    
+    async def test_all_services(self) -> Dict[str, Dict[str, Any]]:
+        """Test all image generation services"""
+        logger.info("Starting comprehensive image generation tests using ISA Model Client...")
+        
+        results = {}
+        
+        # Test each service
+        tests = [
+            ("flux_text_to_image", self.test_flux_text_to_image),
+            ("flux_image_to_image", self.test_flux_kontext_image_to_image),
+            ("sticker_generation", self.test_sticker_generation),
+            ("face_swap", self.test_face_swap)
+        ]
+        
+        for test_name, test_func in tests:
+            logger.info(f"\n{'='*50}")
+            logger.info(f"Running test: {test_name}")
+            logger.info(f"{'='*50}")
+            
+            try:
+                result = await test_func()
+                results[test_name] = result
+                
+                if result.get("status") == "success":
+                    logger.info(f" {test_name} PASSED")
+                else:
+                    logger.error(f"L {test_name} FAILED: {result.get('error', 'Unknown error')}")
+                    
+            except Exception as e:
+                logger.error(f"L {test_name} FAILED with exception: {e}")
+                results[test_name] = {"status": "error", "error": str(e)}
+        
+        # Summary
+        logger.info(f"\n{'='*50}")
+        logger.info("TEST SUMMARY")
+        logger.info(f"{'='*50}")
+        
+        passed = sum(1 for r in results.values() if r.get("status") == "success")
+        total = len(results)
+        
+        logger.info(f"Passed: {passed}/{total}")
+        
+        for test_name, result in results.items():
+            status = " PASS" if result.get("status") == "success" else "L FAIL"
+            logger.info(f"{test_name}: {status}")
+        
+        return results
+    
+    async def get_service_health(self) -> Dict[str, Any]:
+        """Get health status of the client and services"""
+        logger.info("Checking service health...")
+        
+        try:
+            health = await self.client.health_check()
+            return health
+        except Exception as e:
+            logger.error(f"Health check failed: {e}")
+            return {"status": "error", "error": str(e)}
+
+async def main():
+    """Main test function"""
+    tester = ImageGenerationTester()
+    
+    # Get service health
+    logger.info("Checking service health...")
+    health = await tester.get_service_health()
+    logger.info(f"Service health: {health}")
+    
+    # Run all tests
+    results = await tester.test_all_services()
+    
+    # Calculate total cost
+    total_cost = 0.0
+    for test_name, result in results.items():
+        if result.get("status") == "success":
+            cost = result.get("result", {}).get("cost_usd", 0.0)
+            total_cost += cost
+    
+    logger.info(f"\nTotal cost for all tests: ${total_cost:.6f}")
+    
+    return results
+
+if __name__ == "__main__":
+    # Run the tests
+    results = asyncio.run(main())

isa_model/inference/services/llm/base_llm_service.py

@@ -28,6 +28,7 @@ class BaseLLMService(BaseService):
         self, 
         input_data: Union[str, List[Dict[str, str]], Any],
         task: Optional[str] = None,
+        show_reasoning: bool = False,
         **kwargs
     ) -> Dict[str, Any]:
         """
@@ -48,7 +49,7 @@ class BaseLLMService(BaseService):
         
         # ==================== 对话类任务 ====================
         if task == "chat":
-            return await self.chat(input_data, kwargs.get("max_tokens", self.max_tokens))
+            return await self.chat(input_data, kwargs.get("max_tokens", self.max_tokens), show_reasoning=show_reasoning)
         elif task == "complete":
             return await self.complete_text(input_data, kwargs.get("max_tokens", self.max_tokens))
         elif task == "instruct":
@@ -62,7 +63,10 @@ class BaseLLMService(BaseService):
         elif task == "summarize":
             return await self.summarize_text(input_data, kwargs.get("max_length"), kwargs.get("style"))
         elif task == "translate":
-            return await self.translate_text(input_data, kwargs.get("target_language"), kwargs.get("source_language"))
+            target_language = kwargs.get("target_language")
+            if not target_language:
+                raise ValueError("target_language is required for translate task")
+            return await self.translate_text(input_data, target_language, kwargs.get("source_language"))
         
         # ==================== 分析类任务 ====================
         elif task == "analyze":
@@ -91,12 +95,17 @@ class BaseLLMService(BaseService):
             return await self.solve_problem(input_data, kwargs.get("problem_type"))
         elif task == "plan":
             return await self.create_plan(input_data, kwargs.get("plan_type"))
+        elif task == "deep_research":
+            return await self.deep_research(input_data, kwargs.get("research_type"), kwargs.get("search_enabled", True))
         
         # ==================== 工具调用类任务 ====================
         elif task == "tool_call":
             return await self.call_tools(input_data, kwargs.get("available_tools"))
         elif task == "function_call":
-            return await self.call_function(input_data, kwargs.get("function_name"), kwargs.get("parameters"))
+            function_name = kwargs.get("function_name")
+            if not function_name:
+                raise ValueError("function_name is required for function_call task")
+            return await self.call_function(input_data, function_name, kwargs.get("parameters"))
         
         else:
             raise NotImplementedError(f"{self.__class__.__name__} does not support task: {task}")
@@ -106,7 +115,8 @@ class BaseLLMService(BaseService):
     async def chat(
         self,
         input_data: Union[str, List[Dict[str, str]], Any],
-        max_tokens: Optional[int] = None
+        max_tokens: Optional[int] = None,
+        show_reasoning: bool = False
     ) -> Dict[str, Any]:
         """
         对话聊天 - Provider必须实现
@@ -114,6 +124,7 @@ class BaseLLMService(BaseService):
         Args:
             input_data: 输入消息
             max_tokens: 最大生成token数
+            show_reasoning: 是否显示推理过程
             
         Returns:
             Dict containing chat response
@@ -303,6 +314,17 @@ class BaseLLMService(BaseService):
         """
         raise NotImplementedError(f"{self.__class__.__name__} does not support create_plan task")
     
+    async def deep_research(
+        self,
+        input_data: Union[str, Any],
+        research_type: Optional[str] = None,
+        search_enabled: bool = True
+    ) -> Dict[str, Any]:
+        """
+        深度研究 - O-series模型专用任务，支持网络搜索和深入分析
+        """
+        raise NotImplementedError(f"{self.__class__.__name__} does not support deep_research task")
+    
     # ==================== 工具调用类方法 ====================
     
     async def call_tools(
@@ -354,7 +376,7 @@ class BaseLLMService(BaseService):
         """使用适配器管理器转换消息格式"""
         return self.adapter_manager.convert_messages(input_data)
     
-    def _format_response(self, response: str, original_input: Any) -> Union[str, Any]:
+    def _format_response(self, response: Union[str, Any], original_input: Any) -> Union[str, Any]:
         """使用适配器管理器格式化响应"""
         return self.adapter_manager.format_response(response, original_input)
     
@@ -379,7 +401,7 @@ class BaseLLMService(BaseService):
         pass
     
     @abstractmethod
-    async def ainvoke(self, input_data: Union[str, List[Dict[str, str]], Any]) -> Union[str, Any]:
+    async def ainvoke(self, input_data: Union[str, List[Dict[str, str]], Any], show_reasoning: bool = False) -> Union[str, Any]:
         """
         Universal async invocation method that handles different input types
         
@@ -388,6 +410,7 @@ class BaseLLMService(BaseService):
                 - str: Simple text prompt
                 - list: Message history like [{"role": "user", "content": "hello"}]
                 - Any: LangChain message objects or other formats
+            show_reasoning: If True and model supports it, show reasoning process
             
         Returns:
             Model response (string for simple cases, object for complex cases)
@@ -527,6 +550,7 @@ class BaseLLMService(BaseService):
             'reason_about': 'reason',
             'solve_problem': 'solve',
             'create_plan': 'plan',
+            'deep_research': 'deep_research',
             # 工具调用类
             'call_tools': 'tool_call',
             'call_function': 'function_call'