isa-model 0.0.3__py3-none-any.whl → 0.0.8__py3-none-any.whl

isa_model/inference/services/vision/replicate_image_gen_service.py

@@ -17,8 +17,8 @@ import replicate # 导入 replicate 库
 from PIL import Image
 from io import BytesIO
 
-# 调整 BaseService 的导入路径以匹配您的项目结构
-from isa_model.inference.services.base_service import BaseService
+# 调整导入路径以使用正确的基类
+from isa_model.inference.services.vision.base_image_gen_service import BaseImageGenService
 from isa_model.inference.providers.base_provider import BaseProvider
 from isa_model.inference.base import ModelType
 
@@ -26,7 +26,7 @@ from isa_model.inference.base import ModelType
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-class ReplicateVisionService(BaseService):
+class ReplicateImageGenService(BaseImageGenService):
     """
     Replicate Vision服务，用于处理图像生成和分析。
     经过调整，使用原生异步调用并优化了文件处理。
@@ -76,9 +76,9 @@ class ReplicateVisionService(BaseService):
                     raise
         return prepared_input, files_to_close
 
-    async def generate_image(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+    async def _generate_image_internal(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
         """
-        使用Replicate模型生成图像 (已优化为原生异步)
+        内部方法：使用Replicate模型生成图像 (已优化为原生异步)
         """
         prepared_input, files_to_close = await self._prepare_input_files(input_data)
         try:
@@ -146,7 +146,7 @@ class ReplicateVisionService(BaseService):
     
     async def generate_and_save(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
         """生成图像并保存到本地 (此方法无需修改)"""
-        result = await self.generate_image(input_data)
+        result = await self._generate_image_internal(input_data)
         saved_paths = []
         for i, url in enumerate(result["urls"]):
             timestamp = int(time.time())
@@ -182,4 +182,136 @@ class ReplicateVisionService(BaseService):
         logger.info(f"Replicate Vision服务已准备就绪，使用模型: {self.model_name}")
 
     async def unload(self) -> None:
-        logger.info(f"卸载Replicate Vision服务: {self.model_name}")
+        logger.info(f"卸载Replicate Vision服务: {self.model_name}")
+
+    # 实现BaseImageGenService的抽象方法
+    async def generate_image(
+        self, 
+        prompt: 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 a single image from text prompt"""
+        input_data = {
+            "prompt": prompt,
+            "width": width,
+            "height": height,
+            "num_inference_steps": num_inference_steps,
+            "guidance_scale": guidance_scale
+        }
+        
+        if negative_prompt:
+            input_data["negative_prompt"] = negative_prompt
+        if seed:
+            input_data["seed"] = seed
+            
+        return await self._generate_image_internal(input_data)
+
+    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 = 20,
+        guidance_scale: float = 7.5,
+        seed: Optional[int] = None
+    ) -> List[Dict[str, Any]]:
+        """Generate multiple images from text prompt"""
+        results = []
+        for i in range(num_images):
+            current_seed = seed + i if seed else None
+            result = await self.generate_image(
+                prompt, negative_prompt, width, height, 
+                num_inference_steps, guidance_scale, current_seed
+            )
+            results.append(result)
+        return results
+
+    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"""
+        result = await self.generate_image(
+            prompt, negative_prompt, width, height, 
+            num_inference_steps, guidance_scale, seed
+        )
+        
+        # Save the first generated image to the specified path
+        if result.get("urls"):
+            url = result["urls"][0]
+            url_str = str(url) if hasattr(url, "__str__") else url
+            await self._download_image(url_str, output_path)
+            
+            return {
+                "file_path": output_path,
+                "width": width,
+                "height": height,
+                "seed": seed
+            }
+        else:
+            raise ValueError("No image generated")
+
+    async def image_to_image(
+        self,
+        prompt: str,
+        init_image: Union[str, Any],
+        strength: float = 0.8,
+        negative_prompt: Optional[str] = None,
+        num_inference_steps: int = 20,
+        guidance_scale: float = 7.5,
+        seed: Optional[int] = None
+    ) -> Dict[str, Any]:
+        """Generate image based on existing image and prompt"""
+        input_data = {
+            "prompt": prompt,
+            "image": init_image,
+            "strength": strength,
+            "num_inference_steps": num_inference_steps,
+            "guidance_scale": guidance_scale
+        }
+        
+        if negative_prompt:
+            input_data["negative_prompt"] = negative_prompt
+        if seed:
+            input_data["seed"] = seed
+            
+        return await self._generate_image_internal(input_data)
+
+    def get_supported_sizes(self) -> List[Dict[str, int]]:
+        """Get list of supported image dimensions"""
+        return [
+            {"width": 512, "height": 512},
+            {"width": 768, "height": 768},
+            {"width": 1024, "height": 1024},
+            {"width": 768, "height": 1344},
+            {"width": 1344, "height": 768},
+        ]
+
+    def get_model_info(self) -> Dict[str, Any]:
+        """Get information about the image generation model"""
+        return {
+            "name": self.model_name,
+            "max_width": 1344,
+            "max_height": 1344,
+            "supports_negative_prompt": True,
+            "supports_img2img": True
+        }
+
+    async def close(self):
+        """Cleanup resources"""
+        await self.unload()
+

isa_model/training/__init__.py

@@ -1,44 +1,74 @@
 """
-ISA Model Training Framework
+ISA Model Training Module
 
-This module provides unified interfaces for training various types of AI models:
-- LLM training with LlamaFactory
-- Image model training with Flux/LoRA
-- Model evaluation and benchmarking
+Provides unified training capabilities for AI models including:
+- Local training with SFT (Supervised Fine-Tuning)
+- Cloud training on RunPod
+- Model evaluation and management
+- HuggingFace integration
 
-Usage:
-    from isa_model.training import TrainingFactory
+Example usage:
+    ```python
+    from isa_model.training import TrainingFactory, train_gemma
     
-    # Create training factory
-    factory = TrainingFactory()
+    # Quick Gemma training
+    model_path = train_gemma(
+        dataset_path="tatsu-lab/alpaca",
+        model_size="4b",
+        num_epochs=3
+    )
     
-    # Fine-tune Gemma 3:4B
-    model_path = factory.finetune_llm(
-        model_name="gemma:4b",
-        dataset_path="path/to/data.json",
-        training_type="sft"
+    # Advanced training with custom configuration
+    factory = TrainingFactory()
+    model_path = factory.train_model(
+        model_name="google/gemma-2-4b-it",
+        dataset_path="your-dataset.json",
+        use_lora=True,
+        batch_size=4,
+        num_epochs=3
     )
+    ```
 """
 
-from .factory import TrainingFactory, finetune_gemma
-from .engine.llama_factory import (
-    LlamaFactory,
-    LlamaFactoryConfig,
-    SFTConfig,
-    RLConfig,
-    DPOConfig,
-    TrainingStrategy,
-    DatasetFormat
+# Import the new clean factory
+from .factory import TrainingFactory, train_gemma
+
+# Import core components
+from .core import (
+    TrainingConfig,
+    LoRAConfig, 
+    DatasetConfig,
+    BaseTrainer,
+    SFTTrainer,
+    TrainingUtils,
+    DatasetManager
+)
+
+# Import cloud training components
+from .cloud import (
+    RunPodConfig,
+    StorageConfig,
+    JobConfig,
+    TrainingJobOrchestrator
 )
 
 __all__ = [
-    "TrainingFactory",
-    "finetune_gemma",
-    "LlamaFactory", 
-    "LlamaFactoryConfig",
-    "SFTConfig",
-    "RLConfig", 
-    "DPOConfig",
-    "TrainingStrategy",
-    "DatasetFormat"
+    # Main factory
+    'TrainingFactory',
+    'train_gemma',
+    
+    # Core components
+    'TrainingConfig',
+    'LoRAConfig',
+    'DatasetConfig', 
+    'BaseTrainer',
+    'SFTTrainer',
+    'TrainingUtils',
+    'DatasetManager',
+    
+    # Cloud components
+    'RunPodConfig',
+    'StorageConfig',
+    'JobConfig',
+    'TrainingJobOrchestrator'
 ] 

isa_model/training/cloud/__init__.py

@@ -0,0 +1,22 @@
+"""
+Cloud Training Module for ISA Model Framework
+
+This module provides cloud training capabilities including:
+- RunPod integration for on-demand GPU training
+- Remote storage management (S3, GCS, etc.)
+- Training job orchestration and monitoring
+- Automatic resource scaling and management
+"""
+
+from .runpod_trainer import RunPodTrainer
+from .storage_manager import CloudStorageManager
+from .job_orchestrator import TrainingJobOrchestrator
+
+# Import config classes - these are defined in each module that needs them
+# from ..core.config import RunPodConfig, StorageConfig, JobConfig
+
+__all__ = [
+    "RunPodTrainer",
+    "CloudStorageManager", 
+    "TrainingJobOrchestrator"
+]