Spaces:
Sleeping
Sleeping
| import numpy as np | |
| class CustomVRAM: | |
| def __init__(self, global_mem): | |
| self.global_mem = global_mem | |
| self.texture_registry = {} | |
| self.texture_counter = 0 | |
| def load_texture(self, data: np.ndarray, name: str = None) -> str: | |
| if name is None: | |
| name = f"texture_{self.texture_counter}" | |
| self.texture_counter += 1 | |
| # Serialize numpy array to bytes | |
| data_bytes = data.tobytes() | |
| data_shape = data.shape | |
| data_dtype = str(data.dtype) | |
| # Store metadata and data in global memory | |
| # For simplicity, we'll store everything contiguously for now. | |
| # In a real system, this would involve more sophisticated memory management. | |
| # Find a suitable address in global memory (very simplified, no actual allocation logic) | |
| # For this simulation, we'll just use a simple counter for addresses. | |
| # In a real scenario, you'd need a proper memory allocator. | |
| address = self.global_mem.allocate_space(len(data_bytes) + 100) # +100 for metadata | |
| # Store shape, dtype, and then data | |
| # This is a very basic serialization. For production, consider more robust methods. | |
| metadata = f"{data_shape};{data_dtype};{len(data_bytes)}".encode("utf-8") | |
| self.global_mem.write(address, list(metadata)) | |
| self.global_mem.write(address + len(metadata), list(data_bytes)) | |
| self.texture_registry[name] = { | |
| "address": address, | |
| "size": len(data_bytes), | |
| "shape": data_shape, | |
| "dtype": data_dtype, | |
| "metadata_size": len(metadata) | |
| } | |
| return name | |
| def get_texture(self, name: str) -> np.ndarray: | |
| if name not in self.texture_registry: | |
| return None | |
| texture_info = self.texture_registry[name] | |
| address = texture_info["address"] | |
| size = texture_info["size"] | |
| shape = texture_info["shape"] | |
| dtype = texture_info["dtype"] | |
| metadata_size = texture_info["metadata_size"] | |
| # Read data from global memory | |
| data_bytes = bytes(self.global_mem.read(address + metadata_size, size)) | |
| # Deserialize bytes to numpy array | |
| return np.frombuffer(data_bytes, dtype=dtype).reshape(shape) | |
| def has_texture(self, name: str) -> bool: | |
| return name in self.texture_registry | |
| def delete_texture(self, name: str): | |
| if name in self.texture_registry: | |
| # In a real system, you'd deallocate the memory. | |
| # For this simulation, we just remove the entry. | |
| del self.texture_registry[name] | |