| # Thanks to https://github.com/Lightricks/ComfyUI-LTXVideo/blob/master/film_grain.py | |
| import torch | |
| def add_film_grain(images: torch.Tensor, grain_intensity: float = 0, saturation: float = 0.5): | |
| device = images.device | |
| input_was_uint8 = images.dtype == torch.uint8 | |
| if input_was_uint8: | |
| images = images.float().div_(255.0).mul_(2.0).sub_(1.0) | |
| images = images.permute(1, 2, 3, 0) | |
| images.add_(1.0).div_(2.0) | |
| grain = torch.randn_like(images, device=device) | |
| grain[:, :, :, 0] *= 2 | |
| grain[:, :, :, 2] *= 3 | |
| grain = grain * saturation + grain[:, :, :, 1].unsqueeze(3).repeat( | |
| 1, 1, 1, 3 | |
| ) * (1 - saturation) | |
| # Blend the grain with the image | |
| noised_images = images + grain_intensity * grain | |
| noised_images.clamp_(0, 1) | |
| noised_images.sub_(0.5).mul_(2.0) | |
| noised_images = noised_images.permute(3, 0, 1, 2) | |
| if input_was_uint8: | |
| noised_images = noised_images.add(1.0).mul(127.5).clamp(0, 255).to(torch.uint8) | |
| return noised_images | |