Skip to content

Fix concurrent issue of ASTC context #292

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
K0lb3 merged 5 commits into from
Feb 3, 2025
Merged

Fix concurrent issue of ASTC context #292

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
248d37e
fix: type errors in Texture2DConverter
isHarryh Jan 26, 2025
fed2433
fix: incorrect type hint in Texture2DConverter
isHarryh Jan 26, 2025
830a2bd
fix: concurrent issue of astc context
isHarryh Jan 26, 2025
b2bd6bb
fix: type hint use Optional
isHarryh Jan 27, 2025
d39ea26
fix: reduce astc context lock scope
isHarryh Jan 27, 2025
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
73 changes: 40 additions & 33 deletions UnityPy/export/Texture2DConverter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,8 @@
import struct
from copy import copy
from io import BytesIO
from typing import TYPE_CHECKING, Dict, Tuple, Union
from threading import Lock
from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union

import astc_encoder
import texture2ddecoder
Expand Down Expand Up @@ -70,15 +71,6 @@ def image_to_texture2d(
# ASTC
elif target_texture_format.name.startswith("ASTC"):
raw_img = img.tobytes("raw", "RGBA")

block_size = tuple(
map(int, target_texture_format.name.rsplit("_", 1)[1].split("x"))
)

config = astc_encoder.ASTCConfig(
astc_encoder.ASTCProfile.LDR, *block_size, 1, 100
)
context = astc_encoder.ASTCContext(config)
raw_img = astc_encoder.ASTCImage(
astc_encoder.ASTCType.U8, img.width, img.height, 1, raw_img
)
Expand All @@ -88,7 +80,14 @@ def image_to_texture2d(
tex_format = getattr(TF, f"ASTC_RGBA_{block_size[0]}x{block_size[1]}")

swizzle = astc_encoder.ASTCSwizzle.from_str("RGBA")
enc_img = context.compress(raw_img, swizzle)

block_size = tuple(
map(int, target_texture_format.name.rsplit("_", 1)[1].split("x"))
)
context, lock = get_astc_context(block_size)
with lock:
enc_img = context.compress(raw_img, swizzle)

tex_format = target_texture_format
# A
elif target_texture_format == TF.Alpha8:
Expand Down Expand Up @@ -166,7 +165,7 @@ def parse_image_data(
texture_format: Union[int, TextureFormat],
version: tuple,
platform: int,
platform_blob: bytes = None,
platform_blob: Optional[bytes] = None,
flip: bool = True,
) -> Image.Image:
image_data = copy(bytes(image_data))
Expand All @@ -177,7 +176,7 @@ def parse_image_data(

if len(selection) == 0:
raise NotImplementedError(
f"Not implemented texture format: {texture_format.name}"
f"Not implemented texture format: {texture_format}"
Copy link
Owner

@K0lb3 K0lb3 Jan 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, I forgot that texture_format could still be an int.
As the int number isn't that telling imo, it would be better to cast the format to the enum and print it's name if possible, and only show the number if that is not possible.

Copy link
Contributor Author

@isHarryh isHarryh Jan 27, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the int can be cast to the enum, it is very likely to be a known format that has been implemented. In other words, the unimplemented texture format is very likely to be not castable. So showing the int value is enough.

Copy link
Owner

@K0lb3 K0lb3 Feb 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not all of them are implemented atm, but so far I haven't seen any of those in the wild either.
So I guess it's indeed fine to keep it as is as int.

)

if platform == BuildTarget.XBOX360 and texture_format in XBOX_SWAP_FORMATS:
Expand Down Expand Up @@ -230,7 +229,7 @@ def pillow(
mode: str,
codec: str,
args,
swap: tuple = None,
swap: Optional[tuple] = None,
) -> Image.Image:
img = (
Image.frombytes(mode, (width, height), image_data, codec, args)
Expand All @@ -252,32 +251,40 @@ def atc(image_data: bytes, width: int, height: int, alpha: bool) -> Image.Image:
return Image.frombytes("RGBA", (width, height), image_data, "raw", "BGRA")


ASTC_CONTEXTS: Dict[Tuple[int, int], astc_encoder.ASTCContext] = {}


def astc(image_data: bytes, width: int, height: int, block_size: tuple) -> Image.Image:
context = ASTC_CONTEXTS.get(block_size)
if context is None:
config = astc_encoder.ASTCConfig(
astc_encoder.ASTCProfile.LDR,
*block_size,
1,
100,
astc_encoder.ASTCConfigFlags.USE_DECODE_UNORM8,
)
context = ASTC_CONTEXTS[block_size] = astc_encoder.ASTCContext(config)

image = astc_encoder.ASTCImage(astc_encoder.ASTCType.U8, width, height, 1)
texture_size = calculate_astc_compressed_size(width, height, block_size)
if len(image_data) < texture_size:
raise ValueError(f"Invalid ASTC data size: {len(image_data)} < {texture_size}")
context.decompress(
image_data[:texture_size], image, astc_encoder.ASTCSwizzle.from_str("RGBA")
)

context, lock = get_astc_context(block_size)
with lock:
context.decompress(
image_data[:texture_size], image, astc_encoder.ASTCSwizzle.from_str("RGBA")
)

return Image.frombytes("RGBA", (width, height), image.data, "raw", "RGBA")


ASTC_CONTEXTS: Dict[Tuple[int, int], Tuple[astc_encoder.ASTCContext, Lock]] = {}


def get_astc_context(block_size: tuple):
"""Get the ASTC context and its lock using the given `block_size`."""
if block_size not in ASTC_CONTEXTS:
config = astc_encoder.ASTCConfig(
astc_encoder.ASTCProfile.LDR,
*block_size,
block_z=1,
quality=100,
flags=astc_encoder.ASTCConfigFlags.USE_DECODE_UNORM8,
)
context = astc_encoder.ASTCContext(config)
lock = Lock()
ASTC_CONTEXTS[block_size] = (context, lock)
return ASTC_CONTEXTS[block_size]


def calculate_astc_compressed_size(width: int, height: int, block_size: tuple) -> int:
"""Calculate the size of the compressed data for ASTC."""
# calculate the number of blocks
Expand Down Expand Up @@ -327,7 +334,7 @@ def half(
mode: str,
codec: str,
args,
swap: tuple = None,
swap: Optional[tuple] = None,
) -> Image.Image:
# convert half-float to int8
stream = BytesIO(image_data)
Expand Down Expand Up @@ -356,7 +363,7 @@ def rg(
padding = bytes(padding_size)
rgb_data = b"".join(
stream.read(padding_size * 2) + padding
for _ in range(image_data / (2 * padding_size))
for _ in range(len(image_data) // (2 * padding_size))
)
if codec == "RGE":
return half(rgb_data, width, height, mode, "RGB", args)
Expand Down
Loading