Merge pull request #266 from Runware/feature-addtargetMegapixelsParam

Added async flow in imageUpscale, alias for IUpscaleSettings as ISettings

Author: Sirsho1997

runware/base.py

@@ -1274,7 +1274,11 @@ async def _removeImageBackground(
 
     async def imageUpscale(self, upscaleGanPayload: "IImageUpscale") -> "Union[List[IImage], IAsyncTaskResponse]":
         async with self._request_semaphore:
-            return await self._retry_with_reconnect(self._imageUpscale, upscaleGanPayload)
+            return await self._retry_async_with_reconnect(
+                self._imageUpscale,
+                upscaleGanPayload,
+                task_type=ETaskType.IMAGE_UPSCALE.value,
+            )
 
     async def _imageUpscale(self, upscaleGanPayload: IImageUpscale) -> Union[List[IImage], IAsyncTaskResponse]:
         await self.ensureConnection()
@@ -1297,12 +1301,15 @@ async def _upscaleGan(self, upscaleGanPayload: "IImageUpscale") -> "Union[List[I
         taskUUID = getUUID()
         upscaleGanPayload.taskUUID = taskUUID
 
-        # Create a dictionary with mandatory parameters
         task_params = {
             "taskType": ETaskType.IMAGE_UPSCALE.value,
             "taskUUID": taskUUID,
-            "upscaleFactor": upscaleGanPayload.upscaleFactor,
+            "deliveryMethod": upscaleGanPayload.deliveryMethod,
         }
+        if upscaleGanPayload.upscaleFactor is not None:
+            task_params["upscaleFactor"] = upscaleGanPayload.upscaleFactor
+        if upscaleGanPayload.targetMegapixels is not None:
+            task_params["targetMegapixels"] = upscaleGanPayload.targetMegapixels
 
         # Use inputs.image format if inputs is provided, otherwise use inputImage (legacy)
         if upscaleGanPayload.inputs and upscaleGanPayload.inputs.image:
@@ -1347,6 +1354,39 @@ async def _upscaleGan(self, upscaleGanPayload: "IImageUpscale") -> "Union[List[I
                 debug_key="image-upscale-webhook"
             )
 
+        delivery_method_enum = (
+            EDeliveryMethod(upscaleGanPayload.deliveryMethod)
+            if isinstance(upscaleGanPayload.deliveryMethod, str)
+            else upscaleGanPayload.deliveryMethod
+        )
+
+        if delivery_method_enum is EDeliveryMethod.ASYNC:
+            future, should_send = await self._register_pending_operation(
+                taskUUID,
+                expected_results=1,
+                complete_predicate=lambda r: True,
+            )
+            try:
+                if should_send:
+                    await self.send([task_params])
+                    await self._mark_operation_sent(taskUUID)
+                results = await asyncio.wait_for(future, timeout=IMAGE_INITIAL_TIMEOUT / 1000)
+                response = results[0]
+                self._handle_error_response(response)
+                if response.get("status") == "success" or response.get("imageUUID") is not None:
+                    image = createImageFromResponse(response)
+                    return [image]
+                return createAsyncTaskResponse(response)
+            except asyncio.TimeoutError:
+                raise ConnectionError(
+                    f"Timeout waiting for async image upscale acknowledgment | TaskUUID: {taskUUID} | "
+                    f"Timeout: {IMAGE_INITIAL_TIMEOUT}ms"
+                )
+            except RunwareAPIError:
+                raise
+            finally:
+                await self._unregister_pending_operation(taskUUID)
+
         future, should_send = await self._register_pending_operation(
             taskUUID,
             expected_results=1,
@@ -3102,6 +3142,17 @@ def configure_from_task_type(task_type_val: Optional[str]):
                         IMAGE_POLLING_DELAY,
                         f"Image generation timeout after {MAX_POLLS_IMAGE_GENERATION} polls"
                     )
