Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 54 additions & 3 deletions runware/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
56 changes: 31 additions & 25 deletions runware/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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]
Expand Down Expand Up @@ -1224,45 +1250,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):
Expand Down