mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
feat(backend): rename realesrgan class & upscale method
This commit is contained in:
parent
46a2d83b84
commit
7653d21cf5
@ -11,7 +11,7 @@ from pydantic import ConfigDict
|
|||||||
|
|
||||||
from invokeai.app.invocations.primitives import ImageField, ImageOutput
|
from invokeai.app.invocations.primitives import ImageField, ImageOutput
|
||||||
from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
|
from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
|
||||||
from invokeai.backend.image_util.realesrgan.realesrgan import RealESRGANer
|
from invokeai.backend.image_util.realesrgan.realesrgan import RealESRGAN
|
||||||
from invokeai.backend.util.devices import choose_torch_device
|
from invokeai.backend.util.devices import choose_torch_device
|
||||||
|
|
||||||
from .baseinvocation import BaseInvocation, InputField, InvocationContext, WithMetadata, WithWorkflow, invocation
|
from .baseinvocation import BaseInvocation, InputField, InvocationContext, WithMetadata, WithWorkflow, invocation
|
||||||
@ -92,7 +92,7 @@ class ESRGANInvocation(BaseInvocation, WithWorkflow, WithMetadata):
|
|||||||
|
|
||||||
esrgan_model_path = Path(f"core/upscaling/realesrgan/{self.model_name}")
|
esrgan_model_path = Path(f"core/upscaling/realesrgan/{self.model_name}")
|
||||||
|
|
||||||
upsampler = RealESRGANer(
|
upscaler = RealESRGAN(
|
||||||
scale=netscale,
|
scale=netscale,
|
||||||
model_path=models_path / esrgan_model_path,
|
model_path=models_path / esrgan_model_path,
|
||||||
model=rrdbnet_model,
|
model=rrdbnet_model,
|
||||||
@ -107,7 +107,7 @@ class ESRGANInvocation(BaseInvocation, WithWorkflow, WithMetadata):
|
|||||||
# We can pass an `outscale` value here, but it just resizes the image by that factor after
|
# We can pass an `outscale` value here, but it just resizes the image by that factor after
|
||||||
# upscaling, so it's kinda pointless for our purposes. If you want something other than 4x
|
# upscaling, so it's kinda pointless for our purposes. If you want something other than 4x
|
||||||
# upscaling, you'll need to add a resize node after this one.
|
# upscaling, you'll need to add a resize node after this one.
|
||||||
upscaled_image = upsampler.enhance(cv2_image)
|
upscaled_image = upscaler.upscale(cv2_image)
|
||||||
|
|
||||||
# back to PIL
|
# back to PIL
|
||||||
pil_image = Image.fromarray(cv2.cvtColor(upscaled_image, cv2.COLOR_BGR2RGB)).convert("RGBA")
|
pil_image = Image.fromarray(cv2.cvtColor(upscaled_image, cv2.COLOR_BGR2RGB)).convert("RGBA")
|
||||||
|
@ -32,7 +32,7 @@ class ImageMode(str, Enum):
|
|||||||
RGBA = "RGBA"
|
RGBA = "RGBA"
|
||||||
|
|
||||||
|
|
||||||
class RealESRGANer:
|
class RealESRGAN:
|
||||||
"""A helper class for upsampling images with RealESRGAN.
|
"""A helper class for upsampling images with RealESRGAN.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
@ -202,7 +202,7 @@ class RealESRGANer:
|
|||||||
return self.output
|
return self.output
|
||||||
|
|
||||||
@torch.no_grad()
|
@torch.no_grad()
|
||||||
def enhance(self, img: MatLike, esrgan_alpha_upscale: bool = True) -> npt.NDArray[Any]:
|
def upscale(self, img: MatLike, esrgan_alpha_upscale: bool = True) -> npt.NDArray[Any]:
|
||||||
np_img = img.astype(np.float32)
|
np_img = img.astype(np.float32)
|
||||||
alpha: Optional[np.ndarray] = None
|
alpha: Optional[np.ndarray] = None
|
||||||
if np.max(np_img) > 256:
|
if np.max(np_img) > 256:
|
||||||
|
Loading…
Reference in New Issue
Block a user