isa-model 0.2.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

isa_model/inference/services/vision/replicate_image_gen_service.py

@@ -2,168 +2,261 @@
 # -*- coding: utf-8 -*-
 
 """
-Replicate Vision服务
-用于与Replicate API交互，支持图像生成和图像分析
+Replicate 图像生成服务
+支持 flux-schnell (文生图) 和 flux-kontext-pro (图生图) 模型
 """
 
 import os
 import time
 import uuid
 import logging
-from typing import Dict, Any, List, Optional, Union, Tuple
+from typing import Dict, Any, List, Optional, Union
 import asyncio
 import aiohttp
-import replicate # 导入 replicate 库
+import 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
 
 # 设置日志记录
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-class ReplicateVisionService(BaseService):
+class ReplicateImageGenService(BaseImageGenService):
     """
-    Replicate Vision服务，用于处理图像生成和分析。
-    经过调整，使用原生异步调用并优化了文件处理。
+    Replicate 图像生成服务
+    - flux-schnell: 文生图 (t2i) - $3 per 1000 images
+    - flux-kontext-pro: 图生图 (i2i) - $0.04 per image
     """
     
     def __init__(self, provider: BaseProvider, model_name: str):
-        """
-        初始化Replicate Vision服务
-        """
         super().__init__(provider, model_name)
-        # 从 provider 或环境变量获取 API token
-        self.api_token = self.provider.config.get("api_token", os.environ.get("REPLICATE_API_TOKEN"))
-        self.model_type = ModelType.VISION
         
-        # 可选的默认配置
-        self.guidance_scale = self.provider.config.get("guidance_scale", 7.5)
-        self.num_inference_steps = self.provider.config.get("num_inference_steps", 30)
+        # 获取配置
+        provider_config = provider.get_full_config()
+        self.api_token = provider_config.get("api_token") or provider_config.get("replicate_api_token")
         
-        # 生成的图像存储目录
+        if not self.api_token:
+            raise ValueError("Replicate API token not found in provider configuration")
+        
+        # 设置 API token
+        os.environ["REPLICATE_API_TOKEN"] = self.api_token
+        
+        # 生成图像存储目录
         self.output_dir = "generated_images"
         os.makedirs(self.output_dir, exist_ok=True)
         
-        # ★ 调整点: 为 replicate 库设置 API token
-        if self.api_token:
-            # replicate 库会自动从环境变量读取，我们确保它被设置
-            os.environ["REPLICATE_API_TOKEN"] = self.api_token
+        # 统计信息
+        self.last_generation_count = 0
+        self.total_generation_count = 0
+        
+        logger.info(f"Initialized ReplicateImageGenService with model '{self.model_name}'")
+
+    async def generate_image(
+        self, 
+        prompt: str,
+        negative_prompt: Optional[str] = None,
+        width: int = 512,
+        height: int = 512,
+        num_inference_steps: int = 4,
+        guidance_scale: float = 7.5,
+        seed: Optional[int] = None
+    ) -> Dict[str, Any]:
+        """生成单张图像 (文生图)"""
+        
+        if "flux-schnell" in self.model_name:
+            # FLUX Schnell 参数
+            input_data = {
+                "prompt": prompt,
+                "go_fast": True,
+                "megapixels": "1",
+                "num_outputs": 1,
+                "aspect_ratio": "1:1",
+                "output_format": "jpg",
+                "output_quality": 90,
+                "num_inference_steps": 4
+            }
         else:
-            logger.warning("Replicate API token 未找到。服务可能无法正常工作。")
-
-    async def _prepare_input_files(self, input_data: Dict[str, Any]) -> Tuple[Dict[str, Any], List[Any]]:
-        """
-        ★ 新增辅助函数: 准备输入数据，将本地文件路径转换为文件对象。
-        这使得服务能统一处理本地文件和URL。
-        """
-        prepared_input = input_data.copy()
-        files_to_close = []
-        for key, value in prepared_input.items():
-            # 如果值是字符串，且看起来像一个存在的本地文件路径
-            if isinstance(value, str) and not value.startswith(('http://', 'https://')) and os.path.exists(value):
-                logger.info(f"检测到本地文件路径 '{value}'，准备打开文件。")
-                try:
-                    file_handle = open(value, "rb")
-                    prepared_input[key] = file_handle
-                    files_to_close.append(file_handle)
-                except Exception as e:
-                    logger.error(f"打开文件失败 '{value}': {e}")
-                    raise
-        return prepared_input, files_to_close
-
-    async def generate_image(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        使用Replicate模型生成图像 (已优化为原生异步)
-        """
-        prepared_input, files_to_close = await self._prepare_input_files(input_data)
-        try:
-            # 设置默认参数
-            if "guidance_scale" not in prepared_input:
-                prepared_input["guidance_scale"] = self.guidance_scale
-            if "num_inference_steps" not in prepared_input:
-                prepared_input["num_inference_steps"] = self.num_inference_steps
+            # 默认参数
+            input_data = {
+                "prompt": prompt,
+                "width": width,
+                "height": height,
+                "num_inference_steps": num_inference_steps,
+                "guidance_scale": guidance_scale
+            }
             
-            logger.info(f"开始使用模型 {self.model_name} 生成图像 (原生异步)")
+            if negative_prompt:
+                input_data["negative_prompt"] = negative_prompt
+            if seed:
+                input_data["seed"] = seed
+        
+        return await self._generate_internal(input_data)
+
+    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]:
+        """图生图"""
+        
+        if "flux-kontext-pro" in self.model_name:
+            # FLUX Kontext Pro 参数
+            input_data = {
+                "prompt": prompt,
+                "input_image": init_image,
+                "aspect_ratio": "match_input_image",
+                "output_format": "jpg",
+                "safety_tolerance": 2
+            }
+        else:
+            # 默认参数
+            input_data = {
+                "prompt": prompt,
+                "image": init_image,
+                "strength": strength,
+                "num_inference_steps": num_inference_steps,
+                "guidance_scale": guidance_scale
+            }
             
-            # ★ 调整点: 使用原生异步的 replicate.async_run
-            output = await replicate.async_run(self.model_name, input=prepared_input)
+            if negative_prompt:
+                input_data["negative_prompt"] = negative_prompt
+            if seed:
+                input_data["seed"] = seed
+        
+        return await self._generate_internal(input_data)
+
+    async def _generate_internal(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
+        """内部生成方法"""
+        try:
+            logger.info(f"开始使用模型 {self.model_name} 生成图像")
+            
+            # 调用 Replicate API
+            output = await replicate.async_run(self.model_name, input=input_data)
             
-            # 将结果转换为标准格式 (此部分逻辑无需改变)
+            # 处理输出
             if isinstance(output, list):
                 urls = output
             else:
                 urls = [output]
 
+            # 更新统计
+            self.last_generation_count = len(urls)
+            self.total_generation_count += len(urls)
+            
+            # 计算成本
+            cost = self._calculate_cost(len(urls))
+            
+            # 跟踪计费信息
+            from isa_model.inference.billing_tracker import ServiceType
+            self._track_usage(
+                service_type=ServiceType.IMAGE_GENERATION,
+                operation="image_generation",
+                input_units=len(urls),  # 生成的图像数量
+                metadata={
+                    "model": self.model_name,
+                    "prompt": input_data.get("prompt", "")[:100],  # 截取前100字符
+                    "generation_type": "t2i" if "flux-schnell" in self.model_name else "i2i"
+                }
+            )
+            
             result = {
                 "urls": urls,
+                "count": len(urls),
+                "cost_usd": cost,
+                "model": self.model_name,
                 "metadata": {
-                    "model": self.model_name,
-                    "input": input_data # 返回原始输入以供参考
+                    "input": input_data,
+                    "generation_count": len(urls)
                 }
             }
-            logger.info(f"图像生成完成: {result['urls']}")
+            
+            logger.info(f"图像生成完成: {len(urls)} 张图像, 成本: ${cost:.6f}")
             return result
+            
         except Exception as e:
             logger.error(f"图像生成失败: {e}")
             raise
-        finally:
-            # ★ 新增: 确保所有打开的文件都被关闭
-            for f in files_to_close:
-                f.close()
-    
-    async def analyze_image(self, image_path: str, prompt: str) -> Dict[str, Any]:
-        """
-        分析图像 (已优化为原生异步)
-        """
-        input_data = {"image": image_path, "prompt": prompt}
-        prepared_input, files_to_close = await self._prepare_input_files(input_data)
-        try:
-            logger.info(f"开始使用模型 {self.model_name} 分析图像 (原生异步)")
-            # ★ 调整点: 使用原生异步的 replicate.async_run
-            output = await replicate.async_run(self.model_name, input=prepared_input)
+
+    def _calculate_cost(self, image_count: int) -> float:
+        """计算生成成本"""
+        from isa_model.core.model_manager import ModelManager
+        
+        manager = ModelManager()
+        
+        if "flux-schnell" in self.model_name:
+            # $3 per 1000 images
+            return (image_count / 1000) * 3.0
+        elif "flux-kontext-pro" in self.model_name:
+            # $0.04 per image
+            return image_count * 0.04
+        else:
+            # 使用 ModelManager 的定价
+            pricing = manager.get_model_pricing("replicate", self.model_name)
+            return (image_count / 1000) * pricing.get("input", 0.0)
+
+    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 = 4,
+        guidance_scale: float = 7.5,
+        seed: Optional[int] = None
+    ) -> List[Dict[str, Any]]:
+        """生成多张图像"""
+        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 = 4,
+        guidance_scale: float = 7.5,
+        seed: Optional[int] = None
+    ) -> Dict[str, Any]:
+        """生成图像并保存到文件"""
+        result = await self.generate_image(
+            prompt, negative_prompt, width, height, 
+            num_inference_steps, guidance_scale, seed
+        )
+        
+        # 保存第一张图像
+        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)
             
-            result = {
-                "text": "".join(output) if isinstance(output, list) else output,
-                "metadata": {
-                    "model": self.model_name,
-                    "input": input_data
-                }
+            return {
+                "file_path": output_path,
+                "cost_usd": result.get("cost_usd", 0.0),
+                "model": self.model_name
             }
-            logger.info(f"图像分析完成")
-            return result
-        except Exception as e:
-            logger.error(f"图像分析失败: {e}")
-            raise
-        finally:
-            # ★ 新增: 确保所有打开的文件都被关闭
-            for f in files_to_close:
-                f.close()
-    
-    async def generate_and_save(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
-        """生成图像并保存到本地 (此方法无需修改)"""
-        result = await self.generate_image(input_data)
-        saved_paths = []
-        for i, url in enumerate(result["urls"]):
-            timestamp = int(time.time())
-            file_name = f"{self.output_dir}/{timestamp}_{uuid.uuid4().hex[:8]}_{i+1}.png"
-            try:
-                # Convert FileOutput object to string if necessary
-                url_str = str(url) if hasattr(url, "__str__") else url
-                await self._download_image(url_str, file_name)
-                saved_paths.append(file_name)
-                logger.info(f"图像已保存至: {file_name}")
-            except Exception as e:
-                logger.error(f"保存图像失败: {e}")
-        result["saved_paths"] = saved_paths
-        return result
-    
+        else:
+            raise ValueError("No image generated")
+
     async def _download_image(self, url: str, save_path: str) -> None:
-        """异步下载图像并保存 (此方法无需修改)"""
+        """下载图像并保存"""
         try:
             async with aiohttp.ClientSession() as session:
                 async with session.get(url) as response:
@@ -175,11 +268,78 @@ class ReplicateVisionService(BaseService):
             logger.error(f"下载图像时出错: {url}, {e}")
             raise
 
-    # `load` 和 `unload` 方法在Replicate API场景下通常是轻量级的
+    def get_generation_stats(self) -> Dict[str, Any]:
+        """获取生成统计信息"""
+        total_cost = 0.0
+        if "flux-schnell" in self.model_name:
+            total_cost = (self.total_generation_count / 1000) * 3.0
+        elif "flux-kontext-pro" in self.model_name:
+            total_cost = self.total_generation_count * 0.04
+        
+        return {
+            "last_generation_count": self.last_generation_count,
+            "total_generation_count": self.total_generation_count,
+            "total_cost_usd": total_cost,
+            "model": self.model_name
+        }
+
+    def get_supported_sizes(self) -> List[Dict[str, int]]:
+        """获取支持的图像尺寸"""
+        if "flux" in self.model_name:
+            return [
+                {"width": 512, "height": 512},
+                {"width": 768, "height": 768},
+                {"width": 1024, "height": 1024},
+            ]
+        else:
+            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]:
+        """获取模型信息"""
+        if "flux-schnell" in self.model_name:
+            return {
+                "name": self.model_name,
+                "type": "t2i",
+                "cost_per_1000_images": 3.0,
+                "supports_negative_prompt": False,
+                "supports_img2img": False,
+                "max_steps": 4
+            }
+        elif "flux-kontext-pro" in self.model_name:
+            return {
+                "name": self.model_name,
+                "type": "i2i",
+                "cost_per_image": 0.04,
+                "supports_negative_prompt": False,
+                "supports_img2img": True,
+                "max_width": 1024,
+                "max_height": 1024
+            }
+        else:
+            return {
+                "name": self.model_name,
+                "type": "general",
+                "supports_negative_prompt": True,
+                "supports_img2img": True
+            }
+
     async def load(self) -> None:
+        """加载服务"""
         if not self.api_token:
-            raise ValueError("缺少Replicate API令牌，请设置REPLICATE_API_TOKEN环境变量或在provider配置中提供")
-        logger.info(f"Replicate Vision服务已准备就绪，使用模型: {self.model_name}")
+            raise ValueError("缺少 Replicate API 令牌")
+        logger.info(f"Replicate 图像生成服务已准备就绪，使用模型: {self.model_name}")
 
     async def unload(self) -> None:
-        logger.info(f"卸载Replicate Vision服务: {self.model_name}")
+        """卸载服务"""
+        logger.info(f"卸载 Replicate 图像生成服务: {self.model_name}")
+
+    async def close(self):
+        """关闭服务"""
+        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"
+]