import os import time import uuid import logging from typing import List, Optional from .config import Config try: from .image_dashscope import DashScopeClient from .image_seedream import SeedreamClient from .image_gpt import ImageGPT from .image_processor import ImageProcessor except ImportError: from .image_dashscope import DashScopeClient from .image_seedream import SeedreamClient from .image_gpt import ImageGPT from .image_processor import ImageProcessor class ImageClient: def __init__(self, dashscope_api_key: Optional[str] = None, dashscope_base_url: Optional[str] = None, dashscope_local_proxy: Optional[str] = None, gpt_api_key: Optional[str] = None, gpt_base_url: Optional[str] = None, local_proxy: Optional[str] = None, ark_api_key: Optional[str] = None, ark_base_url: Optional[str] = None, ark_local_proxy: Optional[str] = None): """ Unified Image Generation Client Routes requests to DashScope, Seedream, or GPT based on model name. """ self._dashscope_api_key = dashscope_api_key or Config.DASHSCOPE_API_KEY self._dashscope_base_url = dashscope_base_url or Config.DASHSCOPE_BASE_URL self._dashscope_local_proxy = dashscope_local_proxy self._gpt_api_key = gpt_api_key or Config.OPENAI_API_KEY self._gpt_base_url = gpt_base_url or Config.OPENAI_BASE_URL self._gpt_local_proxy = local_proxy or Config.LOCAL_PROXY self._ark_api_key = ark_api_key or Config.ARK_API_KEY self._ark_base_url = ark_base_url or Config.ARK_BASE_URL self._ark_local_proxy = ark_local_proxy self._dashscope_client = None self._seedream_client = None self._gpt_client = None # Initialize Image Processor for downloads self.image_processor = ImageProcessor() # Default save directory self.base_save_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "code", "result", "image_client") @property def dashscope_client(self): """Create DashScope client only when a DashScope model is selected.""" if self._dashscope_client is None: self._dashscope_client = DashScopeClient( api_key=self._dashscope_api_key, base_url=self._dashscope_base_url, local_proxy=self._dashscope_local_proxy, ) return self._dashscope_client @property def seedream_client(self): """Create Seedream client only when a Seedream/ARK model is selected.""" if not self._ark_api_key: raise RuntimeError("ARK_API_KEY not set. Configure ARK only when using Seedream image models.") if self._seedream_client is None: self._seedream_client = SeedreamClient( api_key=self._ark_api_key, base_url=self._ark_base_url, local_proxy=self._ark_local_proxy, ) return self._seedream_client @property def gpt_client(self): """Create OpenAI image client only when a GPT/OpenAI image model is selected.""" if not self._gpt_api_key: raise RuntimeError("OPENAI_API_KEY not set. Configure OpenAI only when using GPT image models.") if self._gpt_client is None: self._gpt_client = ImageGPT( api_key=self._gpt_api_key, base_url=self._gpt_base_url, local_proxy=self._gpt_local_proxy, ) return self._gpt_client def generate_image(self, prompt: str, image_paths: Optional[List[str]] = None, model: str = "wan2.7-image", save_dir: Optional[str] = None, session_id: Optional[str] = None, video_ratio: Optional[str] = "16:9", resolution: Optional[str] = "2K") -> List[str]: """ Generate images based on prompt and optional reference images. Args: prompt: Text prompt for generation. image_paths: List of local file paths or URLs for reference images. model: Model name to determine which provider to use. save_dir: Custom directory to save downloaded images. session_id: Session ID for organizing saved files. video_ratio: Aspect ratio of the video, e.g., "16:9", "9:16", "4:3", "3:4", "1:1". resolution: Resolution string, e.g., "720P", "1080P", "2K", "4K". Returns: List of absolute file paths of the generated images. """ # Determine size from video_ratio and resolution size_map = { "16:9": { "720P": "1280*720", "1080P": "1920*1080", "2K": "2560*1440", "4K": "3840*2160" }, "9:16": { "720P": "720*1280", "1080P": "1080*1920", "2K": "1440*2560", "4K": "2160*3840" }, "4:3": { "720P": "960*720", "1080P": "1440*1080", "2K": "2560*1920", "4K": "3840*2880" }, "3:4": { "720P": "720*960", "1080P": "1080*1440", "2K": "1920*2560", "4K": "2880*3840" }, "1:1": { "720P": "720*720", "1080P": "1080*1080", "2K": "2560*2560", "4K": "3840*3840" } } # Default fallback if ratio or resolution is not found size = size_map.get(video_ratio, size_map["16:9"]).get(resolution, "1920*1080") if not model: model = "wan2.7-image" # Default model if Config.PRINT_MODEL_INPUT: print("---- IMAGE GENERATION REQUEST ----") print(f"Prompt: {prompt}") if image_paths: print(f"Refs: {len(image_paths)}") for p in image_paths: if str(p).startswith("data:"): print(f" - [Base64图片]") else: print(f" - {p}") print(f"Model: {model}") print(f"Video Ratio: {video_ratio}") print(f"Resolution: {resolution}") print(f"Final Size: {size}") if session_id: print(f"Session ID: {session_id}") print("-" * 30) # Determine backend provider is_seedream = "seedream" in model.lower() is_sora = "sora" in model.lower() or "gpt" in model.lower() # Prepare save directory if not save_dir: if session_id: save_dir = os.path.join(self.base_save_dir, session_id) else: save_dir = self.base_save_dir os.makedirs(save_dir, exist_ok=True) generated_local_paths = [] if is_seedream: # --- Seedream Logic --- try: logging.info(f"ImageClient requesting Seedream: {model}") paths = self.seedream_client.generate_image( prompt=prompt, model=model, session_id=session_id or "default", size=size or "2048*2048", image_paths=image_paths ) if paths: generated_local_paths.extend(paths) except Exception as e: logging.error(f"Seedream generation failed: {e}") elif is_sora: # --- GPT/Sora Logic --- try: logging.info(f"ImageClient requesting GPT/Sora: {model}") if image_paths: logging.warning("Sora/GPT model only supports Text-to-Image. Ignoring reference images.") # OpenAI uses 'x' separator, e.g. 1024x1024 # Attempt to map size if needed or just replace '*' gpt_size = size.replace('*', 'x') if size else "1024x1024" path = self.gpt_client.generate_image( prompt=prompt, size=gpt_size, model=model, save_dir=save_dir ) if path and os.path.exists(path): generated_local_paths.append(path) else: logging.error(f"GPT/Sora returned invalid path or download failed: {path}") except Exception as e: logging.error(f"GPT/Sora generation failed: {e}") else: # --- DashScope Logic --- try: logging.info(f"ImageClient requesting DashScope: {model}") if image_paths and len(image_paths) > 0: # Pre-process image paths for DashScope # Convert local paths to file:// URIs if they aren't already URLs # DashScope SDK (via MultiModalConversation) handles file:// formatted_urls = [] for p in image_paths: if p.startswith("http") or p.startswith("file://"): formatted_urls.append(p) else: abs_path = os.path.abspath(p) formatted_urls.append(f"file://{abs_path}") paths = self.dashscope_client.edit_image( prompt=prompt, image_urls=formatted_urls, model=model, size=size, session_id=session_id, save_dir=save_dir ) else: # Text to Image # Assuming default size 1024*1024 or similar paths = self.dashscope_client.generate_image( prompt=prompt, model=model, size=size, session_id=session_id, save_dir=save_dir ) if paths: generated_local_paths.extend(paths) except Exception as e: logging.error(f"DashScope generation failed: {e}") raise RuntimeError(f"DashScope generation failed: {e}") from e return generated_local_paths