diff --git a/runware/base.py b/runware/base.py index b060f80..f9a715d 100644 --- a/runware/base.py +++ b/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, @@ -3100,6 +3140,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 diff --git a/runware/types.py b/runware/types.py index d1fdf3a..2db36b7 100644 --- a/runware/types.py +++ b/runware/types.py @@ -851,6 +851,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): @@ -865,6 +879,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] @@ -1224,34 +1250,13 @@ 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 @@ -1259,10 +1264,11 @@ class IImageUpscale: 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):