isa-model 0.3.4__py3-none-any.whl → 0.3.6__py3-none-any.whl

isa_model/inference/services/{vision → img}/base_image_gen_service.py

@@ -3,7 +3,86 @@ from typing import Dict, Any, List, Union, Optional, BinaryIO
 from isa_model.inference.services.base_service import BaseService
 
 class BaseImageGenService(BaseService):
-    """Base class for image generation services"""
+    """Base class for image generation services with unified task dispatch"""
+    
+    async def invoke(
+        self, 
+        prompt: str,
+        task: Optional[str] = None,
+        **kwargs
+    ) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
+        """
+        统一的任务分发方法 - Base类提供通用实现
+        
+        Args:
+            prompt: 文本提示词
+            task: 任务类型，支持多种图像生成任务
+            **kwargs: 任务特定的附加参数
+            
+        Returns:
+            Dict or List[Dict] containing generation results
+        """
+        task = task or "generate"
+        
+        # ==================== 图像生成类任务 ====================
+        if task == "generate":
+            num_images = kwargs.get("num_images", 1)
+            if num_images == 1:
+                return await self.generate_image(
+                    prompt,
+                    kwargs.get("negative_prompt"),
+                    kwargs.get("width", 512),
+                    kwargs.get("height", 512),
+                    kwargs.get("num_inference_steps", 20),
+                    kwargs.get("guidance_scale", 7.5),
+                    kwargs.get("seed")
+                )
+            else:
+                return await self.generate_images(
+                    prompt,
+                    num_images,
+                    kwargs.get("negative_prompt"),
+                    kwargs.get("width", 512),
+                    kwargs.get("height", 512),
+                    kwargs.get("num_inference_steps", 20),
+                    kwargs.get("guidance_scale", 7.5),
+                    kwargs.get("seed")
+                )
+        elif task == "generate_batch":
+            return await self.generate_images(
+                prompt,
+                kwargs.get("num_images", 4),
+                kwargs.get("negative_prompt"),
+                kwargs.get("width", 512),
+                kwargs.get("height", 512),
+                kwargs.get("num_inference_steps", 20),
+                kwargs.get("guidance_scale", 7.5),
+                kwargs.get("seed")
+            )
+        elif task == "img2img":
+            init_image = kwargs.get("init_image")
+            if not init_image:
+                raise ValueError("img2img task requires init_image parameter")
+            return await self.image_to_image(
+                prompt,
+                init_image,
+                kwargs.get("strength", 0.8),
+                kwargs.get("negative_prompt"),
+                kwargs.get("num_inference_steps", 20),
+                kwargs.get("guidance_scale", 7.5),
+                kwargs.get("seed")
+            )
+        else:
+            raise NotImplementedError(f"{self.__class__.__name__} does not support task: {task}")
+    
+    def get_supported_tasks(self) -> List[str]:
+        """
+        获取支持的任务列表
+        
+        Returns:
+            List of supported task names
+        """
+        return ["generate", "generate_batch", "img2img"]
     
     @abstractmethod
     async def generate_image(
@@ -68,40 +147,6 @@ class BaseImageGenService(BaseService):
         """
         pass
     
-    @abstractmethod
-    async def generate_image_to_file(
-        self, 
-        prompt: str,
-        output_path: str,
-        negative_prompt: Optional[str] = None,
-        width: int = 512,
-        height: int = 512,
-        num_inference_steps: int = 20,
-        guidance_scale: float = 7.5,
-        seed: Optional[int] = None
-    ) -> Dict[str, Any]:
-        """
-        Generate image and save directly to file
-        
-        Args:
-            prompt: Text description of the desired image
-            output_path: Path to save the generated image
-            negative_prompt: Text describing what to avoid in the image
-            width: Image width in pixels
-            height: Image height in pixels
-            num_inference_steps: Number of denoising steps
-            guidance_scale: How closely to follow the prompt
-            seed: Random seed for reproducible results
-            
-        Returns:
-            Dict containing generation results with keys:
-            - file_path: Path to saved image
-            - width: Image width
-            - height: Image height
-            - seed: Seed used for generation
-        """
-        pass
-    
     @abstractmethod
     async def image_to_image(
         self,