Update generator_function/image_function.py

generator_function/image_function.py

@@ -38,16 +38,137 @@ MAX_WORKERS = min(32, (os.cpu_count() or 1) + 4)
 REQUEST_TIMEOUT = 30
 RETRY_ATTEMPTS = 3
 
+
+def get_cost_per_image(
+    model_key: str,
+    width: int | None = None,
+    height: int | None = None,
+) -> float:
+    cfg = MODEL_REGISTRY.get(model_key, {})
+    price = cfg.get("price")
+    unit = cfg.get("price_unit")
+
+    if not price:
+        return 0.0
+
+    if unit == "per_image":
+        return round(price, 4)
+
+    return 0.0
+
 MODEL_REGISTRY: Dict[str, Dict[str, Any]] = {
-    "imagegen-4-ultra": {"id": "google/imagen-4-ultra","aspect_ratios": ["1:1","16:9","9:16","3:4","4:3"],"param_name":"aspect_ratio"},
-    "imagen-4": {"id": "google/imagen-4","aspect_ratios": ["1:1","16:9","9:16","3:4","4:3"],"param_name":"aspect_ratio"},
-    "nano-banana": {"id": "google/nano-banana","aspect_ratios": ["1:1","16:9","9:16","3:4","4:3"],"param_name":"aspect_ratio"},
-    "qwen": {"id": "qwen/qwen-image","aspect_ratios": ["1:1","16:9","9:16","3:4","4:3","3:2","2:3"],"param_name":"aspect_ratio"},
-    "seedream-3": {"id": "bytedance/seedream-3","aspect_ratios": ["1:1","16:9","9:16","3:4","4:3","3:2","2:3","21:9"],"param_name":"aspect_ratio"},
-    "recraft-v3": {"id": "recraft-ai/recraft-v3","aspect_ratios": ["1:1","4:3","3:4","3:2","2:3","16:9","9:16","1:2","2:1","7:5","5:7","4:5","5:4","3:5","5:3"],"param_name":"aspect_ratio"},
-    "photon": {"id": "luma/photon","aspect_ratios": ["1:1","3:4","4:3","9:16","16:9","9:21","21:9"],"param_name":"aspect_ratio"},
-    "ideogram-v3-quality": {"id": "ideogram-ai/ideogram-v3-quality","aspect_ratios": ["1:3","3:1","1:2","2:1","9:16","16:9","10:16","16:10","2:3","3:2","3:4","4:3","4:5","5:4","1:1"],"param_name":"aspect_ratio"},
-    "z-image-turbo": {"id": "prunaai/z-image-turbo","param_name": None,"supports_aspect_ratio": False}}
+    "imagegen-4-ultra": {
+        "id": "google/imagen-4-ultra",
+        "price": 0.06,
+        "price_unit": "per_image",
+        "aspect_ratios": ["1:1","16:9","9:16","3:4","4:3"],
+        "param_name": "aspect_ratio",
+    },
+
+    "imagen-4": {
+        "id": "google/imagen-4",
+        "price": 0.04,
+        "price_unit": "per_image",
+        "aspect_ratios": ["1:1","16:9","9:16","3:4","4:3"],
+        "param_name": "aspect_ratio",
+    },
+
+    "nano-banana": {
+        "id": "google/nano-banana",
+        "price": 0.039,
+        "price_unit": "per_image",
+        "aspect_ratios": ["1:1","16:9","9:16","3:4","4:3"],
+        "param_name": "aspect_ratio",
+    },
+
+    "qwen": {
+        "id": "qwen/qwen-image",
+        "price": 0.025,
+        "price_unit": "per_image",
+        "aspect_ratios": ["1:1","16:9","9:16","3:4","4:3","3:2","2:3"],
+        "param_name": "aspect_ratio",
+    },
+
+    "seedream-3": {
+        "id": "bytedance/seedream-3",
+        "price": 0.03,
+        "price_unit": "per_image",
+        "aspect_ratios": ["1:1","16:9","9:16","3:4","4:3","3:2","2:3","21:9"],
+        "param_name": "aspect_ratio",
+    },
+
+    "recraft-v3": {
+        "id": "recraft-ai/recraft-v3",
+        "price": 0.04,
+        "price_unit": "per_image",
+        "aspect_ratios": [
+            "1:1","4:3","3:4","3:2","2:3","16:9","9:16",
+            "1:2","2:1","7:5","5:7","4:5","5:4","3:5","5:3"
+        ],
+        "param_name": "aspect_ratio",
+    },
+
+    "photon": {
+        "id": "luma/photon",
+        "price": 0.03,
+        "price_unit": "per_image",
+        "aspect_ratios": ["1:1","3:4","4:3","9:16","16:9","9:21","21:9"],
+        "param_name": "aspect_ratio",
+    },
+
+    "ideogram-v3-quality": {
+        "id": "ideogram-ai/ideogram-v3-quality",
+        "price": 0.09,
+        "price_unit": "per_image",
+        "aspect_ratios": [
+            "1:3","3:1","1:2","2:1","9:16","16:9","10:16",
+            "16:10","2:3","3:2","3:4","4:3","4:5","5:4","1:1"
+        ],
+        "param_name": "aspect_ratio",
+    },
+
+    "ideogram-v3-turbo": {
+        "id": "ideogram-ai/ideogram-v3-turbo",
+        "price": 0.03,
+        "price_unit": "per_image",
+        "aspect_ratios": [
+            "1:3","3:1","1:2","2:1","9:16","16:9","10:16",
+            "16:10","2:3","3:2","3:4","4:3","4:5","5:4","1:1"
+        ],
+        "param_name": "aspect_ratio",
+    },
+
+    "flux-2-max": {
+        "id": "black-forest-labs/flux-2-max",
+        "price": 0.04,
+        "price_unit": "per_image",
+        "aspect_ratios": ["1:1","2:3","3:2","3:4","4:3","4:5","5:4","9:16","16:9"],
+        "param_name": "aspect_ratio",
+    },
+
+    "p-image": {
+        "id": "prunaai/p-image",
+        "price": 0.005,
+        "price_unit": "per_image",
+        "type": "custom",
+    },
+
+    "z-image-turbo": {
+        "id": "prunaai/z-image-turbo",
+        "price": 0.01,
+        "price_unit": "per_image",
+        "type": "custom",
+    },
+
+    "bytedance/seedream-4.5": {
+        "id": "bytedance/seedream-4.5",
+        "size": "2k",
+        "price": 0.04,
+        "price_unit": "per_image",
+        "aspect_ratios": ["1:1","2:3","3:2","3:4","4:3","9:16","16:9","21:9"],
+        "param_name": "aspect_ratio",
+    },
+}
 
 _thread_local = threading.local()
 
