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