| 1 | import os |
| 2 | import time |
| 3 | import uuid |
| 4 | import logging |
| 5 | from typing import List, Optional |
| 6 | from .config import Config |
| 7 | |
| 8 | try: |
| 9 | from .image_dashscope import DashScopeClient |
| 10 | from .image_seedream import SeedreamClient |
| 11 | from .image_gpt import ImageGPT |
| 12 | from .image_processor import ImageProcessor |
| 13 | except ImportError: |
| 14 | from .image_dashscope import DashScopeClient |
| 15 | from .image_seedream import SeedreamClient |
| 16 | from .image_gpt import ImageGPT |
| 17 | from .image_processor import ImageProcessor |
| 18 | |
| 19 | class ImageClient: |
| 20 | def __init__(self, |
| 21 | dashscope_api_key: Optional[str] = None, |
| 22 | dashscope_base_url: Optional[str] = None, |
| 23 | dashscope_local_proxy: Optional[str] = None, |
| 24 | gpt_api_key: Optional[str] = None, |
| 25 | gpt_base_url: Optional[str] = None, |
| 26 | local_proxy: Optional[str] = None, |
| 27 | ark_api_key: Optional[str] = None, |
| 28 | ark_base_url: Optional[str] = None, |
| 29 | ark_local_proxy: Optional[str] = None): |
| 30 | """ |
| 31 | Unified Image Generation Client |
| 32 | Routes requests to DashScope, Seedream, or GPT based on model name. |
| 33 | """ |
| 34 | self._dashscope_api_key = dashscope_api_key or Config.DASHSCOPE_API_KEY |
| 35 | self._dashscope_base_url = dashscope_base_url or Config.DASHSCOPE_BASE_URL |
| 36 | self._dashscope_local_proxy = dashscope_local_proxy |
| 37 | |
| 38 | self._gpt_api_key = gpt_api_key or Config.OPENAI_API_KEY |
| 39 | self._gpt_base_url = gpt_base_url or Config.OPENAI_BASE_URL |
| 40 | self._gpt_local_proxy = local_proxy or Config.LOCAL_PROXY |
| 41 | |
| 42 | self._ark_api_key = ark_api_key or Config.ARK_API_KEY |
| 43 | self._ark_base_url = ark_base_url or Config.ARK_BASE_URL |
| 44 | self._ark_local_proxy = ark_local_proxy |
| 45 | |
| 46 | self._dashscope_client = None |
| 47 | self._seedream_client = None |
| 48 | self._gpt_client = None |
| 49 | |
| 50 | # Initialize Image Processor for downloads |
| 51 | self.image_processor = ImageProcessor() |
| 52 | |
| 53 | # Default save directory |
| 54 | self.base_save_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "code", "result", "image_client") |
| 55 | |
| 56 | @property |
| 57 | def dashscope_client(self): |
| 58 | """Create DashScope client only when a DashScope model is selected.""" |
| 59 | if self._dashscope_client is None: |
| 60 | self._dashscope_client = DashScopeClient( |
| 61 | api_key=self._dashscope_api_key, |
| 62 | base_url=self._dashscope_base_url, |
| 63 | local_proxy=self._dashscope_local_proxy, |
| 64 | ) |
| 65 | return self._dashscope_client |
| 66 | |
| 67 | @property |
| 68 | def seedream_client(self): |
| 69 | """Create Seedream client only when a Seedream/ARK model is selected.""" |
| 70 | if not self._ark_api_key: |
| 71 | raise RuntimeError("ARK_API_KEY not set. Configure ARK only when using Seedream image models.") |
| 72 | if self._seedream_client is None: |
| 73 | self._seedream_client = SeedreamClient( |
| 74 | api_key=self._ark_api_key, |
| 75 | base_url=self._ark_base_url, |
| 76 | local_proxy=self._ark_local_proxy, |
| 77 | ) |
| 78 | return self._seedream_client |
| 79 | |
| 80 | @property |
| 81 | def gpt_client(self): |
| 82 | """Create OpenAI image client only when a GPT/OpenAI image model is selected.""" |
| 83 | if not self._gpt_api_key: |
| 84 | raise RuntimeError("OPENAI_API_KEY not set. Configure OpenAI only when using GPT image models.") |
| 85 | if self._gpt_client is None: |
| 86 | self._gpt_client = ImageGPT( |
| 87 | api_key=self._gpt_api_key, |
| 88 | base_url=self._gpt_base_url, |
| 89 | local_proxy=self._gpt_local_proxy, |
| 90 | ) |
| 91 | return self._gpt_client |
| 92 | |
| 93 | def generate_image(self, |
| 94 | prompt: str, |
| 95 | image_paths: Optional[List[str]] = None, |
| 96 | model: str = "wan2.7-image", |
| 97 | save_dir: Optional[str] = None, |
| 98 | session_id: Optional[str] = None, |
| 99 | video_ratio: Optional[str] = "16:9", |
| 100 | resolution: Optional[str] = "2K") -> List[str]: |
| 101 | """ |
| 102 | Generate images based on prompt and optional reference images. |
| 103 | |
| 104 | Args: |
| 105 | prompt: Text prompt for generation. |
| 106 | image_paths: List of local file paths or URLs for reference images. |
| 107 | model: Model name to determine which provider to use. |
| 108 | save_dir: Custom directory to save downloaded images. |
| 109 | session_id: Session ID for organizing saved files. |
| 110 | video_ratio: Aspect ratio of the video, e.g., "16:9", "9:16", "4:3", "3:4", "1:1". |
| 111 | resolution: Resolution string, e.g., "720P", "1080P", "2K", "4K". |
| 112 | |
| 113 | Returns: |
| 114 | List of absolute file paths of the generated images. |
| 115 | """ |
| 116 | # Determine size from video_ratio and resolution |
| 117 | size_map = { |
| 118 | "16:9": { |
| 119 | "720P": "1280*720", |
| 120 | "1080P": "1920*1080", |
| 121 | "2K": "2560*1440", |
| 122 | "4K": "3840*2160" |
| 123 | }, |
| 124 | "9:16": { |
| 125 | "720P": "720*1280", |
| 126 | "1080P": "1080*1920", |
| 127 | "2K": "1440*2560", |
| 128 | "4K": "2160*3840" |
| 129 | }, |
| 130 | "4:3": { |
| 131 | "720P": "960*720", |
| 132 | "1080P": "1440*1080", |
| 133 | "2K": "2560*1920", |
| 134 | "4K": "3840*2880" |
| 135 | }, |
| 136 | "3:4": { |
| 137 | "720P": "720*960", |
| 138 | "1080P": "1080*1440", |
| 139 | "2K": "1920*2560", |
| 140 | "4K": "2880*3840" |
| 141 | }, |
| 142 | "1:1": { |
| 143 | "720P": "720*720", |
| 144 | "1080P": "1080*1080", |
| 145 | "2K": "2560*2560", |
| 146 | "4K": "3840*3840" |
| 147 | } |
| 148 | } |
| 149 | |
| 150 | # Default fallback if ratio or resolution is not found |
| 151 | size = size_map.get(video_ratio, size_map["16:9"]).get(resolution, "1920*1080") |
| 152 | |
| 153 | if not model: |
| 154 | model = "wan2.7-image" # Default model |
| 155 | |
| 156 | if Config.PRINT_MODEL_INPUT: |
| 157 | print("---- IMAGE GENERATION REQUEST ----") |
| 158 | print(f"Prompt: {prompt}") |
| 159 | if image_paths: |
| 160 | print(f"Refs: {len(image_paths)}") |
| 161 | for p in image_paths: |
| 162 | if str(p).startswith("data:"): |
| 163 | print(f" - [Base64图片]") |
| 164 | else: |
| 165 | print(f" - {p}") |
| 166 | print(f"Model: {model}") |
| 167 | print(f"Video Ratio: {video_ratio}") |
| 168 | print(f"Resolution: {resolution}") |
| 169 | print(f"Final Size: {size}") |
| 170 | if session_id: |
| 171 | print(f"Session ID: {session_id}") |
| 172 | print("-" * 30) |
| 173 | |
| 174 | # Determine backend provider |
| 175 | is_seedream = "seedream" in model.lower() |
| 176 | is_sora = "sora" in model.lower() or "gpt" in model.lower() |
| 177 | |
| 178 | # Prepare save directory |
| 179 | if not save_dir: |
| 180 | if session_id: |
| 181 | save_dir = os.path.join(self.base_save_dir, session_id) |
| 182 | else: |
| 183 | save_dir = self.base_save_dir |
| 184 | os.makedirs(save_dir, exist_ok=True) |
| 185 | |
| 186 | generated_local_paths = [] |
| 187 | |
| 188 | if is_seedream: |
| 189 | # --- Seedream Logic --- |
| 190 | try: |
| 191 | logging.info(f"ImageClient requesting Seedream: {model}") |
| 192 | |
| 193 | paths = self.seedream_client.generate_image( |
| 194 | prompt=prompt, |
| 195 | model=model, |
| 196 | session_id=session_id or "default", |
| 197 | size=size or "2048*2048", |
| 198 | image_paths=image_paths |
| 199 | ) |
| 200 | |
| 201 | if paths: |
| 202 | generated_local_paths.extend(paths) |
| 203 | |
| 204 | except Exception as e: |
| 205 | logging.error(f"Seedream generation failed: {e}") |
| 206 | |
| 207 | elif is_sora: |
| 208 | # --- GPT/Sora Logic --- |
| 209 | try: |
| 210 | logging.info(f"ImageClient requesting GPT/Sora: {model}") |
| 211 | if image_paths: |
| 212 | logging.warning("Sora/GPT model only supports Text-to-Image. Ignoring reference images.") |
| 213 | |
| 214 | # OpenAI uses 'x' separator, e.g. 1024x1024 |
| 215 | # Attempt to map size if needed or just replace '*' |
| 216 | gpt_size = size.replace('*', 'x') if size else "1024x1024" |
| 217 | |
| 218 | path = self.gpt_client.generate_image( |
| 219 | prompt=prompt, |
| 220 | size=gpt_size, |
| 221 | model=model, |
| 222 | save_dir=save_dir |
| 223 | ) |
| 224 | |
| 225 | if path and os.path.exists(path): |
| 226 | generated_local_paths.append(path) |
| 227 | else: |
| 228 | logging.error(f"GPT/Sora returned invalid path or download failed: {path}") |
| 229 | |
| 230 | except Exception as e: |
| 231 | logging.error(f"GPT/Sora generation failed: {e}") |
| 232 | |
| 233 | else: |
| 234 | # --- DashScope Logic --- |
| 235 | try: |
| 236 | logging.info(f"ImageClient requesting DashScope: {model}") |
| 237 | |
| 238 | if image_paths and len(image_paths) > 0: |
| 239 | # Pre-process image paths for DashScope |
| 240 | # Convert local paths to file:// URIs if they aren't already URLs |
| 241 | # DashScope SDK (via MultiModalConversation) handles file:// |
| 242 | formatted_urls = [] |
| 243 | for p in image_paths: |
| 244 | if p.startswith("http") or p.startswith("file://"): |
| 245 | formatted_urls.append(p) |
| 246 | else: |
| 247 | abs_path = os.path.abspath(p) |
| 248 | formatted_urls.append(f"file://{abs_path}") |
| 249 | |
| 250 | paths = self.dashscope_client.edit_image( |
| 251 | prompt=prompt, |
| 252 | image_urls=formatted_urls, |
| 253 | model=model, |
| 254 | size=size, |
| 255 | session_id=session_id, |
| 256 | save_dir=save_dir |
| 257 | ) |
| 258 | else: |
| 259 | # Text to Image |
| 260 | # Assuming default size 1024*1024 or similar |
| 261 | paths = self.dashscope_client.generate_image( |
| 262 | prompt=prompt, |
| 263 | model=model, |
| 264 | size=size, |
| 265 | session_id=session_id, |
| 266 | save_dir=save_dir |
| 267 | ) |
| 268 | |
| 269 | if paths: |
| 270 | generated_local_paths.extend(paths) |
| 271 | |
| 272 | except Exception as e: |
| 273 | logging.error(f"DashScope generation failed: {e}") |
| 274 | raise RuntimeError(f"DashScope generation failed: {e}") from e |
| 275 | |
| 276 | return generated_local_paths |
| 277 |