58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
import torch
|
|
from .utils import (
|
|
call_api_with_fallback,
|
|
generation_url_to_tensor,
|
|
resolve_api_key,
|
|
upload_image,
|
|
)
|
|
|
|
|
|
class PrunaImageUpscale:
|
|
"""Upscale an image using the Pruna API (p-image-upscale).
|
|
|
|
Pricing: $0.005/image for 1-4 MP target, $0.01/image for 5-8 MP target.
|
|
"""
|
|
|
|
CATEGORY = "pruna ai"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "upscale"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"target_mp": (
|
|
"INT",
|
|
{"default": 4, "min": 1, "max": 8, "step": 1},
|
|
),
|
|
"enhance_details": ("BOOLEAN", {"default": False}),
|
|
"enhance_realism": ("BOOLEAN", {"default": False}),
|
|
"api_key": ("STRING", {"default": ""}),
|
|
}
|
|
}
|
|
|
|
def upscale(
|
|
self,
|
|
image: torch.Tensor,
|
|
target_mp: int,
|
|
enhance_details: bool,
|
|
enhance_realism: bool,
|
|
api_key: str,
|
|
):
|
|
key = resolve_api_key(api_key)
|
|
image_url = upload_image(image, key)
|
|
payload = {
|
|
"input": {
|
|
"image": image_url,
|
|
"target": target_mp,
|
|
"enhance_details": enhance_details,
|
|
"enhance_realism": enhance_realism,
|
|
"output_format": "png",
|
|
}
|
|
}
|
|
data = call_api_with_fallback("p-image-upscale", payload, key)
|
|
tensor = generation_url_to_tensor(data)
|
|
return (tensor,)
|