Spaces:
Running on Zero
Running on Zero
File size: 2,332 Bytes
b701455 | 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """AutoHDR processor for LightDiffusion-Next.
Applies HDR-like effects to enhance image dynamic range and color.
"""
import logging
from typing import TYPE_CHECKING, Any
import torch
if TYPE_CHECKING:
from src.Core.Context import Context
from src.Core.AbstractModel import AbstractModel
class AutoHDRProcessor:
"""Automatic HDR effect processor.
Wraps src/AutoHDR/ahdr.py as a standardized processor.
"""
@classmethod
def is_enabled(cls, ctx: "Context") -> bool:
"""Check if AutoHDR should be applied."""
return getattr(ctx.generation, "autohdr", True)
@classmethod
def apply(
cls,
image: torch.Tensor,
ctx: "Context" = None,
intensity: float = 1.0,
**kwargs,
) -> torch.Tensor:
"""Apply HDR effects to an image.
Args:
image: Input image tensor
ctx: Optional pipeline context (used for autohdr flag)
intensity: HDR effect intensity
**kwargs: Additional HDR parameters
Returns:
Enhanced image tensor
"""
logger = logging.getLogger(__name__)
try:
from src.AutoHDR import ahdr
hdr = ahdr.HDREffects()
result = hdr.apply_hdr2(image, hdr_intensity=intensity)
# Handle tuple/list return
if isinstance(result, (tuple, list)):
return result[0]
return result
except Exception as e:
logger.warning(f"AutoHDR failed: {e}")
return image
@classmethod
def process(
cls,
ctx: "Context",
model: "AbstractModel" = None,
**kwargs,
) -> "Context":
"""Process context, applying HDR to current_image.
Args:
ctx: Pipeline context with current_image
model: Not used, included for interface compatibility
**kwargs: Additional parameters
Returns:
Context with enhanced current_image
"""
if not cls.is_enabled(ctx):
return ctx
if ctx.current_image is not None:
ctx.current_image = cls.apply(ctx.current_image, ctx)
return ctx
|