isa-model 0.3.3__py3-none-any.whl → 0.3.5__py3-none-any.whl

isa_model/inference/services/vision/replicate_image_gen_service.py

@@ -44,10 +44,6 @@ class ReplicateImageGenService(BaseImageGenService):
         # 设置 API token
         os.environ["REPLICATE_API_TOKEN"] = self.api_token
         
-        # 生成图像存储目录
-        self.output_dir = "generated_images"
-        os.makedirs(self.output_dir, exist_ok=True)
-        
         # 统计信息
         self.last_generation_count = 0
         self.total_generation_count = 0
@@ -133,6 +129,161 @@ class ReplicateImageGenService(BaseImageGenService):
         
         return await self._generate_internal(input_data)
 
+    async def instant_id_generation(
+        self,
+        prompt: str,
+        face_image: Union[str, Any],
+        negative_prompt: Optional[str] = None,
+        num_inference_steps: int = 30,
+        guidance_scale: float = 5.0,
+        seed: Optional[int] = None,
+        identitynet_strength_ratio: float = 0.8,
+        adapter_strength_ratio: float = 0.8
+    ) -> Dict[str, Any]:
+        """InstantID人脸一致性生成"""
+        
+        if "instant-id" in self.model_name:
+            input_data = {
+                "prompt": prompt,
+                "image": face_image,
+                "guidance_scale": guidance_scale,
+                "num_inference_steps": num_inference_steps,
+                "identitynet_strength_ratio": identitynet_strength_ratio,
+                "adapter_strength_ratio": adapter_strength_ratio
+            }
+            
+            if negative_prompt:
+                input_data["negative_prompt"] = negative_prompt
+            if seed:
+                input_data["seed"] = seed
+        else:
+            # 默认InstantID参数
+            input_data = {
+                "prompt": prompt,
+                "face_image": face_image,
+                "negative_prompt": negative_prompt or "",
+                "num_inference_steps": num_inference_steps,
+                "guidance_scale": guidance_scale,
+                "identitynet_strength_ratio": identitynet_strength_ratio,
+                "adapter_strength_ratio": adapter_strength_ratio
+            }
+            
+            if seed:
+                input_data["seed"] = seed
+        
+        return await self._generate_internal(input_data)
+
+    async def consistent_character_generation(
+        self,
+        subject: Union[str, Any],
+        prompt: Optional[str] = None,
+        negative_prompt: Optional[str] = None,
+        number_of_images: int = 4,
+        disable_safety_checker: bool = False
+    ) -> Dict[str, Any]:
+        """一致性角色生成 - 生成同一角色的多种姿态和表情"""
+        
+        if "consistent-character" in self.model_name:
+            input_data = {
+                "subject": subject,
+                "number_of_images": number_of_images,
+                "disable_safety_checker": disable_safety_checker
+            }
+            
+            if prompt:
+                input_data["prompt"] = prompt
+            if negative_prompt:
+                input_data["negative_prompt"] = negative_prompt
+        else:
+            # 默认一致性角色参数
+            input_data = {
+                "subject_image": subject,
+                "prompt": prompt or "portrait, different poses and expressions",
+                "negative_prompt": negative_prompt or "low quality, blurry",
+                "num_images": number_of_images
+            }
+        
+        return await self._generate_internal(input_data)
+
+    async def flux_lora_generation(
+        self,
+        prompt: str,
+        lora_scale: float = 1.0,
+        num_outputs: int = 1,
+        aspect_ratio: str = "1:1",
+        output_format: str = "jpg",
+        guidance_scale: float = 3.5,
+        output_quality: int = 90,
+        num_inference_steps: int = 28,
+        disable_safety_checker: bool = False
+    ) -> Dict[str, Any]:
+        """FLUX LoRA生成 - 使用预训练的LoRA权重"""
+        
+        if any(lora in self.model_name for lora in ["flux-dev-lora", "flux-lora"]):
+            input_data = {
+                "prompt": prompt,
+                "lora_scale": lora_scale,
+                "num_outputs": num_outputs,
+                "aspect_ratio": aspect_ratio,
+                "output_format": output_format,
+                "guidance_scale": guidance_scale,
+                "output_quality": output_quality,
+                "num_inference_steps": num_inference_steps,
+                "disable_safety_checker": disable_safety_checker
+            }
+        else:
+            # 默认LoRA参数
+            input_data = {
+                "prompt": prompt,
+                "lora_strength": lora_scale,
+                "num_images": num_outputs,
+                "guidance_scale": guidance_scale,
+                "num_inference_steps": num_inference_steps
+            }
+        
+        return await self._generate_internal(input_data)
+
+    async def ultimate_upscale(
+        self,
+        image: Union[str, Any],
+        scale: int = 4,
+        scheduler: str = "K_EULER_ANCESTRAL",
+        num_inference_steps: int = 20,
+        guidance_scale: float = 10.0,
+        strength: float = 0.55,
+        hdr: float = 0.0,
+        seed: Optional[int] = None
+    ) -> Dict[str, Any]:
+        """Ultimate SD Upscaler - 专业超分辨率"""
+        
+        if "ultimate" in self.model_name or "upscal" in self.model_name:
+            input_data = {
+                "image": image,
+                "scale": scale,
+                "scheduler": scheduler,
+                "num_inference_steps": num_inference_steps,
+                "guidance_scale": guidance_scale,
+                "strength": strength,
+                "hdr": hdr
+            }
+            
+            if seed:
+                input_data["seed"] = seed
+        else:
+            # 默认超分辨率参数
+            input_data = {
+                "image": image,
+                "upscale_factor": scale,
+                "num_inference_steps": num_inference_steps,
+                "guidance_scale": guidance_scale,
+                "denoising_strength": strength
+            }
+            
+            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:
@@ -141,11 +292,19 @@ class ReplicateImageGenService(BaseImageGenService):
             # 调用 Replicate API
             output = await replicate.async_run(self.model_name, input=input_data)
             
-            # 处理输出
+            # 处理输出 - 转换FileOutput对象为URL字符串
             if isinstance(output, list):
-                urls = output
+                raw_urls = output
             else:
-                urls = [output]
+                raw_urls = [output]
+            
+            # 转换为字符串URL
+            urls = []
+            for url in raw_urls:
+                if hasattr(url, 'url'):
+                    urls.append(str(url.url))  # type: ignore
+                else:
+                    urls.append(str(url))
 
             # 更新统计
             self.last_generation_count = len(urls)
@@ -224,37 +383,6 @@ class ReplicateImageGenService(BaseImageGenService):
             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)
-            
-            return {
-                "file_path": output_path,
-                "cost_usd": result.get("cost_usd", 0.0),
-                "model": self.model_name
-            }
-        else:
-            raise ValueError("No image generated")
-
     async def _download_image(self, url: str, save_path: str) -> None:
         """下载图像并保存"""
         try: