Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,645 Bytes
ebfc6b3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
import io
import subprocess
from pathlib import Path
import torch
from PIL import ExifTags, Image, ImageCms, ImageOps
from PIL.Image import Image as PilImage
from ltx_trainer import logger
def get_gpu_memory_gb(device: torch.device) -> float:
"""Get current GPU memory usage in GB using nvidia-smi"""
try:
device_id = device.index if device.index is not None else 0
result = subprocess.check_output(
[
"nvidia-smi",
"--query-gpu=memory.used",
"--format=csv,nounits,noheader",
"-i",
str(device_id),
],
encoding="utf-8",
)
return float(result.strip()) / 1024 # Convert MB to GB
except (subprocess.CalledProcessError, FileNotFoundError, ValueError) as e:
logger.error(f"Failed to get GPU memory from nvidia-smi: {e}")
# Fallback to torch
return torch.cuda.memory_allocated(device) / 1024**3
def open_image_as_srgb(image_path: str | Path | io.BytesIO) -> PilImage:
"""
Opens an image file, applies rotation (if it's set in metadata) and converts it
to the sRGB color space respecting the original image color space .
"""
exif_colorspace_srgb = 1
with Image.open(image_path) as img_raw:
img = ImageOps.exif_transpose(img_raw)
input_icc_profile = img.info.get("icc_profile")
# Try to convert to sRGB if the image has ICC profile metadata
srgb_profile = ImageCms.createProfile(colorSpace="sRGB")
if input_icc_profile is not None:
input_profile = ImageCms.ImageCmsProfile(io.BytesIO(input_icc_profile))
srgb_img = ImageCms.profileToProfile(img, input_profile, srgb_profile, outputMode="RGB")
else:
# Try fall back to checking EXIF
exif_data = img.getexif()
if exif_data is not None:
# Assume sRGB if no ICC profile and EXIF has no ColorSpace tag
color_space_value = exif_data.get(ExifTags.Base.ColorSpace.value)
if color_space_value is not None and color_space_value != exif_colorspace_srgb:
raise ValueError(
"Image has colorspace tag in EXIF but it isn't set to sRGB,"
" conversion is not supported."
f" EXIF ColorSpace tag value is {color_space_value}",
)
srgb_img = img.convert("RGB")
# Set sRGB profile in metadata since now the image is assumed to be in sRGB.
srgb_profile_data = ImageCms.ImageCmsProfile(srgb_profile).tobytes()
srgb_img.info["icc_profile"] = srgb_profile_data
return srgb_img
|