@@ -109,24 +230,27 @@ def _generate_one(model_key: str,prompt: str, aspect_ratio: str,params: Dict[str
         return []
     for attempt in range(RETRY_ATTEMPTS):
         try:
-            input_payload = {"prompt": prompt}
-            if cfg["id"] == "prunaai/z-image-turbo":
-                input_payload["width"] = params.get("width", 1024)
-                input_payload["height"] = params.get("height", 1024)
-                input_payload["output_quality"] = params.get("output_quality", 90)
-
-            elif cfg.get("param_name"):
-                input_payload[cfg["param_name"]] = aspect_ratio
+            if cfg.get("type") == "custom":
+                input_payload = {
+                    "prompt": prompt,
+                    "width": 1024,
+                    "height": 1024,
+                }
+            else:
+                param_name = cfg.get("param_name", "aspect_ratio")
+                input_payload = {"prompt": prompt,param_name: aspect_ratio}
 
             output = replicate.run(cfg["id"], input=input_payload)
 
+            urls: List[str] = []
             if isinstance(output, list) and output:
-                return [getattr(output[0], "url", str(output[0]))]
-            if isinstance(output, str):
-                return [output]
-            if hasattr(output, "url"):
-                return [output.url]
-
+                urls = [getattr(output[0], "url", str(output[0]))]
+            elif isinstance(output, str):
+                urls = [output]
+            elif hasattr(output, "url"):
+                urls = [output.url]
+            if urls:
+                return urls
         except Exception as e:
             if attempt == RETRY_ATTEMPTS - 1:
                 logger.error(f"replicate run failed: {e}")
@@ -407,4 +531,4 @@ def generate_image(file_path, size, quality, category, sentiment, user_prompt, p
         return image_bytes        
     except Exception as e:
         logger.exception(f"Failed to generate image for {file_path}: {e}")
-        raise
+        raise