+                case (
+                    ETaskType.IMAGE_UPSCALE.value
+                    | ETaskType.IMAGE_VECTORIZE.value
+                    | ETaskType.IMAGE_BACKGROUND_REMOVAL.value
+                ):
+                    return (
+                        IImage,
+                        MAX_POLLS_IMAGE_GENERATION,
+                        IMAGE_POLLING_DELAY,
+                        f"Image task timeout after {MAX_POLLS_IMAGE_GENERATION} polls"
+                    )
                 case (
                     ETaskType.VIDEO_INFERENCE.value
                     | ETaskType.VIDEO_BACKGROUND_REMOVAL.value

runware/types.py

@@ -866,6 +866,20 @@ class ISettings(SerializableMixin):
     expressiveness: Optional[str] = None
     removeBackground: Optional[bool] = None
     backgroundColor: Optional[str] = None
+    # Image upscale 
+    steps: Optional[int] = None
+    seed: Optional[int] = None
+    CFGScale: Optional[float] = None
+    positivePrompt: Optional[str] = None
+    negativePrompt: Optional[str] = None
+    controlNetWeight: Optional[float] = None
+    strength: Optional[float] = None
+    scheduler: Optional[str] = None
+    colorFix: Optional[bool] = None
+    tileDiffusion: Optional[bool] = None
+    clipSkip: Optional[int] = None
+    enhanceDetails: Optional[bool] = None
+    realism: Optional[bool] = None
 
     def __post_init__(self):
         if self.sparseStructure is not None and isinstance(self.sparseStructure, dict):
@@ -880,6 +894,18 @@ def request_key(self) -> str:
         return "settings"
 
 
+@dataclass
+class IUpscaleSettings(ISettings):
+
+    def __post_init__(self):
+        super().__post_init__()
+        warnings.warn(
+            "IUpscaleSettings is deprecated and will be removed in a future release; use ISettings for image upscale settings instead.",
+            DeprecationWarning,
+            stacklevel=3,
+        )
+
+
 @dataclass
 class IInputFrame(SerializableMixin):
     image: Union[str, File]
@@ -1241,45 +1267,25 @@ def __hash__(self):
         return hash((self.taskType, self.taskUUID, self.text, self.cost))
 
 
-@dataclass
-class IUpscaleSettings:
-    # Common parameters across all upscaler models
-    steps: Optional[int] = None  # Quality steps (4-60 depending on model)
-    seed: Optional[int] = None  # Reproducibility toggle
-    CFGScale: Optional[float] = None  # Guidance CFG (3-20 depending on model)
-    positivePrompt: Optional[str] = None  
-    negativePrompt: Optional[str] = None  
-    
-    # Clarity upscaler specific
-    controlNetWeight: Optional[float] = None  # Style preservation/Resemblance (0-1)
-    strength: Optional[float] = None  # Creativity (0-1)
-    scheduler: Optional[str] = None  # Controls noise addition/removal
-    
-    # CCSR and Latent upscaler specific
-    colorFix: Optional[bool] = None  # Color correction (ADAIN/NOFIX)
-    tileDiffusion: Optional[bool] = None  # Tile diffusion for large images
-    
-    # Latent upscaler specific
-    clipSkip: Optional[int] = None  # Skip CLIP layers during guidance (0-2)
-
-
 @dataclass
 class IImageUpscale:
-    upscaleFactor: float  # Changed to float to support decimal values like 1.5
+    upscaleFactor: Optional[float] = None  
+    targetMegapixels: Optional[int] = None
     inputImage: Optional[Union[str, File]] = None
     model: Optional[str] = None  # Model AIR ID (runware:500@1, runware:501@1, runware:502@1, runware:503@1)
-    settings: Optional[Union[IUpscaleSettings, Dict[str, Any]]] = None  # Advanced upscaling settings
+    settings: Optional[Union[ISettings, Dict[str, Any]]] = None
     outputType: Optional[IOutputType] = None
     outputFormat: Optional[IOutputFormat] = None
     includeCost: bool = False
     webhookURL: Optional[str] = None
     providerSettings: Optional[ImageProviderSettings] = None
     safety: Optional[Union[ISafety, Dict[str, Any]]] = None
     inputs: Optional[Union[IInputs, Dict[str, Any]]] = None
+    deliveryMethod: str = "sync"
 
     def __post_init__(self):
         if self.settings is not None and isinstance(self.settings, dict):
-            self.settings = IUpscaleSettings(**self.settings)
+            self.settings = ISettings(**self.settings)
         if self.safety is not None and isinstance(self.safety, dict):
             self.safety = ISafety(**self.safety)
         if self.inputs is not None and isinstance(self.inputs, dict):