Remove unused file and dependencies
This commit is contained in:
wfjsw
@@ -1,27 +0,0 @@
|
||||
FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime
|
||||
|
||||
# Create workdir
|
||||
WORKDIR /invsr
|
||||
|
||||
# Copy from the files to the container
|
||||
COPY . /invsr
|
||||
|
||||
# Install requeriments
|
||||
RUN apt update -y && apt install -y ffmpeg libsm6 libxext6
|
||||
|
||||
RUN pip install -U xformers==0.0.27.post2 --index-url https://download.pytorch.org/whl/cu121
|
||||
RUN pip install -e ".[torch]"
|
||||
RUN pip install -r requirements.txt
|
||||
|
||||
# Clean up all cached files
|
||||
RUN pip cache purge && apt-get clean autoclean && apt-get autoremove --yes && rm -rf /var/lib/{apt,dpkg,cache,log}/
|
||||
|
||||
# Expose gradio port
|
||||
EXPOSE 7860
|
||||
|
||||
# Set listen for 0.0.0.0
|
||||
ENV GRADIO_SERVER_NAME="0.0.0.0"
|
||||
|
||||
# Set python as entrypoint pointing to app.py to run the interface by default
|
||||
ENTRYPOINT [ "python" ]
|
||||
CMD [ "app.py" ]
|
||||
-108
@@ -1,108 +0,0 @@
|
||||
# Arbitrary-steps Image Super-resolution via Diffusion Inversion
|
||||
|
||||
[Zongsheng Yue](https://zsyoaoa.github.io/), [Kang Liao](https://kangliao929.github.io/), [Chen Change Loy](https://www.mmlab-ntu.com/person/ccloy/)
|
||||
|
||||
[](https://arxiv.org/abs/2412.09013) [](https://replicate.com/zsyoaoa/invsr) [](https://huggingface.co/spaces/OAOA/InvSR) <a href="https://colab.research.google.com/drive/1hjgCFnAU4oUUhh9VRfTwsFN1AiIjdcSR?usp=sharing"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google colab logo"></a> 
|
||||
|
||||
<!--[](https://replicate.com/cjwbw/resshift)-->
|
||||
|
||||
|
||||
:star: If you've found InvSR useful for your research or projects, please show your support by starring this repo. Thanks! :hugs:
|
||||
|
||||
---
|
||||
>This study presents a new image super-resolution (SR) technique based on diffusion inversion, aiming at harnessing the rich image priors encapsulated in large pre-trained diffusion models to improve SR performance. We design a \textit{Partial noise Prediction} strategy to construct an intermediate state of the diffusion model, which serves as the starting sampling point. Central to our approach is a deep noise predictor to estimate the optimal noise maps for the forward diffusion process. Once trained, this noise predictor can be used to initialize the sampling process partially along the diffusion trajectory, generating the desirable high-resolution result. Compared to existing approaches, our method offers a flexible and efficient sampling mechanism that supports an arbitrary number of sampling steps, ranging from one to five. Even with a single sampling step, our method demonstrates superior or comparable performance to recent state-of-the-art approaches.
|
||||
><img src="./assets/framework.png" align="middle" width="800">
|
||||
---
|
||||
## Update
|
||||
- **2025.01.08**: Update gradio demo for batch processing.
|
||||
- **2024.12.14**: Add [](https://replicate.com/zsyoaoa/invsr).
|
||||
- **2024.12.13**: Add [](https://huggingface.co/spaces/OAOA/InvSR) and <a href="https://colab.research.google.com/drive/1hjgCFnAU4oUUhh9VRfTwsFN1AiIjdcSR?usp=sharing"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google colab logo"></a>.
|
||||
- **2024.12.11**: Create this repo.
|
||||
|
||||
## Requirements
|
||||
* Python 3.10, Pytorch 2.4.0, [xformers](https://github.com/facebookresearch/xformers) 0.0.27.post2
|
||||
* More detail (See [environment.yaml](environment.yaml))
|
||||
* A suitable [conda](https://conda.io/) environment named `invsr` can be created and activated with:
|
||||
|
||||
```
|
||||
conda create -n invsr python=3.10
|
||||
conda activate invsr
|
||||
pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu121
|
||||
pip install -U xformers==0.0.27.post2 --index-url https://download.pytorch.org/whl/cu121
|
||||
pip install -e ".[torch]"
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
## Applications
|
||||
### :point_right: Real-world Image Super-resolution
|
||||
[<img src="assets/real-7.png" height="235"/>](https://imgsli.com/MzI2MTU5) [<img src="assets/real-1.png" height="235"/>](https://imgsli.com/MzI2MTUx) [<img src="assets/real-2.png" height="235"/>](https://imgsli.com/MzI2MTUy)
|
||||
[<img src="assets/real-4.png" height="361"/>](https://imgsli.com/MzI2MTU0) [<img src="assets/real-6.png" height="361"/>](https://imgsli.com/MzI2MTU3) [<img src="assets/real-5.png" height="361"/>](https://imgsli.com/MzI2MTU1)
|
||||
|
||||
<!--### :point_right: General Image Enhancement-->
|
||||
<!--[<img src="assets/enhance-1.png" height="246.5"/>](https://imgsli.com/MzI2MTYw) [<img src="assets/enhance-2.png" height="246.5"/>](https://imgsli.com/MzI2MTYy) -->
|
||||
<!--[<img src="assets/enhance-3.png" height="207"/>](https://imgsli.com/MzI2MjAx) [<img src="assets/enhance-4.png" height="207"/>](https://imgsli.com/MzI2NTk1) [<img src="assets/enhance-5.png" height="207"/>](https://imgsli.com/MzI2MjA0)-->
|
||||
|
||||
### :point_right: AIGC Image Enhancement
|
||||
[<img src="assets/sdxl-1.png" height="272"/>](https://imgsli.com/MzI2MjQy) [<img src="assets/sdxl-2.png" height="272"/>](https://imgsli.com/MzI2MjQ1) [<img src="assets/sdxl-3.png" height="272"/>](https://imgsli.com/MzI2MjQ3)
|
||||
[<img src="assets/flux-1.png" height="272"/>](https://imgsli.com/MzI2MjQ5) [<img src="assets/flux-2.png" height="272"/>](https://imgsli.com/MzI2MjUw) [<img src="assets/flux-3.png" height="272"/>](https://imgsli.com/MzI2MjUx)
|
||||
|
||||
|
||||
## Inference
|
||||
### :rocket: Fast testing
|
||||
```
|
||||
python inference_invsr.py -i [image folder/image path] -o [result folder] --num_steps 1
|
||||
```
|
||||
1. **To deal with large images, e.g., 1k---->4k, we recommend adding the option** ``--chopping_size 256``.
|
||||
2. Other options:
|
||||
+ Specify the pre-downloaded [SD Turbo](https://huggingface.co/stabilityai/sd-turbo) Model: ``--sd_path``.
|
||||
+ Specify the pre-downloaded noise predictor: ``--started_ckpt_path``.
|
||||
+ The number of sampling steps: ``--num_steps``.
|
||||
+ If your GPU memory is limited, please add the option ``--chopping_bs 1``.
|
||||
|
||||
### :railway_car: Online Demo
|
||||
You can try our method through an online demo:
|
||||
```
|
||||
python app.py
|
||||
```
|
||||
|
||||
### :whale: Now also available in Docker
|
||||
```bash
|
||||
docker compose up -d # Go to http://127.0.0.1:7860/
|
||||
```
|
||||
|
||||
### :airplane: Reproducing our paper results
|
||||
+ Synthetic dataset of ImageNet-Test: [Google Drive](https://drive.google.com/file/d/1PRGrujx3OFilgJ7I6nW7ETIR00wlAl2m/view?usp=sharing).
|
||||
|
||||
+ Real data for image super-resolution: [RealSRV3](https://github.com/csjcai/RealSR) | [RealSet80](testdata/RealSet80)
|
||||
|
||||
+ To reproduce the quantitative results on Imagenet-Test and RealSRV3, please add the color fixing options by ``--color_fix wavelet``.
|
||||
|
||||
## Training
|
||||
### :turtle: Preparing stage
|
||||
1. Download the finetuned LPIPS model from this [link](https://huggingface.co/OAOA/InvSR/resolve/main/vgg16_sdturbo_lpips.pth?download=true) and put it in the folder of "weights".
|
||||
2. Prepare the [config](configs/sd-turbo-sr-ldis.yaml) file:
|
||||
+ SD-Turbo path: configs.sd_pipe.params.cache_dir.
|
||||
+ Training data path: data.train.params.data_source.
|
||||
+ Validation data path: data.val.params.dir_path (low-quality image) and data.val.params.extra_dir_path (high-quality image).
|
||||
+ Batchsize: configs.train.batch and configs.train.microbatch (total batchsize = microbatch * #GPUS * num_grad_accumulation)
|
||||
|
||||
### :dolphin: Begin training
|
||||
```
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --standalone --nproc_per_node=4 --nnodes=1 main.py --save_dir [Logging Folder]
|
||||
```
|
||||
|
||||
### :whale: Resume from interruption
|
||||
```
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --standalone --nproc_per_node=4 --nnodes=1 main.py --save_dir [Logging Folder] --resume save_dir/ckpts/model_xx.pth
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
This project is licensed under [NTU S-Lab License 1.0](LICENSE). Redistribution and use should follow this license.
|
||||
|
||||
## Acknowledgement
|
||||
|
||||
This project is based on [BasicSR](https://github.com/XPixelGroup/BasicSR) and [diffusers](https://github.com/huggingface/diffusers). Thanks for their awesome works.
|
||||
|
||||
### Contact
|
||||
If you have any questions, please feel free to contact me via `zsyzam@gmail.com`.
|
||||
-186
@@ -1,186 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding:utf-8 -*-
|
||||
# Power by Zongsheng Yue 2024-12-11 17:17:41
|
||||
|
||||
import warnings
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
import argparse
|
||||
import numpy as np
|
||||
import gradio as gr
|
||||
from pathlib import Path
|
||||
from omegaconf import OmegaConf
|
||||
from sampler_invsr import InvSamplerSR
|
||||
import os
|
||||
from tqdm import tqdm
|
||||
|
||||
from utils import util_common
|
||||
from utils import util_image
|
||||
from basicsr.utils.download_util import load_file_from_url
|
||||
|
||||
def get_configs(num_steps=1, chopping_size=128, seed=12345):
|
||||
configs = OmegaConf.load("./configs/sample-sd-turbo.yaml")
|
||||
|
||||
if num_steps == 1:
|
||||
configs.timesteps = [200,]
|
||||
elif num_steps == 2:
|
||||
configs.timesteps = [200, 100]
|
||||
elif num_steps == 3:
|
||||
configs.timesteps = [200, 100, 50]
|
||||
elif num_steps == 4:
|
||||
configs.timesteps = [200, 150, 100, 50]
|
||||
elif num_steps == 5:
|
||||
configs.timesteps = [250, 200, 150, 100, 50]
|
||||
else:
|
||||
assert num_steps <= 250
|
||||
configs.timesteps = np.linspace(
|
||||
start=250, stop=0, num=num_steps, endpoint=False, dtype=np.int64()
|
||||
).tolist()
|
||||
print(f'Setting timesteps for inference: {configs.timesteps}')
|
||||
|
||||
configs.sd_path = "./weights"
|
||||
util_common.mkdir(configs.sd_path, delete=False, parents=True)
|
||||
configs.sd_pipe.params.cache_dir = configs.sd_path
|
||||
|
||||
started_ckpt_name = "noise_predictor_sd_turbo_v5.pth"
|
||||
started_ckpt_dir = "./weights"
|
||||
util_common.mkdir(started_ckpt_dir, delete=False, parents=True)
|
||||
started_ckpt_path = Path(started_ckpt_dir) / started_ckpt_name
|
||||
if not started_ckpt_path.exists():
|
||||
load_file_from_url(
|
||||
url="https://huggingface.co/OAOA/InvSR/resolve/main/noise_predictor_sd_turbo_v5.pth",
|
||||
model_dir=started_ckpt_dir,
|
||||
progress=True,
|
||||
file_name=started_ckpt_name,
|
||||
)
|
||||
configs.model_start.ckpt_path = str(started_ckpt_path)
|
||||
|
||||
configs.bs = 1
|
||||
configs.seed = seed
|
||||
configs.basesr.chopping.pch_size = chopping_size
|
||||
configs.basesr.chopping.extra_bs = 4
|
||||
|
||||
return configs
|
||||
|
||||
def predict_single(in_path, num_steps=1, chopping_size=128, seed=12345):
|
||||
configs = get_configs(num_steps=num_steps, chopping_size=chopping_size, seed=seed)
|
||||
sampler = InvSamplerSR(configs)
|
||||
|
||||
out_dir = Path('invsr_output')
|
||||
if not out_dir.exists():
|
||||
out_dir.mkdir()
|
||||
sampler.inference(in_path, out_path=out_dir, bs=1)
|
||||
|
||||
out_path = out_dir / f"{Path(in_path).stem}.png"
|
||||
assert out_path.exists(), 'Super-resolution failed!'
|
||||
im_sr = util_image.imread(out_path, chn="rgb", dtype="uint8")
|
||||
|
||||
return im_sr, str(out_path)
|
||||
|
||||
def process_batch(input_dir, num_steps=1, chopping_size=128, seed=12345, progress=gr.Progress()):
|
||||
input_path = Path(input_dir)
|
||||
output_path = input_path / 'invsr_output'
|
||||
output_path.mkdir(exist_ok=True)
|
||||
|
||||
configs = get_configs(num_steps=num_steps, chopping_size=chopping_size, seed=seed)
|
||||
sampler = InvSamplerSR(configs)
|
||||
|
||||
image_files = list(input_path.glob('*.jpg')) + list(input_path.glob('*.png')) + list(input_path.glob('*.jpeg'))
|
||||
total_files = len(image_files)
|
||||
|
||||
if total_files == 0:
|
||||
return f"No image files found in {input_dir}"
|
||||
|
||||
progress(0, desc="Processing images")
|
||||
for idx, img_path in enumerate(image_files):
|
||||
out_path = output_path / f"{img_path.stem}.png"
|
||||
sampler.inference(str(img_path), out_path=output_path, bs=1)
|
||||
progress((idx + 1)/total_files, desc=f"Processing image {idx + 1}/{total_files}")
|
||||
|
||||
return f"Processed {total_files} images. Results saved in {output_path}"
|
||||
|
||||
title = "Arbitrary-steps Image Super-resolution via Diffusion Inversion"
|
||||
|
||||
article = r"""
|
||||
If you've found InvSR useful for your research or projects, please show your support by ⭐ the <a href='https://github.com/zsyOAOA/InvSR' target='_blank'>Github Repo</a>. Thanks!
|
||||
[](https://github.com/zsyOAOA/InvSR)
|
||||
---
|
||||
If our work is useful for your research, please consider citing:
|
||||
```bibtex
|
||||
@inproceedings{yue2024invsr,
|
||||
title={Arbitrary-steps Image Super-resolution via Diffusion Inversion},
|
||||
author={Yue, Zongsheng and Liao, Kang and Loy, Chen Change},
|
||||
journal={arXiv preprint arXiv:2412.09013},
|
||||
year={2024}
|
||||
}
|
||||
```
|
||||
📋 **License**
|
||||
This project is licensed under <a rel="license" href="https://github.com/zsyOAOA/InvSR/blob/master/LICENSE">S-Lab License 1.0</a>.
|
||||
Redistribution and use for non-commercial purposes should follow this license.
|
||||
📧 **Contact**
|
||||
If you have any questions, please feel free to contact me via <b>zsyzam@gmail.com</b>.
|
||||
|
||||
"""
|
||||
description = r"""
|
||||
<b>Official Gradio demo</b> for <a href='https://github.com/zsyOAOA/InvSR' target='_blank'><b>Arbitrary-steps Image Super-resolution via Diffuion Inversion</b></a>.<br>
|
||||
🔥 InvSR is an image super-resolution method via Diffusion Inversion, supporting arbitrary sampling steps.<br>
|
||||
"""
|
||||
|
||||
with gr.Blocks() as demo:
|
||||
gr.Markdown(f"# {title}")
|
||||
gr.Markdown(description)
|
||||
|
||||
with gr.Tabs():
|
||||
with gr.Tab("Single Image"):
|
||||
with gr.Row():
|
||||
with gr.Column():
|
||||
input_image = gr.Image(type="filepath", label="Input: Low Quality Image")
|
||||
num_steps = gr.Dropdown(
|
||||
choices=[1,2,3,4,5],
|
||||
value=1,
|
||||
label="Number of steps",
|
||||
)
|
||||
chopping_size = gr.Dropdown(
|
||||
choices=[128, 256],
|
||||
value=128,
|
||||
label="Chopping size",
|
||||
)
|
||||
seed = gr.Number(value=12345, precision=0, label="Random seed")
|
||||
process_btn = gr.Button("Process")
|
||||
|
||||
with gr.Column():
|
||||
output_image = gr.Image(type="numpy", label="Output: High Quality Image")
|
||||
output_file = gr.File(label="Download the output")
|
||||
|
||||
process_btn.click(
|
||||
fn=predict_single,
|
||||
inputs=[input_image, num_steps, chopping_size, seed],
|
||||
outputs=[output_image, output_file]
|
||||
)
|
||||
|
||||
with gr.Tab("Batch Processing"):
|
||||
input_dir = gr.Textbox(label="Input Directory Path")
|
||||
batch_num_steps = gr.Dropdown(
|
||||
choices=[1,2,3,4,5],
|
||||
value=1,
|
||||
label="Number of steps",
|
||||
)
|
||||
batch_chopping_size = gr.Dropdown(
|
||||
choices=[128, 256],
|
||||
value=128,
|
||||
label="Chopping size",
|
||||
)
|
||||
batch_seed = gr.Number(value=12345, precision=0, label="Random seed")
|
||||
batch_btn = gr.Button("Process Folder")
|
||||
output_text = gr.Textbox(label="Processing Status")
|
||||
|
||||
batch_btn.click(
|
||||
fn=process_batch,
|
||||
inputs=[input_dir, batch_num_steps, batch_chopping_size, batch_seed],
|
||||
outputs=output_text
|
||||
)
|
||||
|
||||
gr.Markdown(article)
|
||||
|
||||
demo.queue(max_size=5)
|
||||
demo.launch(share=False)
|
||||
@@ -1,4 +0,0 @@
|
||||
# https://github.com/xinntao/BasicSR
|
||||
# flake8: noqa
|
||||
from .data import *
|
||||
from .utils import *
|
||||
@@ -1,101 +0,0 @@
|
||||
import importlib
|
||||
import numpy as np
|
||||
import random
|
||||
import torch
|
||||
import torch.utils.data
|
||||
from copy import deepcopy
|
||||
from functools import partial
|
||||
from os import path as osp
|
||||
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.prefetch_dataloader import PrefetchDataLoader
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils import get_root_logger, scandir
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils.dist_util import get_dist_info
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils.registry import DATASET_REGISTRY
|
||||
|
||||
__all__ = ['build_dataset', 'build_dataloader']
|
||||
|
||||
# automatically scan and import dataset modules for registry
|
||||
# scan all the files under the data folder with '_dataset' in file names
|
||||
data_folder = osp.dirname(osp.abspath(__file__))
|
||||
dataset_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(data_folder) if v.endswith('_dataset.py')]
|
||||
# import all the dataset modules
|
||||
_dataset_modules = [importlib.import_module(f'custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.{file_name}') for file_name in dataset_filenames]
|
||||
|
||||
|
||||
def build_dataset(dataset_opt):
|
||||
"""Build dataset from options.
|
||||
|
||||
Args:
|
||||
dataset_opt (dict): Configuration for dataset. It must contain:
|
||||
name (str): Dataset name.
|
||||
type (str): Dataset type.
|
||||
"""
|
||||
dataset_opt = deepcopy(dataset_opt)
|
||||
dataset = DATASET_REGISTRY.get(dataset_opt['type'])(dataset_opt)
|
||||
logger = get_root_logger()
|
||||
logger.info(f'Dataset [{dataset.__class__.__name__}] - {dataset_opt["name"]} is built.')
|
||||
return dataset
|
||||
|
||||
|
||||
def build_dataloader(dataset, dataset_opt, num_gpu=1, dist=False, sampler=None, seed=None):
|
||||
"""Build dataloader.
|
||||
|
||||
Args:
|
||||
dataset (torch.utils.data.Dataset): Dataset.
|
||||
dataset_opt (dict): Dataset options. It contains the following keys:
|
||||
phase (str): 'train' or 'val'.
|
||||
num_worker_per_gpu (int): Number of workers for each GPU.
|
||||
batch_size_per_gpu (int): Training batch size for each GPU.
|
||||
num_gpu (int): Number of GPUs. Used only in the train phase.
|
||||
Default: 1.
|
||||
dist (bool): Whether in distributed training. Used only in the train
|
||||
phase. Default: False.
|
||||
sampler (torch.utils.data.sampler): Data sampler. Default: None.
|
||||
seed (int | None): Seed. Default: None
|
||||
"""
|
||||
phase = dataset_opt['phase']
|
||||
rank, _ = get_dist_info()
|
||||
if phase == 'train':
|
||||
if dist: # distributed training
|
||||
batch_size = dataset_opt['batch_size_per_gpu']
|
||||
num_workers = dataset_opt['num_worker_per_gpu']
|
||||
else: # non-distributed training
|
||||
multiplier = 1 if num_gpu == 0 else num_gpu
|
||||
batch_size = dataset_opt['batch_size_per_gpu'] * multiplier
|
||||
num_workers = dataset_opt['num_worker_per_gpu'] * multiplier
|
||||
dataloader_args = dict(
|
||||
dataset=dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
num_workers=num_workers,
|
||||
sampler=sampler,
|
||||
drop_last=True)
|
||||
if sampler is None:
|
||||
dataloader_args['shuffle'] = True
|
||||
dataloader_args['worker_init_fn'] = partial(
|
||||
worker_init_fn, num_workers=num_workers, rank=rank, seed=seed) if seed is not None else None
|
||||
elif phase in ['val', 'test']: # validation
|
||||
dataloader_args = dict(dataset=dataset, batch_size=1, shuffle=False, num_workers=0)
|
||||
else:
|
||||
raise ValueError(f"Wrong dataset phase: {phase}. Supported ones are 'train', 'val' and 'test'.")
|
||||
|
||||
dataloader_args['pin_memory'] = dataset_opt.get('pin_memory', False)
|
||||
dataloader_args['persistent_workers'] = dataset_opt.get('persistent_workers', False)
|
||||
|
||||
prefetch_mode = dataset_opt.get('prefetch_mode')
|
||||
if prefetch_mode == 'cpu': # CPUPrefetcher
|
||||
num_prefetch_queue = dataset_opt.get('num_prefetch_queue', 1)
|
||||
logger = get_root_logger()
|
||||
logger.info(f'Use {prefetch_mode} prefetch dataloader: num_prefetch_queue = {num_prefetch_queue}')
|
||||
return PrefetchDataLoader(num_prefetch_queue=num_prefetch_queue, **dataloader_args)
|
||||
else:
|
||||
# prefetch_mode=None: Normal dataloader
|
||||
# prefetch_mode='cuda': dataloader for CUDAPrefetcher
|
||||
return torch.utils.data.DataLoader(**dataloader_args)
|
||||
|
||||
|
||||
def worker_init_fn(worker_id, num_workers, rank, seed):
|
||||
# Set the worker seed to num_workers * rank + worker_id + seed
|
||||
worker_seed = num_workers * rank + worker_id + seed
|
||||
np.random.seed(worker_seed)
|
||||
random.seed(worker_seed)
|
||||
@@ -1,48 +0,0 @@
|
||||
import math
|
||||
import torch
|
||||
from torch.utils.data.sampler import Sampler
|
||||
|
||||
|
||||
class EnlargedSampler(Sampler):
|
||||
"""Sampler that restricts data loading to a subset of the dataset.
|
||||
|
||||
Modified from torch.utils.data.distributed.DistributedSampler
|
||||
Support enlarging the dataset for iteration-based training, for saving
|
||||
time when restart the dataloader after each epoch
|
||||
|
||||
Args:
|
||||
dataset (torch.utils.data.Dataset): Dataset used for sampling.
|
||||
num_replicas (int | None): Number of processes participating in
|
||||
the training. It is usually the world_size.
|
||||
rank (int | None): Rank of the current process within num_replicas.
|
||||
ratio (int): Enlarging ratio. Default: 1.
|
||||
"""
|
||||
|
||||
def __init__(self, dataset, num_replicas, rank, ratio=1):
|
||||
self.dataset = dataset
|
||||
self.num_replicas = num_replicas
|
||||
self.rank = rank
|
||||
self.epoch = 0
|
||||
self.num_samples = math.ceil(len(self.dataset) * ratio / self.num_replicas)
|
||||
self.total_size = self.num_samples * self.num_replicas
|
||||
|
||||
def __iter__(self):
|
||||
# deterministically shuffle based on epoch
|
||||
g = torch.Generator()
|
||||
g.manual_seed(self.epoch)
|
||||
indices = torch.randperm(self.total_size, generator=g).tolist()
|
||||
|
||||
dataset_size = len(self.dataset)
|
||||
indices = [v % dataset_size for v in indices]
|
||||
|
||||
# subsample
|
||||
indices = indices[self.rank:self.total_size:self.num_replicas]
|
||||
assert len(indices) == self.num_samples
|
||||
|
||||
return iter(indices)
|
||||
|
||||
def __len__(self):
|
||||
return self.num_samples
|
||||
|
||||
def set_epoch(self, epoch):
|
||||
self.epoch = epoch
|
||||
@@ -1,315 +0,0 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from os import path as osp
|
||||
from torch.nn import functional as F
|
||||
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.transforms import mod_crop
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils import img2tensor, scandir
|
||||
|
||||
|
||||
def read_img_seq(path, require_mod_crop=False, scale=1, return_imgname=False):
|
||||
"""Read a sequence of images from a given folder path.
|
||||
|
||||
Args:
|
||||
path (list[str] | str): List of image paths or image folder path.
|
||||
require_mod_crop (bool): Require mod crop for each image.
|
||||
Default: False.
|
||||
scale (int): Scale factor for mod_crop. Default: 1.
|
||||
return_imgname(bool): Whether return image names. Default False.
|
||||
|
||||
Returns:
|
||||
Tensor: size (t, c, h, w), RGB, [0, 1].
|
||||
list[str]: Returned image name list.
|
||||
"""
|
||||
if isinstance(path, list):
|
||||
img_paths = path
|
||||
else:
|
||||
img_paths = sorted(list(scandir(path, full_path=True)))
|
||||
imgs = [cv2.imread(v).astype(np.float32) / 255. for v in img_paths]
|
||||
|
||||
if require_mod_crop:
|
||||
imgs = [mod_crop(img, scale) for img in imgs]
|
||||
imgs = img2tensor(imgs, bgr2rgb=True, float32=True)
|
||||
imgs = torch.stack(imgs, dim=0)
|
||||
|
||||
if return_imgname:
|
||||
imgnames = [osp.splitext(osp.basename(path))[0] for path in img_paths]
|
||||
return imgs, imgnames
|
||||
else:
|
||||
return imgs
|
||||
|
||||
|
||||
def generate_frame_indices(crt_idx, max_frame_num, num_frames, padding='reflection'):
|
||||
"""Generate an index list for reading `num_frames` frames from a sequence
|
||||
of images.
|
||||
|
||||
Args:
|
||||
crt_idx (int): Current center index.
|
||||
max_frame_num (int): Max number of the sequence of images (from 1).
|
||||
num_frames (int): Reading num_frames frames.
|
||||
padding (str): Padding mode, one of
|
||||
'replicate' | 'reflection' | 'reflection_circle' | 'circle'
|
||||
Examples: current_idx = 0, num_frames = 5
|
||||
The generated frame indices under different padding mode:
|
||||
replicate: [0, 0, 0, 1, 2]
|
||||
reflection: [2, 1, 0, 1, 2]
|
||||
reflection_circle: [4, 3, 0, 1, 2]
|
||||
circle: [3, 4, 0, 1, 2]
|
||||
|
||||
Returns:
|
||||
list[int]: A list of indices.
|
||||
"""
|
||||
assert num_frames % 2 == 1, 'num_frames should be an odd number.'
|
||||
assert padding in ('replicate', 'reflection', 'reflection_circle', 'circle'), f'Wrong padding mode: {padding}.'
|
||||
|
||||
max_frame_num = max_frame_num - 1 # start from 0
|
||||
num_pad = num_frames // 2
|
||||
|
||||
indices = []
|
||||
for i in range(crt_idx - num_pad, crt_idx + num_pad + 1):
|
||||
if i < 0:
|
||||
if padding == 'replicate':
|
||||
pad_idx = 0
|
||||
elif padding == 'reflection':
|
||||
pad_idx = -i
|
||||
elif padding == 'reflection_circle':
|
||||
pad_idx = crt_idx + num_pad - i
|
||||
else:
|
||||
pad_idx = num_frames + i
|
||||
elif i > max_frame_num:
|
||||
if padding == 'replicate':
|
||||
pad_idx = max_frame_num
|
||||
elif padding == 'reflection':
|
||||
pad_idx = max_frame_num * 2 - i
|
||||
elif padding == 'reflection_circle':
|
||||
pad_idx = (crt_idx - num_pad) - (i - max_frame_num)
|
||||
else:
|
||||
pad_idx = i - num_frames
|
||||
else:
|
||||
pad_idx = i
|
||||
indices.append(pad_idx)
|
||||
return indices
|
||||
|
||||
|
||||
def paired_paths_from_lmdb(folders, keys):
|
||||
"""Generate paired paths from lmdb files.
|
||||
|
||||
Contents of lmdb. Taking the `lq.lmdb` for example, the file structure is:
|
||||
|
||||
::
|
||||
|
||||
lq.lmdb
|
||||
├── data.mdb
|
||||
├── lock.mdb
|
||||
├── meta_info.txt
|
||||
|
||||
The data.mdb and lock.mdb are standard lmdb files and you can refer to
|
||||
https://lmdb.readthedocs.io/en/release/ for more details.
|
||||
|
||||
The meta_info.txt is a specified txt file to record the meta information
|
||||
of our datasets. It will be automatically created when preparing
|
||||
datasets by our provided dataset tools.
|
||||
Each line in the txt file records
|
||||
1)image name (with extension),
|
||||
2)image shape,
|
||||
3)compression level, separated by a white space.
|
||||
Example: `baboon.png (120,125,3) 1`
|
||||
|
||||
We use the image name without extension as the lmdb key.
|
||||
Note that we use the same key for the corresponding lq and gt images.
|
||||
|
||||
Args:
|
||||
folders (list[str]): A list of folder path. The order of list should
|
||||
be [input_folder, gt_folder].
|
||||
keys (list[str]): A list of keys identifying folders. The order should
|
||||
be in consistent with folders, e.g., ['lq', 'gt'].
|
||||
Note that this key is different from lmdb keys.
|
||||
|
||||
Returns:
|
||||
list[str]: Returned path list.
|
||||
"""
|
||||
assert len(folders) == 2, ('The len of folders should be 2 with [input_folder, gt_folder]. '
|
||||
f'But got {len(folders)}')
|
||||
assert len(keys) == 2, f'The len of keys should be 2 with [input_key, gt_key]. But got {len(keys)}'
|
||||
input_folder, gt_folder = folders
|
||||
input_key, gt_key = keys
|
||||
|
||||
if not (input_folder.endswith('.lmdb') and gt_folder.endswith('.lmdb')):
|
||||
raise ValueError(f'{input_key} folder and {gt_key} folder should both in lmdb '
|
||||
f'formats. But received {input_key}: {input_folder}; '
|
||||
f'{gt_key}: {gt_folder}')
|
||||
# ensure that the two meta_info files are the same
|
||||
with open(osp.join(input_folder, 'meta_info.txt')) as fin:
|
||||
input_lmdb_keys = [line.split('.')[0] for line in fin]
|
||||
with open(osp.join(gt_folder, 'meta_info.txt')) as fin:
|
||||
gt_lmdb_keys = [line.split('.')[0] for line in fin]
|
||||
if set(input_lmdb_keys) != set(gt_lmdb_keys):
|
||||
raise ValueError(f'Keys in {input_key}_folder and {gt_key}_folder are different.')
|
||||
else:
|
||||
paths = []
|
||||
for lmdb_key in sorted(input_lmdb_keys):
|
||||
paths.append(dict([(f'{input_key}_path', lmdb_key), (f'{gt_key}_path', lmdb_key)]))
|
||||
return paths
|
||||
|
||||
|
||||
def paired_paths_from_meta_info_file(folders, keys, meta_info_file, filename_tmpl):
|
||||
"""Generate paired paths from an meta information file.
|
||||
|
||||
Each line in the meta information file contains the image names and
|
||||
image shape (usually for gt), separated by a white space.
|
||||
|
||||
Example of an meta information file:
|
||||
```
|
||||
0001_s001.png (480,480,3)
|
||||
0001_s002.png (480,480,3)
|
||||
```
|
||||
|
||||
Args:
|
||||
folders (list[str]): A list of folder path. The order of list should
|
||||
be [input_folder, gt_folder].
|
||||
keys (list[str]): A list of keys identifying folders. The order should
|
||||
be in consistent with folders, e.g., ['lq', 'gt'].
|
||||
meta_info_file (str): Path to the meta information file.
|
||||
filename_tmpl (str): Template for each filename. Note that the
|
||||
template excludes the file extension. Usually the filename_tmpl is
|
||||
for files in the input folder.
|
||||
|
||||
Returns:
|
||||
list[str]: Returned path list.
|
||||
"""
|
||||
assert len(folders) == 2, ('The len of folders should be 2 with [input_folder, gt_folder]. '
|
||||
f'But got {len(folders)}')
|
||||
assert len(keys) == 2, f'The len of keys should be 2 with [input_key, gt_key]. But got {len(keys)}'
|
||||
input_folder, gt_folder = folders
|
||||
input_key, gt_key = keys
|
||||
|
||||
with open(meta_info_file, 'r') as fin:
|
||||
gt_names = [line.strip().split(' ')[0] for line in fin]
|
||||
|
||||
paths = []
|
||||
for gt_name in gt_names:
|
||||
basename, ext = osp.splitext(osp.basename(gt_name))
|
||||
input_name = f'{filename_tmpl.format(basename)}{ext}'
|
||||
input_path = osp.join(input_folder, input_name)
|
||||
gt_path = osp.join(gt_folder, gt_name)
|
||||
paths.append(dict([(f'{input_key}_path', input_path), (f'{gt_key}_path', gt_path)]))
|
||||
return paths
|
||||
|
||||
|
||||
def paired_paths_from_folder(folders, keys, filename_tmpl):
|
||||
"""Generate paired paths from folders.
|
||||
|
||||
Args:
|
||||
folders (list[str]): A list of folder path. The order of list should
|
||||
be [input_folder, gt_folder].
|
||||
keys (list[str]): A list of keys identifying folders. The order should
|
||||
be in consistent with folders, e.g., ['lq', 'gt'].
|
||||
filename_tmpl (str): Template for each filename. Note that the
|
||||
template excludes the file extension. Usually the filename_tmpl is
|
||||
for files in the input folder.
|
||||
|
||||
Returns:
|
||||
list[str]: Returned path list.
|
||||
"""
|
||||
assert len(folders) == 2, ('The len of folders should be 2 with [input_folder, gt_folder]. '
|
||||
f'But got {len(folders)}')
|
||||
assert len(keys) == 2, f'The len of keys should be 2 with [input_key, gt_key]. But got {len(keys)}'
|
||||
input_folder, gt_folder = folders
|
||||
input_key, gt_key = keys
|
||||
|
||||
input_paths = list(scandir(input_folder))
|
||||
gt_paths = list(scandir(gt_folder))
|
||||
assert len(input_paths) == len(gt_paths), (f'{input_key} and {gt_key} datasets have different number of images: '
|
||||
f'{len(input_paths)}, {len(gt_paths)}.')
|
||||
paths = []
|
||||
for gt_path in gt_paths:
|
||||
basename, ext = osp.splitext(osp.basename(gt_path))
|
||||
input_name = f'{filename_tmpl.format(basename)}{ext}'
|
||||
input_path = osp.join(input_folder, input_name)
|
||||
assert input_name in input_paths, f'{input_name} is not in {input_key}_paths.'
|
||||
gt_path = osp.join(gt_folder, gt_path)
|
||||
paths.append(dict([(f'{input_key}_path', input_path), (f'{gt_key}_path', gt_path)]))
|
||||
return paths
|
||||
|
||||
|
||||
def paths_from_folder(folder):
|
||||
"""Generate paths from folder.
|
||||
|
||||
Args:
|
||||
folder (str): Folder path.
|
||||
|
||||
Returns:
|
||||
list[str]: Returned path list.
|
||||
"""
|
||||
|
||||
paths = list(scandir(folder))
|
||||
paths = [osp.join(folder, path) for path in paths]
|
||||
return paths
|
||||
|
||||
|
||||
def paths_from_lmdb(folder):
|
||||
"""Generate paths from lmdb.
|
||||
|
||||
Args:
|
||||
folder (str): Folder path.
|
||||
|
||||
Returns:
|
||||
list[str]: Returned path list.
|
||||
"""
|
||||
if not folder.endswith('.lmdb'):
|
||||
raise ValueError(f'Folder {folder}folder should in lmdb format.')
|
||||
with open(osp.join(folder, 'meta_info.txt')) as fin:
|
||||
paths = [line.split('.')[0] for line in fin]
|
||||
return paths
|
||||
|
||||
|
||||
def generate_gaussian_kernel(kernel_size=13, sigma=1.6):
|
||||
"""Generate Gaussian kernel used in `duf_downsample`.
|
||||
|
||||
Args:
|
||||
kernel_size (int): Kernel size. Default: 13.
|
||||
sigma (float): Sigma of the Gaussian kernel. Default: 1.6.
|
||||
|
||||
Returns:
|
||||
np.array: The Gaussian kernel.
|
||||
"""
|
||||
from scipy.ndimage import filters as filters
|
||||
kernel = np.zeros((kernel_size, kernel_size))
|
||||
# set element at the middle to one, a dirac delta
|
||||
kernel[kernel_size // 2, kernel_size // 2] = 1
|
||||
# gaussian-smooth the dirac, resulting in a gaussian filter
|
||||
return filters.gaussian_filter(kernel, sigma)
|
||||
|
||||
|
||||
def duf_downsample(x, kernel_size=13, scale=4):
|
||||
"""Downsamping with Gaussian kernel used in the DUF official code.
|
||||
|
||||
Args:
|
||||
x (Tensor): Frames to be downsampled, with shape (b, t, c, h, w).
|
||||
kernel_size (int): Kernel size. Default: 13.
|
||||
scale (int): Downsampling factor. Supported scale: (2, 3, 4).
|
||||
Default: 4.
|
||||
|
||||
Returns:
|
||||
Tensor: DUF downsampled frames.
|
||||
"""
|
||||
assert scale in (2, 3, 4), f'Only support scale (2, 3, 4), but got {scale}.'
|
||||
|
||||
squeeze_flag = False
|
||||
if x.ndim == 4:
|
||||
squeeze_flag = True
|
||||
x = x.unsqueeze(0)
|
||||
b, t, c, h, w = x.size()
|
||||
x = x.view(-1, 1, h, w)
|
||||
pad_w, pad_h = kernel_size // 2 + scale * 2, kernel_size // 2 + scale * 2
|
||||
x = F.pad(x, (pad_w, pad_w, pad_h, pad_h), 'reflect')
|
||||
|
||||
gaussian_filter = generate_gaussian_kernel(kernel_size, 0.4 * scale)
|
||||
gaussian_filter = torch.from_numpy(gaussian_filter).type_as(x).unsqueeze(0).unsqueeze(0)
|
||||
x = F.conv2d(x, gaussian_filter, stride=scale)
|
||||
x = x[:, :, 2:-2, 2:-2]
|
||||
x = x.view(b, t, c, x.size(2), x.size(3))
|
||||
if squeeze_flag:
|
||||
x = x.squeeze(0)
|
||||
return x
|
||||
@@ -1,765 +0,0 @@
|
||||
import cv2
|
||||
import math
|
||||
import numpy as np
|
||||
import random
|
||||
import torch
|
||||
from scipy import special
|
||||
from scipy.stats import multivariate_normal
|
||||
# from torchvision.transforms.functional_tensor import rgb_to_grayscale
|
||||
from torchvision.transforms.functional import rgb_to_grayscale
|
||||
|
||||
# -------------------------------------------------------------------- #
|
||||
# --------------------------- blur kernels --------------------------- #
|
||||
# -------------------------------------------------------------------- #
|
||||
|
||||
|
||||
# --------------------------- util functions --------------------------- #
|
||||
def sigma_matrix2(sig_x, sig_y, theta):
|
||||
"""Calculate the rotated sigma matrix (two dimensional matrix).
|
||||
|
||||
Args:
|
||||
sig_x (float):
|
||||
sig_y (float):
|
||||
theta (float): Radian measurement.
|
||||
|
||||
Returns:
|
||||
ndarray: Rotated sigma matrix.
|
||||
"""
|
||||
d_matrix = np.array([[sig_x**2, 0], [0, sig_y**2]])
|
||||
u_matrix = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]])
|
||||
return np.dot(u_matrix, np.dot(d_matrix, u_matrix.T))
|
||||
|
||||
|
||||
def mesh_grid(kernel_size):
|
||||
"""Generate the mesh grid, centering at zero.
|
||||
|
||||
Args:
|
||||
kernel_size (int):
|
||||
|
||||
Returns:
|
||||
xy (ndarray): with the shape (kernel_size, kernel_size, 2)
|
||||
xx (ndarray): with the shape (kernel_size, kernel_size)
|
||||
yy (ndarray): with the shape (kernel_size, kernel_size)
|
||||
"""
|
||||
ax = np.arange(-kernel_size // 2 + 1., kernel_size // 2 + 1.)
|
||||
xx, yy = np.meshgrid(ax, ax)
|
||||
xy = np.hstack((xx.reshape((kernel_size * kernel_size, 1)), yy.reshape(kernel_size * kernel_size,
|
||||
1))).reshape(kernel_size, kernel_size, 2)
|
||||
return xy, xx, yy
|
||||
|
||||
|
||||
def pdf2(sigma_matrix, grid):
|
||||
"""Calculate PDF of the bivariate Gaussian distribution.
|
||||
|
||||
Args:
|
||||
sigma_matrix (ndarray): with the shape (2, 2)
|
||||
grid (ndarray): generated by :func:`mesh_grid`,
|
||||
with the shape (K, K, 2), K is the kernel size.
|
||||
|
||||
Returns:
|
||||
kernel (ndarrray): un-normalized kernel.
|
||||
"""
|
||||
inverse_sigma = np.linalg.inv(sigma_matrix)
|
||||
kernel = np.exp(-0.5 * np.sum(np.dot(grid, inverse_sigma) * grid, 2))
|
||||
return kernel
|
||||
|
||||
|
||||
def cdf2(d_matrix, grid):
|
||||
"""Calculate the CDF of the standard bivariate Gaussian distribution.
|
||||
Used in skewed Gaussian distribution.
|
||||
|
||||
Args:
|
||||
d_matrix (ndarrasy): skew matrix.
|
||||
grid (ndarray): generated by :func:`mesh_grid`,
|
||||
with the shape (K, K, 2), K is the kernel size.
|
||||
|
||||
Returns:
|
||||
cdf (ndarray): skewed cdf.
|
||||
"""
|
||||
rv = multivariate_normal([0, 0], [[1, 0], [0, 1]])
|
||||
grid = np.dot(grid, d_matrix)
|
||||
cdf = rv.cdf(grid)
|
||||
return cdf
|
||||
|
||||
|
||||
def bivariate_Gaussian(kernel_size, sig_x, sig_y, theta, grid=None, isotropic=True):
|
||||
"""Generate a bivariate isotropic or anisotropic Gaussian kernel.
|
||||
|
||||
In the isotropic mode, only `sig_x` is used. `sig_y` and `theta` is ignored.
|
||||
|
||||
Args:
|
||||
kernel_size (int):
|
||||
sig_x (float):
|
||||
sig_y (float):
|
||||
theta (float): Radian measurement.
|
||||
grid (ndarray, optional): generated by :func:`mesh_grid`,
|
||||
with the shape (K, K, 2), K is the kernel size. Default: None
|
||||
isotropic (bool):
|
||||
|
||||
Returns:
|
||||
kernel (ndarray): normalized kernel.
|
||||
"""
|
||||
if grid is None:
|
||||
grid, _, _ = mesh_grid(kernel_size)
|
||||
if isotropic:
|
||||
sigma_matrix = np.array([[sig_x**2, 0], [0, sig_x**2]])
|
||||
else:
|
||||
sigma_matrix = sigma_matrix2(sig_x, sig_y, theta)
|
||||
kernel = pdf2(sigma_matrix, grid)
|
||||
kernel = kernel / np.sum(kernel)
|
||||
return kernel
|
||||
|
||||
|
||||
def bivariate_generalized_Gaussian(kernel_size, sig_x, sig_y, theta, beta, grid=None, isotropic=True):
|
||||
"""Generate a bivariate generalized Gaussian kernel.
|
||||
|
||||
``Paper: Parameter Estimation For Multivariate Generalized Gaussian Distributions``
|
||||
|
||||
In the isotropic mode, only `sig_x` is used. `sig_y` and `theta` is ignored.
|
||||
|
||||
Args:
|
||||
kernel_size (int):
|
||||
sig_x (float):
|
||||
sig_y (float):
|
||||
theta (float): Radian measurement.
|
||||
beta (float): shape parameter, beta = 1 is the normal distribution.
|
||||
grid (ndarray, optional): generated by :func:`mesh_grid`,
|
||||
with the shape (K, K, 2), K is the kernel size. Default: None
|
||||
|
||||
Returns:
|
||||
kernel (ndarray): normalized kernel.
|
||||
"""
|
||||
if grid is None:
|
||||
grid, _, _ = mesh_grid(kernel_size)
|
||||
if isotropic:
|
||||
sigma_matrix = np.array([[sig_x**2, 0], [0, sig_x**2]])
|
||||
else:
|
||||
sigma_matrix = sigma_matrix2(sig_x, sig_y, theta)
|
||||
inverse_sigma = np.linalg.inv(sigma_matrix)
|
||||
kernel = np.exp(-0.5 * np.power(np.sum(np.dot(grid, inverse_sigma) * grid, 2), beta))
|
||||
kernel = kernel / np.sum(kernel)
|
||||
return kernel
|
||||
|
||||
|
||||
def bivariate_plateau(kernel_size, sig_x, sig_y, theta, beta, grid=None, isotropic=True):
|
||||
"""Generate a plateau-like anisotropic kernel.
|
||||
|
||||
1 / (1+x^(beta))
|
||||
|
||||
Reference: https://stats.stackexchange.com/questions/203629/is-there-a-plateau-shaped-distribution
|
||||
|
||||
In the isotropic mode, only `sig_x` is used. `sig_y` and `theta` is ignored.
|
||||
|
||||
Args:
|
||||
kernel_size (int):
|
||||
sig_x (float):
|
||||
sig_y (float):
|
||||
theta (float): Radian measurement.
|
||||
beta (float): shape parameter, beta = 1 is the normal distribution.
|
||||
grid (ndarray, optional): generated by :func:`mesh_grid`,
|
||||
with the shape (K, K, 2), K is the kernel size. Default: None
|
||||
|
||||
Returns:
|
||||
kernel (ndarray): normalized kernel.
|
||||
"""
|
||||
if grid is None:
|
||||
grid, _, _ = mesh_grid(kernel_size)
|
||||
if isotropic:
|
||||
sigma_matrix = np.array([[sig_x**2, 0], [0, sig_x**2]])
|
||||
else:
|
||||
sigma_matrix = sigma_matrix2(sig_x, sig_y, theta)
|
||||
inverse_sigma = np.linalg.inv(sigma_matrix)
|
||||
kernel = np.reciprocal(np.power(np.sum(np.dot(grid, inverse_sigma) * grid, 2), beta) + 1)
|
||||
kernel = kernel / np.sum(kernel)
|
||||
return kernel
|
||||
|
||||
|
||||
def random_bivariate_Gaussian(kernel_size,
|
||||
sigma_x_range,
|
||||
sigma_y_range,
|
||||
rotation_range,
|
||||
noise_range=None,
|
||||
isotropic=True):
|
||||
"""Randomly generate bivariate isotropic or anisotropic Gaussian kernels.
|
||||
|
||||
In the isotropic mode, only `sigma_x_range` is used. `sigma_y_range` and `rotation_range` is ignored.
|
||||
|
||||
Args:
|
||||
kernel_size (int):
|
||||
sigma_x_range (tuple): [0.6, 5]
|
||||
sigma_y_range (tuple): [0.6, 5]
|
||||
rotation range (tuple): [-math.pi, math.pi]
|
||||
noise_range(tuple, optional): multiplicative kernel noise,
|
||||
[0.75, 1.25]. Default: None
|
||||
|
||||
Returns:
|
||||
kernel (ndarray):
|
||||
"""
|
||||
assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
|
||||
assert sigma_x_range[0] < sigma_x_range[1], 'Wrong sigma_x_range.'
|
||||
sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])
|
||||
if isotropic is False:
|
||||
assert sigma_y_range[0] < sigma_y_range[1], 'Wrong sigma_y_range.'
|
||||
assert rotation_range[0] < rotation_range[1], 'Wrong rotation_range.'
|
||||
sigma_y = np.random.uniform(sigma_y_range[0], sigma_y_range[1])
|
||||
rotation = np.random.uniform(rotation_range[0], rotation_range[1])
|
||||
else:
|
||||
sigma_y = sigma_x
|
||||
rotation = 0
|
||||
|
||||
kernel = bivariate_Gaussian(kernel_size, sigma_x, sigma_y, rotation, isotropic=isotropic)
|
||||
|
||||
# add multiplicative noise
|
||||
if noise_range is not None:
|
||||
assert noise_range[0] < noise_range[1], 'Wrong noise range.'
|
||||
noise = np.random.uniform(noise_range[0], noise_range[1], size=kernel.shape)
|
||||
kernel = kernel * noise
|
||||
kernel = kernel / np.sum(kernel)
|
||||
return kernel
|
||||
|
||||
|
||||
def random_bivariate_generalized_Gaussian(kernel_size,
|
||||
sigma_x_range,
|
||||
sigma_y_range,
|
||||
rotation_range,
|
||||
beta_range,
|
||||
noise_range=None,
|
||||
isotropic=True):
|
||||
"""Randomly generate bivariate generalized Gaussian kernels.
|
||||
|
||||
In the isotropic mode, only `sigma_x_range` is used. `sigma_y_range` and `rotation_range` is ignored.
|
||||
|
||||
Args:
|
||||
kernel_size (int):
|
||||
sigma_x_range (tuple): [0.6, 5]
|
||||
sigma_y_range (tuple): [0.6, 5]
|
||||
rotation range (tuple): [-math.pi, math.pi]
|
||||
beta_range (tuple): [0.5, 8]
|
||||
noise_range(tuple, optional): multiplicative kernel noise,
|
||||
[0.75, 1.25]. Default: None
|
||||
|
||||
Returns:
|
||||
kernel (ndarray):
|
||||
"""
|
||||
assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
|
||||
assert sigma_x_range[0] < sigma_x_range[1], 'Wrong sigma_x_range.'
|
||||
sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])
|
||||
if isotropic is False:
|
||||
assert sigma_y_range[0] < sigma_y_range[1], 'Wrong sigma_y_range.'
|
||||
assert rotation_range[0] < rotation_range[1], 'Wrong rotation_range.'
|
||||
sigma_y = np.random.uniform(sigma_y_range[0], sigma_y_range[1])
|
||||
rotation = np.random.uniform(rotation_range[0], rotation_range[1])
|
||||
else:
|
||||
sigma_y = sigma_x
|
||||
rotation = 0
|
||||
|
||||
# assume beta_range[0] < 1 < beta_range[1]
|
||||
if np.random.uniform() < 0.5:
|
||||
beta = np.random.uniform(beta_range[0], 1)
|
||||
else:
|
||||
beta = np.random.uniform(1, beta_range[1])
|
||||
|
||||
kernel = bivariate_generalized_Gaussian(kernel_size, sigma_x, sigma_y, rotation, beta, isotropic=isotropic)
|
||||
|
||||
# add multiplicative noise
|
||||
if noise_range is not None:
|
||||
assert noise_range[0] < noise_range[1], 'Wrong noise range.'
|
||||
noise = np.random.uniform(noise_range[0], noise_range[1], size=kernel.shape)
|
||||
kernel = kernel * noise
|
||||
kernel = kernel / np.sum(kernel)
|
||||
return kernel
|
||||
|
||||
|
||||
def random_bivariate_plateau(kernel_size,
|
||||
sigma_x_range,
|
||||
sigma_y_range,
|
||||
rotation_range,
|
||||
beta_range,
|
||||
noise_range=None,
|
||||
isotropic=True):
|
||||
"""Randomly generate bivariate plateau kernels.
|
||||
|
||||
In the isotropic mode, only `sigma_x_range` is used. `sigma_y_range` and `rotation_range` is ignored.
|
||||
|
||||
Args:
|
||||
kernel_size (int):
|
||||
sigma_x_range (tuple): [0.6, 5]
|
||||
sigma_y_range (tuple): [0.6, 5]
|
||||
rotation range (tuple): [-math.pi/2, math.pi/2]
|
||||
beta_range (tuple): [1, 4]
|
||||
noise_range(tuple, optional): multiplicative kernel noise,
|
||||
[0.75, 1.25]. Default: None
|
||||
|
||||
Returns:
|
||||
kernel (ndarray):
|
||||
"""
|
||||
assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
|
||||
assert sigma_x_range[0] < sigma_x_range[1], 'Wrong sigma_x_range.'
|
||||
sigma_x = np.random.uniform(sigma_x_range[0], sigma_x_range[1])
|
||||
if isotropic is False:
|
||||
assert sigma_y_range[0] < sigma_y_range[1], 'Wrong sigma_y_range.'
|
||||
assert rotation_range[0] < rotation_range[1], 'Wrong rotation_range.'
|
||||
sigma_y = np.random.uniform(sigma_y_range[0], sigma_y_range[1])
|
||||
rotation = np.random.uniform(rotation_range[0], rotation_range[1])
|
||||
else:
|
||||
sigma_y = sigma_x
|
||||
rotation = 0
|
||||
|
||||
# TODO: this may be not proper
|
||||
if np.random.uniform() < 0.5:
|
||||
beta = np.random.uniform(beta_range[0], 1)
|
||||
else:
|
||||
beta = np.random.uniform(1, beta_range[1])
|
||||
|
||||
kernel = bivariate_plateau(kernel_size, sigma_x, sigma_y, rotation, beta, isotropic=isotropic)
|
||||
# add multiplicative noise
|
||||
if noise_range is not None:
|
||||
assert noise_range[0] < noise_range[1], 'Wrong noise range.'
|
||||
noise = np.random.uniform(noise_range[0], noise_range[1], size=kernel.shape)
|
||||
kernel = kernel * noise
|
||||
kernel = kernel / np.sum(kernel)
|
||||
|
||||
return kernel
|
||||
|
||||
|
||||
def random_mixed_kernels(kernel_list,
|
||||
kernel_prob,
|
||||
kernel_size=21,
|
||||
sigma_x_range=(0.6, 5),
|
||||
sigma_y_range=(0.6, 5),
|
||||
rotation_range=(-math.pi, math.pi),
|
||||
betag_range=(0.5, 8),
|
||||
betap_range=(0.5, 8),
|
||||
noise_range=None):
|
||||
"""Randomly generate mixed kernels.
|
||||
|
||||
Args:
|
||||
kernel_list (tuple): a list name of kernel types,
|
||||
support ['iso', 'aniso', 'skew', 'generalized', 'plateau_iso',
|
||||
'plateau_aniso']
|
||||
kernel_prob (tuple): corresponding kernel probability for each
|
||||
kernel type
|
||||
kernel_size (int):
|
||||
sigma_x_range (tuple): [0.6, 5]
|
||||
sigma_y_range (tuple): [0.6, 5]
|
||||
rotation range (tuple): [-math.pi, math.pi]
|
||||
beta_range (tuple): [0.5, 8]
|
||||
noise_range(tuple, optional): multiplicative kernel noise,
|
||||
[0.75, 1.25]. Default: None
|
||||
|
||||
Returns:
|
||||
kernel (ndarray):
|
||||
"""
|
||||
kernel_type = random.choices(kernel_list, kernel_prob)[0]
|
||||
if kernel_type == 'iso':
|
||||
kernel = random_bivariate_Gaussian(
|
||||
kernel_size, sigma_x_range, sigma_y_range, rotation_range, noise_range=noise_range, isotropic=True)
|
||||
elif kernel_type == 'aniso':
|
||||
kernel = random_bivariate_Gaussian(
|
||||
kernel_size, sigma_x_range, sigma_y_range, rotation_range, noise_range=noise_range, isotropic=False)
|
||||
elif kernel_type == 'generalized_iso':
|
||||
kernel = random_bivariate_generalized_Gaussian(
|
||||
kernel_size,
|
||||
sigma_x_range,
|
||||
sigma_y_range,
|
||||
rotation_range,
|
||||
betag_range,
|
||||
noise_range=noise_range,
|
||||
isotropic=True)
|
||||
elif kernel_type == 'generalized_aniso':
|
||||
kernel = random_bivariate_generalized_Gaussian(
|
||||
kernel_size,
|
||||
sigma_x_range,
|
||||
sigma_y_range,
|
||||
rotation_range,
|
||||
betag_range,
|
||||
noise_range=noise_range,
|
||||
isotropic=False)
|
||||
elif kernel_type == 'plateau_iso':
|
||||
kernel = random_bivariate_plateau(
|
||||
kernel_size, sigma_x_range, sigma_y_range, rotation_range, betap_range, noise_range=None, isotropic=True)
|
||||
elif kernel_type == 'plateau_aniso':
|
||||
kernel = random_bivariate_plateau(
|
||||
kernel_size, sigma_x_range, sigma_y_range, rotation_range, betap_range, noise_range=None, isotropic=False)
|
||||
return kernel
|
||||
|
||||
|
||||
np.seterr(divide='ignore', invalid='ignore')
|
||||
|
||||
|
||||
def circular_lowpass_kernel(cutoff, kernel_size, pad_to=0):
|
||||
"""2D sinc filter
|
||||
|
||||
Reference: https://dsp.stackexchange.com/questions/58301/2-d-circularly-symmetric-low-pass-filter
|
||||
|
||||
Args:
|
||||
cutoff (float): cutoff frequency in radians (pi is max)
|
||||
kernel_size (int): horizontal and vertical size, must be odd.
|
||||
pad_to (int): pad kernel size to desired size, must be odd or zero.
|
||||
"""
|
||||
assert kernel_size % 2 == 1, 'Kernel size must be an odd number.'
|
||||
kernel = np.fromfunction(
|
||||
lambda x, y: cutoff * special.j1(cutoff * np.sqrt(
|
||||
(x - (kernel_size - 1) / 2)**2 + (y - (kernel_size - 1) / 2)**2)) / (2 * np.pi * np.sqrt(
|
||||
(x - (kernel_size - 1) / 2)**2 + (y - (kernel_size - 1) / 2)**2)), [kernel_size, kernel_size])
|
||||
kernel[(kernel_size - 1) // 2, (kernel_size - 1) // 2] = cutoff**2 / (4 * np.pi)
|
||||
kernel = kernel / np.sum(kernel)
|
||||
if pad_to > kernel_size:
|
||||
pad_size = (pad_to - kernel_size) // 2
|
||||
kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size)))
|
||||
return kernel
|
||||
|
||||
|
||||
# ------------------------------------------------------------- #
|
||||
# --------------------------- noise --------------------------- #
|
||||
# ------------------------------------------------------------- #
|
||||
|
||||
# ----------------------- Gaussian Noise ----------------------- #
|
||||
|
||||
|
||||
def generate_gaussian_noise(img, sigma=10, gray_noise=False):
|
||||
"""Generate Gaussian noise.
|
||||
|
||||
Args:
|
||||
img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
|
||||
sigma (float): Noise scale (measured in range 255). Default: 10.
|
||||
|
||||
Returns:
|
||||
(Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],
|
||||
float32.
|
||||
"""
|
||||
if gray_noise:
|
||||
noise = np.float32(np.random.randn(*(img.shape[0:2]))) * sigma / 255.
|
||||
noise = np.expand_dims(noise, axis=2).repeat(3, axis=2)
|
||||
else:
|
||||
noise = np.float32(np.random.randn(*(img.shape))) * sigma / 255.
|
||||
return noise
|
||||
|
||||
|
||||
def add_gaussian_noise(img, sigma=10, clip=True, rounds=False, gray_noise=False):
|
||||
"""Add Gaussian noise.
|
||||
|
||||
Args:
|
||||
img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
|
||||
sigma (float): Noise scale (measured in range 255). Default: 10.
|
||||
|
||||
Returns:
|
||||
(Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],
|
||||
float32.
|
||||
"""
|
||||
noise = generate_gaussian_noise(img, sigma, gray_noise)
|
||||
out = img + noise
|
||||
if clip and rounds:
|
||||
out = np.clip((out * 255.0).round(), 0, 255) / 255.
|
||||
elif clip:
|
||||
out = np.clip(out, 0, 1)
|
||||
elif rounds:
|
||||
out = (out * 255.0).round() / 255.
|
||||
return out
|
||||
|
||||
|
||||
def generate_gaussian_noise_pt(img, sigma=10, gray_noise=0):
|
||||
"""Add Gaussian noise (PyTorch version).
|
||||
|
||||
Args:
|
||||
img (Tensor): Shape (b, c, h, w), range[0, 1], float32.
|
||||
scale (float | Tensor): Noise scale. Default: 1.0.
|
||||
|
||||
Returns:
|
||||
(Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],
|
||||
float32.
|
||||
"""
|
||||
b, _, h, w = img.size()
|
||||
if not isinstance(sigma, (float, int)):
|
||||
sigma = sigma.view(img.size(0), 1, 1, 1)
|
||||
if isinstance(gray_noise, (float, int)):
|
||||
cal_gray_noise = gray_noise > 0
|
||||
else:
|
||||
gray_noise = gray_noise.view(b, 1, 1, 1)
|
||||
cal_gray_noise = torch.sum(gray_noise) > 0
|
||||
|
||||
if cal_gray_noise:
|
||||
noise_gray = torch.randn(*img.size()[2:4], dtype=img.dtype, device=img.device) * sigma / 255.
|
||||
noise_gray = noise_gray.view(b, 1, h, w)
|
||||
|
||||
# always calculate color noise
|
||||
noise = torch.randn(*img.size(), dtype=img.dtype, device=img.device) * sigma / 255.
|
||||
|
||||
if cal_gray_noise:
|
||||
noise = noise * (1 - gray_noise) + noise_gray * gray_noise
|
||||
return noise
|
||||
|
||||
|
||||
def add_gaussian_noise_pt(img, sigma=10, gray_noise=0, clip=True, rounds=False):
|
||||
"""Add Gaussian noise (PyTorch version).
|
||||
|
||||
Args:
|
||||
img (Tensor): Shape (b, c, h, w), range[0, 1], float32.
|
||||
scale (float | Tensor): Noise scale. Default: 1.0.
|
||||
|
||||
Returns:
|
||||
(Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],
|
||||
float32.
|
||||
"""
|
||||
noise = generate_gaussian_noise_pt(img, sigma, gray_noise)
|
||||
out = img + noise
|
||||
if clip and rounds:
|
||||
out = torch.clamp((out * 255.0).round(), 0, 255) / 255.
|
||||
elif clip:
|
||||
out = torch.clamp(out, 0, 1)
|
||||
elif rounds:
|
||||
out = (out * 255.0).round() / 255.
|
||||
return out
|
||||
|
||||
|
||||
# ----------------------- Random Gaussian Noise ----------------------- #
|
||||
def random_generate_gaussian_noise(img, sigma_range=(0, 10), gray_prob=0):
|
||||
sigma = np.random.uniform(sigma_range[0], sigma_range[1])
|
||||
if np.random.uniform() < gray_prob:
|
||||
gray_noise = True
|
||||
else:
|
||||
gray_noise = False
|
||||
return generate_gaussian_noise(img, sigma, gray_noise)
|
||||
|
||||
|
||||
def random_add_gaussian_noise(img, sigma_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):
|
||||
noise = random_generate_gaussian_noise(img, sigma_range, gray_prob)
|
||||
out = img + noise
|
||||
if clip and rounds:
|
||||
out = np.clip((out * 255.0).round(), 0, 255) / 255.
|
||||
elif clip:
|
||||
out = np.clip(out, 0, 1)
|
||||
elif rounds:
|
||||
out = (out * 255.0).round() / 255.
|
||||
return out
|
||||
|
||||
|
||||
def random_generate_gaussian_noise_pt(img, sigma_range=(0, 10), gray_prob=0):
|
||||
sigma = torch.rand(
|
||||
img.size(0), dtype=img.dtype, device=img.device) * (sigma_range[1] - sigma_range[0]) + sigma_range[0]
|
||||
gray_noise = torch.rand(img.size(0), dtype=img.dtype, device=img.device)
|
||||
gray_noise = (gray_noise < gray_prob).float()
|
||||
return generate_gaussian_noise_pt(img, sigma, gray_noise)
|
||||
|
||||
|
||||
def random_add_gaussian_noise_pt(img, sigma_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):
|
||||
noise = random_generate_gaussian_noise_pt(img, sigma_range, gray_prob)
|
||||
out = img + noise
|
||||
if clip and rounds:
|
||||
out = torch.clamp((out * 255.0).round(), 0, 255) / 255.
|
||||
elif clip:
|
||||
out = torch.clamp(out, 0, 1)
|
||||
elif rounds:
|
||||
out = (out * 255.0).round() / 255.
|
||||
return out
|
||||
|
||||
|
||||
# ----------------------- Poisson (Shot) Noise ----------------------- #
|
||||
|
||||
|
||||
def generate_poisson_noise(img, scale=1.0, gray_noise=False):
|
||||
"""Generate poisson noise.
|
||||
|
||||
Reference: https://github.com/scikit-image/scikit-image/blob/main/skimage/util/noise.py#L37-L219
|
||||
|
||||
Args:
|
||||
img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
|
||||
scale (float): Noise scale. Default: 1.0.
|
||||
gray_noise (bool): Whether generate gray noise. Default: False.
|
||||
|
||||
Returns:
|
||||
(Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],
|
||||
float32.
|
||||
"""
|
||||
if gray_noise:
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||||
# round and clip image for counting vals correctly
|
||||
img = np.clip((img * 255.0).round(), 0, 255) / 255.
|
||||
vals = len(np.unique(img))
|
||||
vals = 2**np.ceil(np.log2(vals))
|
||||
out = np.float32(np.random.poisson(img * vals) / float(vals))
|
||||
noise = out - img
|
||||
if gray_noise:
|
||||
noise = np.repeat(noise[:, :, np.newaxis], 3, axis=2)
|
||||
return noise * scale
|
||||
|
||||
|
||||
def add_poisson_noise(img, scale=1.0, clip=True, rounds=False, gray_noise=False):
|
||||
"""Add poisson noise.
|
||||
|
||||
Args:
|
||||
img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
|
||||
scale (float): Noise scale. Default: 1.0.
|
||||
gray_noise (bool): Whether generate gray noise. Default: False.
|
||||
|
||||
Returns:
|
||||
(Numpy array): Returned noisy image, shape (h, w, c), range[0, 1],
|
||||
float32.
|
||||
"""
|
||||
noise = generate_poisson_noise(img, scale, gray_noise)
|
||||
out = img + noise
|
||||
if clip and rounds:
|
||||
out = np.clip((out * 255.0).round(), 0, 255) / 255.
|
||||
elif clip:
|
||||
out = np.clip(out, 0, 1)
|
||||
elif rounds:
|
||||
out = (out * 255.0).round() / 255.
|
||||
return out
|
||||
|
||||
|
||||
def generate_poisson_noise_pt(img, scale=1.0, gray_noise=0):
|
||||
"""Generate a batch of poisson noise (PyTorch version)
|
||||
|
||||
Args:
|
||||
img (Tensor): Input image, shape (b, c, h, w), range [0, 1], float32.
|
||||
scale (float | Tensor): Noise scale. Number or Tensor with shape (b).
|
||||
Default: 1.0.
|
||||
gray_noise (float | Tensor): 0-1 number or Tensor with shape (b).
|
||||
0 for False, 1 for True. Default: 0.
|
||||
|
||||
Returns:
|
||||
(Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],
|
||||
float32.
|
||||
"""
|
||||
b, _, h, w = img.size()
|
||||
if isinstance(gray_noise, (float, int)):
|
||||
cal_gray_noise = gray_noise > 0
|
||||
else:
|
||||
gray_noise = gray_noise.view(b, 1, 1, 1)
|
||||
cal_gray_noise = torch.sum(gray_noise) > 0
|
||||
if cal_gray_noise:
|
||||
img_gray = rgb_to_grayscale(img, num_output_channels=1)
|
||||
# round and clip image for counting vals correctly
|
||||
img_gray = torch.clamp((img_gray * 255.0).round(), 0, 255) / 255.
|
||||
# use for-loop to get the unique values for each sample
|
||||
vals_list = [len(torch.unique(img_gray[i, :, :, :])) for i in range(b)]
|
||||
vals_list = [2**np.ceil(np.log2(vals)) for vals in vals_list]
|
||||
vals = img_gray.new_tensor(vals_list).view(b, 1, 1, 1)
|
||||
out = torch.poisson(img_gray * vals) / vals
|
||||
noise_gray = out - img_gray
|
||||
noise_gray = noise_gray.expand(b, 3, h, w)
|
||||
|
||||
# always calculate color noise
|
||||
# round and clip image for counting vals correctly
|
||||
img = torch.clamp((img * 255.0).round(), 0, 255) / 255.
|
||||
# use for-loop to get the unique values for each sample
|
||||
vals_list = [len(torch.unique(img[i, :, :, :])) for i in range(b)]
|
||||
vals_list = [2**np.ceil(np.log2(vals)) for vals in vals_list]
|
||||
vals = img.new_tensor(vals_list).view(b, 1, 1, 1)
|
||||
out = torch.poisson(img * vals) / vals
|
||||
noise = out - img
|
||||
if cal_gray_noise:
|
||||
noise = noise * (1 - gray_noise) + noise_gray * gray_noise
|
||||
if not isinstance(scale, (float, int)):
|
||||
scale = scale.view(b, 1, 1, 1)
|
||||
return noise * scale
|
||||
|
||||
|
||||
def add_poisson_noise_pt(img, scale=1.0, clip=True, rounds=False, gray_noise=0):
|
||||
"""Add poisson noise to a batch of images (PyTorch version).
|
||||
|
||||
Args:
|
||||
img (Tensor): Input image, shape (b, c, h, w), range [0, 1], float32.
|
||||
scale (float | Tensor): Noise scale. Number or Tensor with shape (b).
|
||||
Default: 1.0.
|
||||
gray_noise (float | Tensor): 0-1 number or Tensor with shape (b).
|
||||
0 for False, 1 for True. Default: 0.
|
||||
|
||||
Returns:
|
||||
(Tensor): Returned noisy image, shape (b, c, h, w), range[0, 1],
|
||||
float32.
|
||||
"""
|
||||
noise = generate_poisson_noise_pt(img, scale, gray_noise)
|
||||
out = img + noise
|
||||
if clip and rounds:
|
||||
out = torch.clamp((out * 255.0).round(), 0, 255) / 255.
|
||||
elif clip:
|
||||
out = torch.clamp(out, 0, 1)
|
||||
elif rounds:
|
||||
out = (out * 255.0).round() / 255.
|
||||
return out
|
||||
|
||||
|
||||
# ----------------------- Random Poisson (Shot) Noise ----------------------- #
|
||||
|
||||
|
||||
def random_generate_poisson_noise(img, scale_range=(0, 1.0), gray_prob=0):
|
||||
scale = np.random.uniform(scale_range[0], scale_range[1])
|
||||
if np.random.uniform() < gray_prob:
|
||||
gray_noise = True
|
||||
else:
|
||||
gray_noise = False
|
||||
return generate_poisson_noise(img, scale, gray_noise)
|
||||
|
||||
|
||||
def random_add_poisson_noise(img, scale_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):
|
||||
noise = random_generate_poisson_noise(img, scale_range, gray_prob)
|
||||
out = img + noise
|
||||
if clip and rounds:
|
||||
out = np.clip((out * 255.0).round(), 0, 255) / 255.
|
||||
elif clip:
|
||||
out = np.clip(out, 0, 1)
|
||||
elif rounds:
|
||||
out = (out * 255.0).round() / 255.
|
||||
return out
|
||||
|
||||
|
||||
def random_generate_poisson_noise_pt(img, scale_range=(0, 1.0), gray_prob=0):
|
||||
scale = torch.rand(
|
||||
img.size(0), dtype=img.dtype, device=img.device) * (scale_range[1] - scale_range[0]) + scale_range[0]
|
||||
gray_noise = torch.rand(img.size(0), dtype=img.dtype, device=img.device)
|
||||
gray_noise = (gray_noise < gray_prob).float()
|
||||
return generate_poisson_noise_pt(img, scale, gray_noise)
|
||||
|
||||
|
||||
def random_add_poisson_noise_pt(img, scale_range=(0, 1.0), gray_prob=0, clip=True, rounds=False):
|
||||
noise = random_generate_poisson_noise_pt(img, scale_range, gray_prob)
|
||||
out = img + noise
|
||||
if clip and rounds:
|
||||
out = torch.clamp((out * 255.0).round(), 0, 255) / 255.
|
||||
elif clip:
|
||||
out = torch.clamp(out, 0, 1)
|
||||
elif rounds:
|
||||
out = (out * 255.0).round() / 255.
|
||||
return out
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------ #
|
||||
# --------------------------- JPEG compression --------------------------- #
|
||||
# ------------------------------------------------------------------------ #
|
||||
|
||||
|
||||
def add_jpg_compression(img, quality=90):
|
||||
"""Add JPG compression artifacts.
|
||||
|
||||
Args:
|
||||
img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
|
||||
quality (float): JPG compression quality. 0 for lowest quality, 100 for
|
||||
best quality. Default: 90.
|
||||
|
||||
Returns:
|
||||
(Numpy array): Returned image after JPG, shape (h, w, c), range[0, 1],
|
||||
float32.
|
||||
"""
|
||||
img = np.clip(img, 0, 1)
|
||||
encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), int(quality)]
|
||||
_, encimg = cv2.imencode('.jpg', img * 255., encode_param)
|
||||
img = np.float32(cv2.imdecode(encimg, 1)) / 255.
|
||||
return img
|
||||
|
||||
|
||||
def random_add_jpg_compression(img, quality_range=(90, 100)):
|
||||
"""Randomly add JPG compression artifacts.
|
||||
|
||||
Args:
|
||||
img (Numpy array): Input image, shape (h, w, c), range [0, 1], float32.
|
||||
quality_range (tuple[float] | list[float]): JPG compression quality
|
||||
range. 0 for lowest quality, 100 for best quality.
|
||||
Default: (90, 100).
|
||||
|
||||
Returns:
|
||||
(Numpy array): Returned image after JPG, shape (h, w, c), range[0, 1],
|
||||
float32.
|
||||
"""
|
||||
quality = np.random.uniform(quality_range[0], quality_range[1])
|
||||
return add_jpg_compression(img, quality)
|
||||
@@ -1,80 +0,0 @@
|
||||
import random
|
||||
import time
|
||||
from os import path as osp
|
||||
from torch.utils import data as data
|
||||
from torchvision.transforms.functional import normalize
|
||||
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.transforms import augment
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils import FileClient, get_root_logger, imfrombytes, img2tensor
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils.registry import DATASET_REGISTRY
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class FFHQDataset(data.Dataset):
|
||||
"""FFHQ dataset for StyleGAN.
|
||||
|
||||
Args:
|
||||
opt (dict): Config for train datasets. It contains the following keys:
|
||||
dataroot_gt (str): Data root path for gt.
|
||||
io_backend (dict): IO backend type and other kwarg.
|
||||
mean (list | tuple): Image mean.
|
||||
std (list | tuple): Image std.
|
||||
use_hflip (bool): Whether to horizontally flip.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, opt):
|
||||
super(FFHQDataset, self).__init__()
|
||||
self.opt = opt
|
||||
# file client (io backend)
|
||||
self.file_client = None
|
||||
self.io_backend_opt = opt['io_backend']
|
||||
|
||||
self.gt_folder = opt['dataroot_gt']
|
||||
self.mean = opt['mean']
|
||||
self.std = opt['std']
|
||||
|
||||
if self.io_backend_opt['type'] == 'lmdb':
|
||||
self.io_backend_opt['db_paths'] = self.gt_folder
|
||||
if not self.gt_folder.endswith('.lmdb'):
|
||||
raise ValueError("'dataroot_gt' should end with '.lmdb', but received {self.gt_folder}")
|
||||
with open(osp.join(self.gt_folder, 'meta_info.txt')) as fin:
|
||||
self.paths = [line.split('.')[0] for line in fin]
|
||||
else:
|
||||
# FFHQ has 70000 images in total
|
||||
self.paths = [osp.join(self.gt_folder, f'{v:08d}.png') for v in range(70000)]
|
||||
|
||||
def __getitem__(self, index):
|
||||
if self.file_client is None:
|
||||
self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
|
||||
|
||||
# load gt image
|
||||
gt_path = self.paths[index]
|
||||
# avoid errors caused by high latency in reading files
|
||||
retry = 3
|
||||
while retry > 0:
|
||||
try:
|
||||
img_bytes = self.file_client.get(gt_path)
|
||||
except Exception as e:
|
||||
logger = get_root_logger()
|
||||
logger.warning(f'File client error: {e}, remaining retry times: {retry - 1}')
|
||||
# change another file to read
|
||||
index = random.randint(0, self.__len__())
|
||||
gt_path = self.paths[index]
|
||||
time.sleep(1) # sleep 1s for occasional server congestion
|
||||
else:
|
||||
break
|
||||
finally:
|
||||
retry -= 1
|
||||
img_gt = imfrombytes(img_bytes, float32=True)
|
||||
|
||||
# random horizontal flip
|
||||
img_gt = augment(img_gt, hflip=self.opt['use_hflip'], rotation=False)
|
||||
# BGR to RGB, HWC to CHW, numpy to tensor
|
||||
img_gt = img2tensor(img_gt, bgr2rgb=True, float32=True)
|
||||
# normalize
|
||||
normalize(img_gt, self.mean, self.std, inplace=True)
|
||||
return {'gt': img_gt, 'gt_path': gt_path}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.paths)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,4 +0,0 @@
|
||||
000 100 (720,1280,3)
|
||||
011 100 (720,1280,3)
|
||||
015 100 (720,1280,3)
|
||||
020 100 (720,1280,3)
|
||||
@@ -1,270 +0,0 @@
|
||||
000 100 (720,1280,3)
|
||||
001 100 (720,1280,3)
|
||||
002 100 (720,1280,3)
|
||||
003 100 (720,1280,3)
|
||||
004 100 (720,1280,3)
|
||||
005 100 (720,1280,3)
|
||||
006 100 (720,1280,3)
|
||||
007 100 (720,1280,3)
|
||||
008 100 (720,1280,3)
|
||||
009 100 (720,1280,3)
|
||||
010 100 (720,1280,3)
|
||||
011 100 (720,1280,3)
|
||||
012 100 (720,1280,3)
|
||||
013 100 (720,1280,3)
|
||||
014 100 (720,1280,3)
|
||||
015 100 (720,1280,3)
|
||||
016 100 (720,1280,3)
|
||||
017 100 (720,1280,3)
|
||||
018 100 (720,1280,3)
|
||||
019 100 (720,1280,3)
|
||||
020 100 (720,1280,3)
|
||||
021 100 (720,1280,3)
|
||||
022 100 (720,1280,3)
|
||||
023 100 (720,1280,3)
|
||||
024 100 (720,1280,3)
|
||||
025 100 (720,1280,3)
|
||||
026 100 (720,1280,3)
|
||||
027 100 (720,1280,3)
|
||||
028 100 (720,1280,3)
|
||||
029 100 (720,1280,3)
|
||||
030 100 (720,1280,3)
|
||||
031 100 (720,1280,3)
|
||||
032 100 (720,1280,3)
|
||||
033 100 (720,1280,3)
|
||||
034 100 (720,1280,3)
|
||||
035 100 (720,1280,3)
|
||||
036 100 (720,1280,3)
|
||||
037 100 (720,1280,3)
|
||||
038 100 (720,1280,3)
|
||||
039 100 (720,1280,3)
|
||||
040 100 (720,1280,3)
|
||||
041 100 (720,1280,3)
|
||||
042 100 (720,1280,3)
|
||||
043 100 (720,1280,3)
|
||||
044 100 (720,1280,3)
|
||||
045 100 (720,1280,3)
|
||||
046 100 (720,1280,3)
|
||||
047 100 (720,1280,3)
|
||||
048 100 (720,1280,3)
|
||||
049 100 (720,1280,3)
|
||||
050 100 (720,1280,3)
|
||||
051 100 (720,1280,3)
|
||||
052 100 (720,1280,3)
|
||||
053 100 (720,1280,3)
|
||||
054 100 (720,1280,3)
|
||||
055 100 (720,1280,3)
|
||||
056 100 (720,1280,3)
|
||||
057 100 (720,1280,3)
|
||||
058 100 (720,1280,3)
|
||||
059 100 (720,1280,3)
|
||||
060 100 (720,1280,3)
|
||||
061 100 (720,1280,3)
|
||||
062 100 (720,1280,3)
|
||||
063 100 (720,1280,3)
|
||||
064 100 (720,1280,3)
|
||||
065 100 (720,1280,3)
|
||||
066 100 (720,1280,3)
|
||||
067 100 (720,1280,3)
|
||||
068 100 (720,1280,3)
|
||||
069 100 (720,1280,3)
|
||||
070 100 (720,1280,3)
|
||||
071 100 (720,1280,3)
|
||||
072 100 (720,1280,3)
|
||||
073 100 (720,1280,3)
|
||||
074 100 (720,1280,3)
|
||||
075 100 (720,1280,3)
|
||||
076 100 (720,1280,3)
|
||||
077 100 (720,1280,3)
|
||||
078 100 (720,1280,3)
|
||||
079 100 (720,1280,3)
|
||||
080 100 (720,1280,3)
|
||||
081 100 (720,1280,3)
|
||||
082 100 (720,1280,3)
|
||||
083 100 (720,1280,3)
|
||||
084 100 (720,1280,3)
|
||||
085 100 (720,1280,3)
|
||||
086 100 (720,1280,3)
|
||||
087 100 (720,1280,3)
|
||||
088 100 (720,1280,3)
|
||||
089 100 (720,1280,3)
|
||||
090 100 (720,1280,3)
|
||||
091 100 (720,1280,3)
|
||||
092 100 (720,1280,3)
|
||||
093 100 (720,1280,3)
|
||||
094 100 (720,1280,3)
|
||||
095 100 (720,1280,3)
|
||||
096 100 (720,1280,3)
|
||||
097 100 (720,1280,3)
|
||||
098 100 (720,1280,3)
|
||||
099 100 (720,1280,3)
|
||||
100 100 (720,1280,3)
|
||||
101 100 (720,1280,3)
|
||||
102 100 (720,1280,3)
|
||||
103 100 (720,1280,3)
|
||||
104 100 (720,1280,3)
|
||||
105 100 (720,1280,3)
|
||||
106 100 (720,1280,3)
|
||||
107 100 (720,1280,3)
|
||||
108 100 (720,1280,3)
|
||||
109 100 (720,1280,3)
|
||||
110 100 (720,1280,3)
|
||||
111 100 (720,1280,3)
|
||||
112 100 (720,1280,3)
|
||||
113 100 (720,1280,3)
|
||||
114 100 (720,1280,3)
|
||||
115 100 (720,1280,3)
|
||||
116 100 (720,1280,3)
|
||||
117 100 (720,1280,3)
|
||||
118 100 (720,1280,3)
|
||||
119 100 (720,1280,3)
|
||||
120 100 (720,1280,3)
|
||||
121 100 (720,1280,3)
|
||||
122 100 (720,1280,3)
|
||||
123 100 (720,1280,3)
|
||||
124 100 (720,1280,3)
|
||||
125 100 (720,1280,3)
|
||||
126 100 (720,1280,3)
|
||||
127 100 (720,1280,3)
|
||||
128 100 (720,1280,3)
|
||||
129 100 (720,1280,3)
|
||||
130 100 (720,1280,3)
|
||||
131 100 (720,1280,3)
|
||||
132 100 (720,1280,3)
|
||||
133 100 (720,1280,3)
|
||||
134 100 (720,1280,3)
|
||||
135 100 (720,1280,3)
|
||||
136 100 (720,1280,3)
|
||||
137 100 (720,1280,3)
|
||||
138 100 (720,1280,3)
|
||||
139 100 (720,1280,3)
|
||||
140 100 (720,1280,3)
|
||||
141 100 (720,1280,3)
|
||||
142 100 (720,1280,3)
|
||||
143 100 (720,1280,3)
|
||||
144 100 (720,1280,3)
|
||||
145 100 (720,1280,3)
|
||||
146 100 (720,1280,3)
|
||||
147 100 (720,1280,3)
|
||||
148 100 (720,1280,3)
|
||||
149 100 (720,1280,3)
|
||||
150 100 (720,1280,3)
|
||||
151 100 (720,1280,3)
|
||||
152 100 (720,1280,3)
|
||||
153 100 (720,1280,3)
|
||||
154 100 (720,1280,3)
|
||||
155 100 (720,1280,3)
|
||||
156 100 (720,1280,3)
|
||||
157 100 (720,1280,3)
|
||||
158 100 (720,1280,3)
|
||||
159 100 (720,1280,3)
|
||||
160 100 (720,1280,3)
|
||||
161 100 (720,1280,3)
|
||||
162 100 (720,1280,3)
|
||||
163 100 (720,1280,3)
|
||||
164 100 (720,1280,3)
|
||||
165 100 (720,1280,3)
|
||||
166 100 (720,1280,3)
|
||||
167 100 (720,1280,3)
|
||||
168 100 (720,1280,3)
|
||||
169 100 (720,1280,3)
|
||||
170 100 (720,1280,3)
|
||||
171 100 (720,1280,3)
|
||||
172 100 (720,1280,3)
|
||||
173 100 (720,1280,3)
|
||||
174 100 (720,1280,3)
|
||||
175 100 (720,1280,3)
|
||||
176 100 (720,1280,3)
|
||||
177 100 (720,1280,3)
|
||||
178 100 (720,1280,3)
|
||||
179 100 (720,1280,3)
|
||||
180 100 (720,1280,3)
|
||||
181 100 (720,1280,3)
|
||||
182 100 (720,1280,3)
|
||||
183 100 (720,1280,3)
|
||||
184 100 (720,1280,3)
|
||||
185 100 (720,1280,3)
|
||||
186 100 (720,1280,3)
|
||||
187 100 (720,1280,3)
|
||||
188 100 (720,1280,3)
|
||||
189 100 (720,1280,3)
|
||||
190 100 (720,1280,3)
|
||||
191 100 (720,1280,3)
|
||||
192 100 (720,1280,3)
|
||||
193 100 (720,1280,3)
|
||||
194 100 (720,1280,3)
|
||||
195 100 (720,1280,3)
|
||||
196 100 (720,1280,3)
|
||||
197 100 (720,1280,3)
|
||||
198 100 (720,1280,3)
|
||||
199 100 (720,1280,3)
|
||||
200 100 (720,1280,3)
|
||||
201 100 (720,1280,3)
|
||||
202 100 (720,1280,3)
|
||||
203 100 (720,1280,3)
|
||||
204 100 (720,1280,3)
|
||||
205 100 (720,1280,3)
|
||||
206 100 (720,1280,3)
|
||||
207 100 (720,1280,3)
|
||||
208 100 (720,1280,3)
|
||||
209 100 (720,1280,3)
|
||||
210 100 (720,1280,3)
|
||||
211 100 (720,1280,3)
|
||||
212 100 (720,1280,3)
|
||||
213 100 (720,1280,3)
|
||||
214 100 (720,1280,3)
|
||||
215 100 (720,1280,3)
|
||||
216 100 (720,1280,3)
|
||||
217 100 (720,1280,3)
|
||||
218 100 (720,1280,3)
|
||||
219 100 (720,1280,3)
|
||||
220 100 (720,1280,3)
|
||||
221 100 (720,1280,3)
|
||||
222 100 (720,1280,3)
|
||||
223 100 (720,1280,3)
|
||||
224 100 (720,1280,3)
|
||||
225 100 (720,1280,3)
|
||||
226 100 (720,1280,3)
|
||||
227 100 (720,1280,3)
|
||||
228 100 (720,1280,3)
|
||||
229 100 (720,1280,3)
|
||||
230 100 (720,1280,3)
|
||||
231 100 (720,1280,3)
|
||||
232 100 (720,1280,3)
|
||||
233 100 (720,1280,3)
|
||||
234 100 (720,1280,3)
|
||||
235 100 (720,1280,3)
|
||||
236 100 (720,1280,3)
|
||||
237 100 (720,1280,3)
|
||||
238 100 (720,1280,3)
|
||||
239 100 (720,1280,3)
|
||||
240 100 (720,1280,3)
|
||||
241 100 (720,1280,3)
|
||||
242 100 (720,1280,3)
|
||||
243 100 (720,1280,3)
|
||||
244 100 (720,1280,3)
|
||||
245 100 (720,1280,3)
|
||||
246 100 (720,1280,3)
|
||||
247 100 (720,1280,3)
|
||||
248 100 (720,1280,3)
|
||||
249 100 (720,1280,3)
|
||||
250 100 (720,1280,3)
|
||||
251 100 (720,1280,3)
|
||||
252 100 (720,1280,3)
|
||||
253 100 (720,1280,3)
|
||||
254 100 (720,1280,3)
|
||||
255 100 (720,1280,3)
|
||||
256 100 (720,1280,3)
|
||||
257 100 (720,1280,3)
|
||||
258 100 (720,1280,3)
|
||||
259 100 (720,1280,3)
|
||||
260 100 (720,1280,3)
|
||||
261 100 (720,1280,3)
|
||||
262 100 (720,1280,3)
|
||||
263 100 (720,1280,3)
|
||||
264 100 (720,1280,3)
|
||||
265 100 (720,1280,3)
|
||||
266 100 (720,1280,3)
|
||||
267 100 (720,1280,3)
|
||||
268 100 (720,1280,3)
|
||||
269 100 (720,1280,3)
|
||||
@@ -1,4 +0,0 @@
|
||||
240 100 (720,1280,3)
|
||||
241 100 (720,1280,3)
|
||||
246 100 (720,1280,3)
|
||||
257 100 (720,1280,3)
|
||||
@@ -1,30 +0,0 @@
|
||||
240 100 (720,1280,3)
|
||||
241 100 (720,1280,3)
|
||||
242 100 (720,1280,3)
|
||||
243 100 (720,1280,3)
|
||||
244 100 (720,1280,3)
|
||||
245 100 (720,1280,3)
|
||||
246 100 (720,1280,3)
|
||||
247 100 (720,1280,3)
|
||||
248 100 (720,1280,3)
|
||||
249 100 (720,1280,3)
|
||||
250 100 (720,1280,3)
|
||||
251 100 (720,1280,3)
|
||||
252 100 (720,1280,3)
|
||||
253 100 (720,1280,3)
|
||||
254 100 (720,1280,3)
|
||||
255 100 (720,1280,3)
|
||||
256 100 (720,1280,3)
|
||||
257 100 (720,1280,3)
|
||||
258 100 (720,1280,3)
|
||||
259 100 (720,1280,3)
|
||||
260 100 (720,1280,3)
|
||||
261 100 (720,1280,3)
|
||||
262 100 (720,1280,3)
|
||||
263 100 (720,1280,3)
|
||||
264 100 (720,1280,3)
|
||||
265 100 (720,1280,3)
|
||||
266 100 (720,1280,3)
|
||||
267 100 (720,1280,3)
|
||||
268 100 (720,1280,3)
|
||||
269 100 (720,1280,3)
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -1,106 +0,0 @@
|
||||
from torch.utils import data as data
|
||||
from torchvision.transforms.functional import normalize
|
||||
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.data_util import paired_paths_from_folder, paired_paths_from_lmdb, paired_paths_from_meta_info_file
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.transforms import augment, paired_random_crop
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils import FileClient, bgr2ycbcr, imfrombytes, img2tensor
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils.registry import DATASET_REGISTRY
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class PairedImageDataset(data.Dataset):
|
||||
"""Paired image dataset for image restoration.
|
||||
|
||||
Read LQ (Low Quality, e.g. LR (Low Resolution), blurry, noisy, etc) and GT image pairs.
|
||||
|
||||
There are three modes:
|
||||
|
||||
1. **lmdb**: Use lmdb files. If opt['io_backend'] == lmdb.
|
||||
2. **meta_info_file**: Use meta information file to generate paths. \
|
||||
If opt['io_backend'] != lmdb and opt['meta_info_file'] is not None.
|
||||
3. **folder**: Scan folders to generate paths. The rest.
|
||||
|
||||
Args:
|
||||
opt (dict): Config for train datasets. It contains the following keys:
|
||||
dataroot_gt (str): Data root path for gt.
|
||||
dataroot_lq (str): Data root path for lq.
|
||||
meta_info_file (str): Path for meta information file.
|
||||
io_backend (dict): IO backend type and other kwarg.
|
||||
filename_tmpl (str): Template for each filename. Note that the template excludes the file extension.
|
||||
Default: '{}'.
|
||||
gt_size (int): Cropped patched size for gt patches.
|
||||
use_hflip (bool): Use horizontal flips.
|
||||
use_rot (bool): Use rotation (use vertical flip and transposing h and w for implementation).
|
||||
scale (bool): Scale, which will be added automatically.
|
||||
phase (str): 'train' or 'val'.
|
||||
"""
|
||||
|
||||
def __init__(self, opt):
|
||||
super(PairedImageDataset, self).__init__()
|
||||
self.opt = opt
|
||||
# file client (io backend)
|
||||
self.file_client = None
|
||||
self.io_backend_opt = opt['io_backend']
|
||||
self.mean = opt['mean'] if 'mean' in opt else None
|
||||
self.std = opt['std'] if 'std' in opt else None
|
||||
|
||||
self.gt_folder, self.lq_folder = opt['dataroot_gt'], opt['dataroot_lq']
|
||||
if 'filename_tmpl' in opt:
|
||||
self.filename_tmpl = opt['filename_tmpl']
|
||||
else:
|
||||
self.filename_tmpl = '{}'
|
||||
|
||||
if self.io_backend_opt['type'] == 'lmdb':
|
||||
self.io_backend_opt['db_paths'] = [self.lq_folder, self.gt_folder]
|
||||
self.io_backend_opt['client_keys'] = ['lq', 'gt']
|
||||
self.paths = paired_paths_from_lmdb([self.lq_folder, self.gt_folder], ['lq', 'gt'])
|
||||
elif 'meta_info_file' in self.opt and self.opt['meta_info_file'] is not None:
|
||||
self.paths = paired_paths_from_meta_info_file([self.lq_folder, self.gt_folder], ['lq', 'gt'],
|
||||
self.opt['meta_info_file'], self.filename_tmpl)
|
||||
else:
|
||||
self.paths = paired_paths_from_folder([self.lq_folder, self.gt_folder], ['lq', 'gt'], self.filename_tmpl)
|
||||
|
||||
def __getitem__(self, index):
|
||||
if self.file_client is None:
|
||||
self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
|
||||
|
||||
scale = self.opt['scale']
|
||||
|
||||
# Load gt and lq images. Dimension order: HWC; channel order: BGR;
|
||||
# image range: [0, 1], float32.
|
||||
gt_path = self.paths[index]['gt_path']
|
||||
img_bytes = self.file_client.get(gt_path, 'gt')
|
||||
img_gt = imfrombytes(img_bytes, float32=True)
|
||||
lq_path = self.paths[index]['lq_path']
|
||||
img_bytes = self.file_client.get(lq_path, 'lq')
|
||||
img_lq = imfrombytes(img_bytes, float32=True)
|
||||
|
||||
# augmentation for training
|
||||
if self.opt['phase'] == 'train':
|
||||
gt_size = self.opt['gt_size']
|
||||
# random crop
|
||||
img_gt, img_lq = paired_random_crop(img_gt, img_lq, gt_size, scale, gt_path)
|
||||
# flip, rotation
|
||||
img_gt, img_lq = augment([img_gt, img_lq], self.opt['use_hflip'], self.opt['use_rot'])
|
||||
|
||||
# color space transform
|
||||
if 'color' in self.opt and self.opt['color'] == 'y':
|
||||
img_gt = bgr2ycbcr(img_gt, y_only=True)[..., None]
|
||||
img_lq = bgr2ycbcr(img_lq, y_only=True)[..., None]
|
||||
|
||||
# crop the unmatched GT images during validation or testing, especially for SR benchmark datasets
|
||||
# TODO: It is better to update the datasets, rather than force to crop
|
||||
if self.opt['phase'] != 'train':
|
||||
img_gt = img_gt[0:img_lq.shape[0] * scale, 0:img_lq.shape[1] * scale, :]
|
||||
|
||||
# BGR to RGB, HWC to CHW, numpy to tensor
|
||||
img_gt, img_lq = img2tensor([img_gt, img_lq], bgr2rgb=True, float32=True)
|
||||
# normalize
|
||||
if self.mean is not None or self.std is not None:
|
||||
normalize(img_lq, self.mean, self.std, inplace=True)
|
||||
normalize(img_gt, self.mean, self.std, inplace=True)
|
||||
|
||||
return {'lq': img_lq, 'gt': img_gt, 'lq_path': lq_path, 'gt_path': gt_path}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.paths)
|
||||
@@ -1,122 +0,0 @@
|
||||
import queue as Queue
|
||||
import threading
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
|
||||
class PrefetchGenerator(threading.Thread):
|
||||
"""A general prefetch generator.
|
||||
|
||||
Reference: https://stackoverflow.com/questions/7323664/python-generator-pre-fetch
|
||||
|
||||
Args:
|
||||
generator: Python generator.
|
||||
num_prefetch_queue (int): Number of prefetch queue.
|
||||
"""
|
||||
|
||||
def __init__(self, generator, num_prefetch_queue):
|
||||
threading.Thread.__init__(self)
|
||||
self.queue = Queue.Queue(num_prefetch_queue)
|
||||
self.generator = generator
|
||||
self.daemon = True
|
||||
self.start()
|
||||
|
||||
def run(self):
|
||||
for item in self.generator:
|
||||
self.queue.put(item)
|
||||
self.queue.put(None)
|
||||
|
||||
def __next__(self):
|
||||
next_item = self.queue.get()
|
||||
if next_item is None:
|
||||
raise StopIteration
|
||||
return next_item
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
|
||||
class PrefetchDataLoader(DataLoader):
|
||||
"""Prefetch version of dataloader.
|
||||
|
||||
Reference: https://github.com/IgorSusmelj/pytorch-styleguide/issues/5#
|
||||
|
||||
TODO:
|
||||
Need to test on single gpu and ddp (multi-gpu). There is a known issue in
|
||||
ddp.
|
||||
|
||||
Args:
|
||||
num_prefetch_queue (int): Number of prefetch queue.
|
||||
kwargs (dict): Other arguments for dataloader.
|
||||
"""
|
||||
|
||||
def __init__(self, num_prefetch_queue, **kwargs):
|
||||
self.num_prefetch_queue = num_prefetch_queue
|
||||
super(PrefetchDataLoader, self).__init__(**kwargs)
|
||||
|
||||
def __iter__(self):
|
||||
return PrefetchGenerator(super().__iter__(), self.num_prefetch_queue)
|
||||
|
||||
|
||||
class CPUPrefetcher():
|
||||
"""CPU prefetcher.
|
||||
|
||||
Args:
|
||||
loader: Dataloader.
|
||||
"""
|
||||
|
||||
def __init__(self, loader):
|
||||
self.ori_loader = loader
|
||||
self.loader = iter(loader)
|
||||
|
||||
def next(self):
|
||||
try:
|
||||
return next(self.loader)
|
||||
except StopIteration:
|
||||
return None
|
||||
|
||||
def reset(self):
|
||||
self.loader = iter(self.ori_loader)
|
||||
|
||||
|
||||
class CUDAPrefetcher():
|
||||
"""CUDA prefetcher.
|
||||
|
||||
Reference: https://github.com/NVIDIA/apex/issues/304#
|
||||
|
||||
It may consume more GPU memory.
|
||||
|
||||
Args:
|
||||
loader: Dataloader.
|
||||
opt (dict): Options.
|
||||
"""
|
||||
|
||||
def __init__(self, loader, opt):
|
||||
self.ori_loader = loader
|
||||
self.loader = iter(loader)
|
||||
self.opt = opt
|
||||
self.stream = torch.cuda.Stream()
|
||||
self.device = torch.device('cuda' if opt['num_gpu'] != 0 else 'cpu')
|
||||
self.preload()
|
||||
|
||||
def preload(self):
|
||||
try:
|
||||
self.batch = next(self.loader) # self.batch is a dict
|
||||
except StopIteration:
|
||||
self.batch = None
|
||||
return None
|
||||
# put tensors to gpu
|
||||
with torch.cuda.stream(self.stream):
|
||||
for k, v in self.batch.items():
|
||||
if torch.is_tensor(v):
|
||||
self.batch[k] = self.batch[k].to(device=self.device, non_blocking=True)
|
||||
|
||||
def next(self):
|
||||
torch.cuda.current_stream().wait_stream(self.stream)
|
||||
batch = self.batch
|
||||
self.preload()
|
||||
return batch
|
||||
|
||||
def reset(self):
|
||||
self.loader = iter(self.ori_loader)
|
||||
self.preload()
|
||||
@@ -1,384 +0,0 @@
|
||||
import cv2
|
||||
import math
|
||||
import numpy as np
|
||||
import os
|
||||
import os.path as osp
|
||||
import random
|
||||
import time
|
||||
import torch
|
||||
from pathlib import Path
|
||||
|
||||
import albumentations
|
||||
|
||||
import torch.nn.functional as F
|
||||
from torch.utils import data as data
|
||||
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils import DiffJPEG
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.degradations import circular_lowpass_kernel, random_mixed_kernels
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.transforms import augment
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils import FileClient, get_root_logger, imfrombytes, img2tensor
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils.registry import DATASET_REGISTRY
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils.img_process_util import filter2D
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.transforms import paired_random_crop, random_crop
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.degradations import random_add_gaussian_noise_pt, random_add_poisson_noise_pt
|
||||
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.utils import util_image
|
||||
|
||||
def readline_txt(txt_file):
|
||||
txt_file = [txt_file, ] if isinstance(txt_file, str) else txt_file
|
||||
out = []
|
||||
for txt_file_current in txt_file:
|
||||
with open(txt_file_current, 'r') as ff:
|
||||
out.extend([x[:-1] for x in ff.readlines()])
|
||||
|
||||
return out
|
||||
|
||||
@DATASET_REGISTRY.register(suffix='basicsr')
|
||||
class RealESRGANDataset(data.Dataset):
|
||||
"""Dataset used for Real-ESRGAN model:
|
||||
Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
|
||||
|
||||
It loads gt (Ground-Truth) images, and augments them.
|
||||
It also generates blur kernels and sinc kernels for generating low-quality images.
|
||||
Note that the low-quality images are processed in tensors on GPUS for faster processing.
|
||||
|
||||
Args:
|
||||
opt (dict): Config for train datasets. It contains the following keys:
|
||||
dataroot_gt (str): Data root path for gt.
|
||||
meta_info (str): Path for meta information file.
|
||||
io_backend (dict): IO backend type and other kwarg.
|
||||
use_hflip (bool): Use horizontal flips.
|
||||
use_rot (bool): Use rotation (use vertical flip and transposing h and w for implementation).
|
||||
Please see more options in the codes.
|
||||
"""
|
||||
|
||||
def __init__(self, opt, mode='training'):
|
||||
super(RealESRGANDataset, self).__init__()
|
||||
self.opt = opt
|
||||
self.file_client = None
|
||||
self.io_backend_opt = opt['io_backend']
|
||||
|
||||
# file client (lmdb io backend)
|
||||
self.image_paths = []
|
||||
self.text_paths = []
|
||||
self.moment_paths = []
|
||||
if opt.get('data_source', None) is not None:
|
||||
for ii in range(len(opt['data_source'])):
|
||||
configs = opt['data_source'].get(f'source{ii+1}')
|
||||
root_path = Path(configs.root_path)
|
||||
im_folder = root_path / configs.image_path
|
||||
im_ext = configs.im_ext
|
||||
image_stems = sorted([x.stem for x in im_folder.glob(f"*.{im_ext}")])
|
||||
if configs.get('length', None) is not None:
|
||||
assert configs.length < len(image_stems)
|
||||
image_stems = image_stems[:configs.length]
|
||||
|
||||
if configs.get("text_path", None) is not None:
|
||||
text_folder = root_path / configs.text_path
|
||||
text_stems = [x.stem for x in text_folder.glob("*.txt")]
|
||||
image_stems = sorted(list(set(image_stems).intersection(set(text_stems))))
|
||||
self.text_paths.extend([str(text_folder / f"{x}.txt") for x in image_stems])
|
||||
else:
|
||||
self.text_paths.extend([None, ] * len(image_stems))
|
||||
|
||||
self.image_paths.extend([str(im_folder / f"{x}.{im_ext}") for x in image_stems])
|
||||
|
||||
if configs.get("moment_path", None) is not None:
|
||||
moment_folder = root_path / configs.moment_path
|
||||
self.moment_paths.extend([str(moment_folder / f"{x}.npy") for x in image_stems])
|
||||
else:
|
||||
self.moment_paths.extend([None, ] * len(image_stems))
|
||||
|
||||
# blur settings for the first degradation
|
||||
self.blur_kernel_size = opt['blur_kernel_size']
|
||||
self.kernel_list = opt['kernel_list']
|
||||
self.kernel_prob = opt['kernel_prob'] # a list for each kernel probability
|
||||
self.blur_sigma = opt['blur_sigma']
|
||||
self.betag_range = opt['betag_range'] # betag used in generalized Gaussian blur kernels
|
||||
self.betap_range = opt['betap_range'] # betap used in plateau blur kernels
|
||||
self.sinc_prob = opt['sinc_prob'] # the probability for sinc filters
|
||||
|
||||
# blur settings for the second degradation
|
||||
self.blur_kernel_size2 = opt['blur_kernel_size2']
|
||||
self.kernel_list2 = opt['kernel_list2']
|
||||
self.kernel_prob2 = opt['kernel_prob2']
|
||||
self.blur_sigma2 = opt['blur_sigma2']
|
||||
self.betag_range2 = opt['betag_range2']
|
||||
self.betap_range2 = opt['betap_range2']
|
||||
self.sinc_prob2 = opt['sinc_prob2']
|
||||
|
||||
# a final sinc filter
|
||||
self.final_sinc_prob = opt['final_sinc_prob']
|
||||
|
||||
self.kernel_range1 = [x for x in range(3, opt['blur_kernel_size'], 2)] # kernel size ranges from 7 to 21
|
||||
self.kernel_range2 = [x for x in range(3, opt['blur_kernel_size2'], 2)] # kernel size ranges from 7 to 21
|
||||
# TODO: kernel range is now hard-coded, should be in the configure file
|
||||
# convolving with pulse tensor brings no blurry effect
|
||||
self.pulse_tensor = torch.zeros(opt['blur_kernel_size2'], opt['blur_kernel_size2']).float()
|
||||
self.pulse_tensor[opt['blur_kernel_size2']//2, opt['blur_kernel_size2']//2] = 1
|
||||
|
||||
self.mode = mode
|
||||
|
||||
def __getitem__(self, index):
|
||||
if self.file_client is None:
|
||||
self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
|
||||
|
||||
# -------------------------------- Load gt images -------------------------------- #
|
||||
# Shape: (h, w, c); channel order: BGR; image range: [0, 1], float32.
|
||||
gt_path = self.image_paths[index]
|
||||
# avoid errors caused by high latency in reading files
|
||||
retry = 3
|
||||
while retry > 0:
|
||||
try:
|
||||
img_bytes = self.file_client.get(gt_path, 'gt')
|
||||
img_gt = imfrombytes(img_bytes, float32=True)
|
||||
except:
|
||||
index = random.randint(0, self.__len__())
|
||||
gt_path = self.image_paths[index]
|
||||
time.sleep(1) # sleep 1s for occasional server congestion
|
||||
finally:
|
||||
retry -= 1
|
||||
if self.mode == 'testing':
|
||||
if not hasattr(self, 'test_aug'):
|
||||
self.test_aug = albumentations.Compose([
|
||||
albumentations.SmallestMaxSize(
|
||||
max_size=self.opt['gt_size'],
|
||||
interpolation=cv2.INTER_AREA,
|
||||
),
|
||||
albumentations.CenterCrop(self.opt['gt_size'], self.opt['gt_size']),
|
||||
])
|
||||
img_gt = self.test_aug(image=img_gt)['image']
|
||||
elif self.mode == 'training':
|
||||
# -------------------- Do augmentation for training: flip, rotation -------------------- #
|
||||
if self.opt['use_hflip'] or self.opt['use_rot']:
|
||||
img_gt = augment(img_gt, self.opt['use_hflip'], self.opt['use_rot'])
|
||||
|
||||
h, w = img_gt.shape[0:2]
|
||||
gt_size = self.opt['gt_size']
|
||||
|
||||
# resize or pad
|
||||
if not self.opt['random_crop']:
|
||||
if not min(h, w) == gt_size:
|
||||
if not hasattr(self, 'smallest_resizer'):
|
||||
self.smallest_resizer = util_image.SmallestMaxSize(
|
||||
max_size=gt_size, pass_resize=False,
|
||||
)
|
||||
img_gt = self.smallest_resizer(img_gt)
|
||||
|
||||
# center crop
|
||||
if not hasattr(self, 'center_cropper'):
|
||||
self.center_cropper = albumentations.CenterCrop(gt_size, gt_size)
|
||||
img_gt = self.center_cropper(image=img_gt)['image']
|
||||
else:
|
||||
img_gt = random_crop(img_gt, self.opt['gt_size'])
|
||||
else:
|
||||
raise ValueError(f'Unexpected value {self.mode} for mode parameter')
|
||||
|
||||
# ------------------------ Generate kernels (used in the first degradation) ------------------------ #
|
||||
kernel_size = random.choice(self.kernel_range1)
|
||||
if np.random.uniform() < self.opt['sinc_prob']:
|
||||
# this sinc filter setting is for kernels ranging from [7, 21]
|
||||
if kernel_size < 13:
|
||||
omega_c = np.random.uniform(np.pi / 3, np.pi)
|
||||
else:
|
||||
omega_c = np.random.uniform(np.pi / 5, np.pi)
|
||||
kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False)
|
||||
else:
|
||||
kernel = random_mixed_kernels(
|
||||
self.kernel_list,
|
||||
self.kernel_prob,
|
||||
kernel_size,
|
||||
self.blur_sigma,
|
||||
self.blur_sigma, [-math.pi, math.pi],
|
||||
self.betag_range,
|
||||
self.betap_range,
|
||||
noise_range=None)
|
||||
# pad kernel
|
||||
pad_size = (self.blur_kernel_size - kernel_size) // 2
|
||||
kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size)))
|
||||
|
||||
# ------------------------ Generate kernels (used in the second degradation) ------------------------ #
|
||||
kernel_size = random.choice(self.kernel_range2)
|
||||
if np.random.uniform() < self.opt['sinc_prob2']:
|
||||
if kernel_size < 13:
|
||||
omega_c = np.random.uniform(np.pi / 3, np.pi)
|
||||
else:
|
||||
omega_c = np.random.uniform(np.pi / 5, np.pi)
|
||||
kernel2 = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False)
|
||||
else:
|
||||
kernel2 = random_mixed_kernels(
|
||||
self.kernel_list2,
|
||||
self.kernel_prob2,
|
||||
kernel_size,
|
||||
self.blur_sigma2,
|
||||
self.blur_sigma2, [-math.pi, math.pi],
|
||||
self.betag_range2,
|
||||
self.betap_range2,
|
||||
noise_range=None)
|
||||
|
||||
# pad kernel
|
||||
pad_size = (self.blur_kernel_size2 - kernel_size) // 2
|
||||
kernel2 = np.pad(kernel2, ((pad_size, pad_size), (pad_size, pad_size)))
|
||||
|
||||
# ------------------------------------- the final sinc kernel ------------------------------------- #
|
||||
if np.random.uniform() < self.opt['final_sinc_prob']:
|
||||
kernel_size = random.choice(self.kernel_range2)
|
||||
omega_c = np.random.uniform(np.pi / 3, np.pi)
|
||||
sinc_kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=self.blur_kernel_size2)
|
||||
sinc_kernel = torch.FloatTensor(sinc_kernel)
|
||||
else:
|
||||
sinc_kernel = self.pulse_tensor
|
||||
|
||||
# BGR to RGB, HWC to CHW, numpy to tensor
|
||||
img_gt = img2tensor([img_gt], bgr2rgb=True, float32=True)[0]
|
||||
kernel = torch.FloatTensor(kernel)
|
||||
kernel2 = torch.FloatTensor(kernel2)
|
||||
|
||||
if self.text_paths[index] is None or self.opt['random_crop']:
|
||||
prompt = ""
|
||||
else:
|
||||
with open(self.text_paths[index], 'r') as ff:
|
||||
prompt = ff.read()
|
||||
if self.opt.max_token_length is not None:
|
||||
prompt = prompt[:self.opt.max_token_length]
|
||||
|
||||
return_d = {
|
||||
'gt': img_gt,
|
||||
'gt_path': gt_path,
|
||||
'txt': prompt,
|
||||
'kernel1': kernel,
|
||||
'kernel2': kernel2,
|
||||
'sinc_kernel': sinc_kernel,
|
||||
}
|
||||
if self.moment_paths[index] is not None and (not self.opt['random_crop']):
|
||||
return_d['gt_moment'] = np.load(self.moment_paths[index])
|
||||
|
||||
return return_d
|
||||
|
||||
def __len__(self):
|
||||
return len(self.image_paths)
|
||||
|
||||
def degrade_fun(self, conf_degradation, im_gt, kernel1, kernel2, sinc_kernel):
|
||||
if not hasattr(self, 'jpeger'):
|
||||
self.jpeger = DiffJPEG(differentiable=False) # simulate JPEG compression artifacts
|
||||
|
||||
ori_h, ori_w = im_gt.size()[2:4]
|
||||
sf = conf_degradation.sf
|
||||
|
||||
# ----------------------- The first degradation process ----------------------- #
|
||||
# blur
|
||||
out = filter2D(im_gt, kernel1)
|
||||
# random resize
|
||||
updown_type = random.choices(
|
||||
['up', 'down', 'keep'],
|
||||
conf_degradation['resize_prob'],
|
||||
)[0]
|
||||
if updown_type == 'up':
|
||||
scale = random.uniform(1, conf_degradation['resize_range'][1])
|
||||
elif updown_type == 'down':
|
||||
scale = random.uniform(conf_degradation['resize_range'][0], 1)
|
||||
else:
|
||||
scale = 1
|
||||
mode = random.choice(['area', 'bilinear', 'bicubic'])
|
||||
out = F.interpolate(out, scale_factor=scale, mode=mode)
|
||||
# add noise
|
||||
gray_noise_prob = conf_degradation['gray_noise_prob']
|
||||
if random.random() < conf_degradation['gaussian_noise_prob']:
|
||||
out = random_add_gaussian_noise_pt(
|
||||
out,
|
||||
sigma_range=conf_degradation['noise_range'],
|
||||
clip=True,
|
||||
rounds=False,
|
||||
gray_prob=gray_noise_prob,
|
||||
)
|
||||
else:
|
||||
out = random_add_poisson_noise_pt(
|
||||
out,
|
||||
scale_range=conf_degradation['poisson_scale_range'],
|
||||
gray_prob=gray_noise_prob,
|
||||
clip=True,
|
||||
rounds=False)
|
||||
# JPEG compression
|
||||
jpeg_p = out.new_zeros(out.size(0)).uniform_(*conf_degradation['jpeg_range'])
|
||||
out = torch.clamp(out, 0, 1) # clamp to [0, 1], otherwise JPEGer will result in unpleasant artifacts
|
||||
out = self.jpeger(out, quality=jpeg_p)
|
||||
|
||||
# ----------------------- The second degradation process ----------------------- #
|
||||
# blur
|
||||
if random.random() < conf_degradation['second_order_prob']:
|
||||
if random.random() < conf_degradation['second_blur_prob']:
|
||||
out = filter2D(out, kernel2)
|
||||
# random resize
|
||||
updown_type = random.choices(
|
||||
['up', 'down', 'keep'],
|
||||
conf_degradation['resize_prob2'],
|
||||
)[0]
|
||||
if updown_type == 'up':
|
||||
scale = random.uniform(1, conf_degradation['resize_range2'][1])
|
||||
elif updown_type == 'down':
|
||||
scale = random.uniform(conf_degradation['resize_range2'][0], 1)
|
||||
else:
|
||||
scale = 1
|
||||
mode = random.choice(['area', 'bilinear', 'bicubic'])
|
||||
out = F.interpolate(
|
||||
out,
|
||||
size=(int(ori_h / sf * scale), int(ori_w / sf * scale)),
|
||||
mode=mode,
|
||||
)
|
||||
# add noise
|
||||
gray_noise_prob = conf_degradation['gray_noise_prob2']
|
||||
if random.random() < conf_degradation['gaussian_noise_prob2']:
|
||||
out = random_add_gaussian_noise_pt(
|
||||
out,
|
||||
sigma_range=conf_degradation['noise_range2'],
|
||||
clip=True,
|
||||
rounds=False,
|
||||
gray_prob=gray_noise_prob,
|
||||
)
|
||||
else:
|
||||
out = random_add_poisson_noise_pt(
|
||||
out,
|
||||
scale_range=conf_degradation['poisson_scale_range2'],
|
||||
gray_prob=gray_noise_prob,
|
||||
clip=True,
|
||||
rounds=False,
|
||||
)
|
||||
|
||||
# JPEG compression + the final sinc filter
|
||||
# We also need to resize images to desired sizes. We group [resize back + sinc filter] together
|
||||
# as one operation.
|
||||
# We consider two orders:
|
||||
# 1. [resize back + sinc filter] + JPEG compression
|
||||
# 2. JPEG compression + [resize back + sinc filter]
|
||||
# Empirically, we find other combinations (sinc + JPEG + Resize) will introduce twisted lines.
|
||||
if random.random() < 0.5:
|
||||
# resize back + the final sinc filter
|
||||
mode = random.choice(['area', 'bilinear', 'bicubic'])
|
||||
out = F.interpolate(
|
||||
out,
|
||||
size=(ori_h // sf, ori_w // sf),
|
||||
mode=mode,
|
||||
)
|
||||
out = filter2D(out, sinc_kernel)
|
||||
# JPEG compression
|
||||
jpeg_p = out.new_zeros(out.size(0)).uniform_(*conf_degradation['jpeg_range2'])
|
||||
out = torch.clamp(out, 0, 1)
|
||||
out = self.jpeger(out, quality=jpeg_p)
|
||||
else:
|
||||
# JPEG compression
|
||||
jpeg_p = out.new_zeros(out.size(0)).uniform_(*conf_degradation['jpeg_range2'])
|
||||
out = torch.clamp(out, 0, 1)
|
||||
out = self.jpeger(out, quality=jpeg_p)
|
||||
# resize back + the final sinc filter
|
||||
mode = random.choice(['area', 'bilinear', 'bicubic'])
|
||||
out = F.interpolate(
|
||||
out,
|
||||
size=(ori_h // sf, ori_w // sf),
|
||||
mode=mode,
|
||||
)
|
||||
out = filter2D(out, sinc_kernel)
|
||||
|
||||
# clamp and round
|
||||
im_lq = torch.clamp((out * 255.0).round(), 0, 255) / 255.
|
||||
|
||||
return {'lq':im_lq.contiguous(), 'gt':im_gt}
|
||||
@@ -1,106 +0,0 @@
|
||||
import os
|
||||
from torch.utils import data as data
|
||||
from torchvision.transforms.functional import normalize
|
||||
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.data_util import paired_paths_from_folder, paired_paths_from_lmdb
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.transforms import augment, paired_random_crop
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils import FileClient, imfrombytes, img2tensor
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils.registry import DATASET_REGISTRY
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register(suffix='basicsr')
|
||||
class RealESRGANPairedDataset(data.Dataset):
|
||||
"""Paired image dataset for image restoration.
|
||||
|
||||
Read LQ (Low Quality, e.g. LR (Low Resolution), blurry, noisy, etc) and GT image pairs.
|
||||
|
||||
There are three modes:
|
||||
|
||||
1. **lmdb**: Use lmdb files. If opt['io_backend'] == lmdb.
|
||||
2. **meta_info_file**: Use meta information file to generate paths. \
|
||||
If opt['io_backend'] != lmdb and opt['meta_info_file'] is not None.
|
||||
3. **folder**: Scan folders to generate paths. The rest.
|
||||
|
||||
Args:
|
||||
opt (dict): Config for train datasets. It contains the following keys:
|
||||
dataroot_gt (str): Data root path for gt.
|
||||
dataroot_lq (str): Data root path for lq.
|
||||
meta_info (str): Path for meta information file.
|
||||
io_backend (dict): IO backend type and other kwarg.
|
||||
filename_tmpl (str): Template for each filename. Note that the template excludes the file extension.
|
||||
Default: '{}'.
|
||||
gt_size (int): Cropped patched size for gt patches.
|
||||
use_hflip (bool): Use horizontal flips.
|
||||
use_rot (bool): Use rotation (use vertical flip and transposing h and w for implementation).
|
||||
scale (bool): Scale, which will be added automatically.
|
||||
phase (str): 'train' or 'val'.
|
||||
"""
|
||||
|
||||
def __init__(self, opt):
|
||||
super(RealESRGANPairedDataset, self).__init__()
|
||||
self.opt = opt
|
||||
self.file_client = None
|
||||
self.io_backend_opt = opt['io_backend']
|
||||
# mean and std for normalizing the input images
|
||||
self.mean = opt['mean'] if 'mean' in opt else None
|
||||
self.std = opt['std'] if 'std' in opt else None
|
||||
|
||||
self.gt_folder, self.lq_folder = opt['dataroot_gt'], opt['dataroot_lq']
|
||||
self.filename_tmpl = opt['filename_tmpl'] if 'filename_tmpl' in opt else '{}'
|
||||
|
||||
# file client (lmdb io backend)
|
||||
if self.io_backend_opt['type'] == 'lmdb':
|
||||
self.io_backend_opt['db_paths'] = [self.lq_folder, self.gt_folder]
|
||||
self.io_backend_opt['client_keys'] = ['lq', 'gt']
|
||||
self.paths = paired_paths_from_lmdb([self.lq_folder, self.gt_folder], ['lq', 'gt'])
|
||||
elif 'meta_info' in self.opt and self.opt['meta_info'] is not None:
|
||||
# disk backend with meta_info
|
||||
# Each line in the meta_info describes the relative path to an image
|
||||
with open(self.opt['meta_info']) as fin:
|
||||
paths = [line.strip() for line in fin]
|
||||
self.paths = []
|
||||
for path in paths:
|
||||
gt_path, lq_path = path.split(', ')
|
||||
gt_path = os.path.join(self.gt_folder, gt_path)
|
||||
lq_path = os.path.join(self.lq_folder, lq_path)
|
||||
self.paths.append(dict([('gt_path', gt_path), ('lq_path', lq_path)]))
|
||||
else:
|
||||
# disk backend
|
||||
# it will scan the whole folder to get meta info
|
||||
# it will be time-consuming for folders with too many files. It is recommended using an extra meta txt file
|
||||
self.paths = paired_paths_from_folder([self.lq_folder, self.gt_folder], ['lq', 'gt'], self.filename_tmpl)
|
||||
|
||||
def __getitem__(self, index):
|
||||
if self.file_client is None:
|
||||
self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
|
||||
|
||||
scale = self.opt['scale']
|
||||
|
||||
# Load gt and lq images. Dimension order: HWC; channel order: BGR;
|
||||
# image range: [0, 1], float32.
|
||||
gt_path = self.paths[index]['gt_path']
|
||||
img_bytes = self.file_client.get(gt_path, 'gt')
|
||||
img_gt = imfrombytes(img_bytes, float32=True)
|
||||
lq_path = self.paths[index]['lq_path']
|
||||
img_bytes = self.file_client.get(lq_path, 'lq')
|
||||
img_lq = imfrombytes(img_bytes, float32=True)
|
||||
|
||||
# augmentation for training
|
||||
if self.opt['phase'] == 'train':
|
||||
gt_size = self.opt['gt_size']
|
||||
# random crop
|
||||
img_gt, img_lq = paired_random_crop(img_gt, img_lq, gt_size, scale, gt_path)
|
||||
# flip, rotation
|
||||
img_gt, img_lq = augment([img_gt, img_lq], self.opt['use_hflip'], self.opt['use_rot'])
|
||||
|
||||
# BGR to RGB, HWC to CHW, numpy to tensor
|
||||
img_gt, img_lq = img2tensor([img_gt, img_lq], bgr2rgb=True, float32=True)
|
||||
# normalize
|
||||
if self.mean is not None or self.std is not None:
|
||||
normalize(img_lq, self.mean, self.std, inplace=True)
|
||||
normalize(img_gt, self.mean, self.std, inplace=True)
|
||||
|
||||
return {'lq': img_lq, 'gt': img_gt, 'lq_path': lq_path, 'gt_path': gt_path}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.paths)
|
||||
@@ -1,352 +0,0 @@
|
||||
import numpy as np
|
||||
import random
|
||||
import torch
|
||||
from pathlib import Path
|
||||
from torch.utils import data as data
|
||||
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.transforms import augment, paired_random_crop
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils import FileClient, get_root_logger, imfrombytes, img2tensor
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils.flow_util import dequantize_flow
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils.registry import DATASET_REGISTRY
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class REDSDataset(data.Dataset):
|
||||
"""REDS dataset for training.
|
||||
|
||||
The keys are generated from a meta info txt file.
|
||||
basicsr/data/meta_info/meta_info_REDS_GT.txt
|
||||
|
||||
Each line contains:
|
||||
1. subfolder (clip) name; 2. frame number; 3. image shape, separated by
|
||||
a white space.
|
||||
Examples:
|
||||
000 100 (720,1280,3)
|
||||
001 100 (720,1280,3)
|
||||
...
|
||||
|
||||
Key examples: "000/00000000"
|
||||
GT (gt): Ground-Truth;
|
||||
LQ (lq): Low-Quality, e.g., low-resolution/blurry/noisy/compressed frames.
|
||||
|
||||
Args:
|
||||
opt (dict): Config for train dataset. It contains the following keys:
|
||||
dataroot_gt (str): Data root path for gt.
|
||||
dataroot_lq (str): Data root path for lq.
|
||||
dataroot_flow (str, optional): Data root path for flow.
|
||||
meta_info_file (str): Path for meta information file.
|
||||
val_partition (str): Validation partition types. 'REDS4' or 'official'.
|
||||
io_backend (dict): IO backend type and other kwarg.
|
||||
num_frame (int): Window size for input frames.
|
||||
gt_size (int): Cropped patched size for gt patches.
|
||||
interval_list (list): Interval list for temporal augmentation.
|
||||
random_reverse (bool): Random reverse input frames.
|
||||
use_hflip (bool): Use horizontal flips.
|
||||
use_rot (bool): Use rotation (use vertical flip and transposing h and w for implementation).
|
||||
scale (bool): Scale, which will be added automatically.
|
||||
"""
|
||||
|
||||
def __init__(self, opt):
|
||||
super(REDSDataset, self).__init__()
|
||||
self.opt = opt
|
||||
self.gt_root, self.lq_root = Path(opt['dataroot_gt']), Path(opt['dataroot_lq'])
|
||||
self.flow_root = Path(opt['dataroot_flow']) if opt['dataroot_flow'] is not None else None
|
||||
assert opt['num_frame'] % 2 == 1, (f'num_frame should be odd number, but got {opt["num_frame"]}')
|
||||
self.num_frame = opt['num_frame']
|
||||
self.num_half_frames = opt['num_frame'] // 2
|
||||
|
||||
self.keys = []
|
||||
with open(opt['meta_info_file'], 'r') as fin:
|
||||
for line in fin:
|
||||
folder, frame_num, _ = line.split(' ')
|
||||
self.keys.extend([f'{folder}/{i:08d}' for i in range(int(frame_num))])
|
||||
|
||||
# remove the video clips used in validation
|
||||
if opt['val_partition'] == 'REDS4':
|
||||
val_partition = ['000', '011', '015', '020']
|
||||
elif opt['val_partition'] == 'official':
|
||||
val_partition = [f'{v:03d}' for v in range(240, 270)]
|
||||
else:
|
||||
raise ValueError(f'Wrong validation partition {opt["val_partition"]}.'
|
||||
f"Supported ones are ['official', 'REDS4'].")
|
||||
self.keys = [v for v in self.keys if v.split('/')[0] not in val_partition]
|
||||
|
||||
# file client (io backend)
|
||||
self.file_client = None
|
||||
self.io_backend_opt = opt['io_backend']
|
||||
self.is_lmdb = False
|
||||
if self.io_backend_opt['type'] == 'lmdb':
|
||||
self.is_lmdb = True
|
||||
if self.flow_root is not None:
|
||||
self.io_backend_opt['db_paths'] = [self.lq_root, self.gt_root, self.flow_root]
|
||||
self.io_backend_opt['client_keys'] = ['lq', 'gt', 'flow']
|
||||
else:
|
||||
self.io_backend_opt['db_paths'] = [self.lq_root, self.gt_root]
|
||||
self.io_backend_opt['client_keys'] = ['lq', 'gt']
|
||||
|
||||
# temporal augmentation configs
|
||||
self.interval_list = opt['interval_list']
|
||||
self.random_reverse = opt['random_reverse']
|
||||
interval_str = ','.join(str(x) for x in opt['interval_list'])
|
||||
logger = get_root_logger()
|
||||
logger.info(f'Temporal augmentation interval list: [{interval_str}]; '
|
||||
f'random reverse is {self.random_reverse}.')
|
||||
|
||||
def __getitem__(self, index):
|
||||
if self.file_client is None:
|
||||
self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
|
||||
|
||||
scale = self.opt['scale']
|
||||
gt_size = self.opt['gt_size']
|
||||
key = self.keys[index]
|
||||
clip_name, frame_name = key.split('/') # key example: 000/00000000
|
||||
center_frame_idx = int(frame_name)
|
||||
|
||||
# determine the neighboring frames
|
||||
interval = random.choice(self.interval_list)
|
||||
|
||||
# ensure not exceeding the borders
|
||||
start_frame_idx = center_frame_idx - self.num_half_frames * interval
|
||||
end_frame_idx = center_frame_idx + self.num_half_frames * interval
|
||||
# each clip has 100 frames starting from 0 to 99
|
||||
while (start_frame_idx < 0) or (end_frame_idx > 99):
|
||||
center_frame_idx = random.randint(0, 99)
|
||||
start_frame_idx = (center_frame_idx - self.num_half_frames * interval)
|
||||
end_frame_idx = center_frame_idx + self.num_half_frames * interval
|
||||
frame_name = f'{center_frame_idx:08d}'
|
||||
neighbor_list = list(range(start_frame_idx, end_frame_idx + 1, interval))
|
||||
# random reverse
|
||||
if self.random_reverse and random.random() < 0.5:
|
||||
neighbor_list.reverse()
|
||||
|
||||
assert len(neighbor_list) == self.num_frame, (f'Wrong length of neighbor list: {len(neighbor_list)}')
|
||||
|
||||
# get the GT frame (as the center frame)
|
||||
if self.is_lmdb:
|
||||
img_gt_path = f'{clip_name}/{frame_name}'
|
||||
else:
|
||||
img_gt_path = self.gt_root / clip_name / f'{frame_name}.png'
|
||||
img_bytes = self.file_client.get(img_gt_path, 'gt')
|
||||
img_gt = imfrombytes(img_bytes, float32=True)
|
||||
|
||||
# get the neighboring LQ frames
|
||||
img_lqs = []
|
||||
for neighbor in neighbor_list:
|
||||
if self.is_lmdb:
|
||||
img_lq_path = f'{clip_name}/{neighbor:08d}'
|
||||
else:
|
||||
img_lq_path = self.lq_root / clip_name / f'{neighbor:08d}.png'
|
||||
img_bytes = self.file_client.get(img_lq_path, 'lq')
|
||||
img_lq = imfrombytes(img_bytes, float32=True)
|
||||
img_lqs.append(img_lq)
|
||||
|
||||
# get flows
|
||||
if self.flow_root is not None:
|
||||
img_flows = []
|
||||
# read previous flows
|
||||
for i in range(self.num_half_frames, 0, -1):
|
||||
if self.is_lmdb:
|
||||
flow_path = f'{clip_name}/{frame_name}_p{i}'
|
||||
else:
|
||||
flow_path = (self.flow_root / clip_name / f'{frame_name}_p{i}.png')
|
||||
img_bytes = self.file_client.get(flow_path, 'flow')
|
||||
cat_flow = imfrombytes(img_bytes, flag='grayscale', float32=False) # uint8, [0, 255]
|
||||
dx, dy = np.split(cat_flow, 2, axis=0)
|
||||
flow = dequantize_flow(dx, dy, max_val=20, denorm=False) # we use max_val 20 here.
|
||||
img_flows.append(flow)
|
||||
# read next flows
|
||||
for i in range(1, self.num_half_frames + 1):
|
||||
if self.is_lmdb:
|
||||
flow_path = f'{clip_name}/{frame_name}_n{i}'
|
||||
else:
|
||||
flow_path = (self.flow_root / clip_name / f'{frame_name}_n{i}.png')
|
||||
img_bytes = self.file_client.get(flow_path, 'flow')
|
||||
cat_flow = imfrombytes(img_bytes, flag='grayscale', float32=False) # uint8, [0, 255]
|
||||
dx, dy = np.split(cat_flow, 2, axis=0)
|
||||
flow = dequantize_flow(dx, dy, max_val=20, denorm=False) # we use max_val 20 here.
|
||||
img_flows.append(flow)
|
||||
|
||||
# for random crop, here, img_flows and img_lqs have the same
|
||||
# spatial size
|
||||
img_lqs.extend(img_flows)
|
||||
|
||||
# randomly crop
|
||||
img_gt, img_lqs = paired_random_crop(img_gt, img_lqs, gt_size, scale, img_gt_path)
|
||||
if self.flow_root is not None:
|
||||
img_lqs, img_flows = img_lqs[:self.num_frame], img_lqs[self.num_frame:]
|
||||
|
||||
# augmentation - flip, rotate
|
||||
img_lqs.append(img_gt)
|
||||
if self.flow_root is not None:
|
||||
img_results, img_flows = augment(img_lqs, self.opt['use_hflip'], self.opt['use_rot'], img_flows)
|
||||
else:
|
||||
img_results = augment(img_lqs, self.opt['use_hflip'], self.opt['use_rot'])
|
||||
|
||||
img_results = img2tensor(img_results)
|
||||
img_lqs = torch.stack(img_results[0:-1], dim=0)
|
||||
img_gt = img_results[-1]
|
||||
|
||||
if self.flow_root is not None:
|
||||
img_flows = img2tensor(img_flows)
|
||||
# add the zero center flow
|
||||
img_flows.insert(self.num_half_frames, torch.zeros_like(img_flows[0]))
|
||||
img_flows = torch.stack(img_flows, dim=0)
|
||||
|
||||
# img_lqs: (t, c, h, w)
|
||||
# img_flows: (t, 2, h, w)
|
||||
# img_gt: (c, h, w)
|
||||
# key: str
|
||||
if self.flow_root is not None:
|
||||
return {'lq': img_lqs, 'flow': img_flows, 'gt': img_gt, 'key': key}
|
||||
else:
|
||||
return {'lq': img_lqs, 'gt': img_gt, 'key': key}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.keys)
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class REDSRecurrentDataset(data.Dataset):
|
||||
"""REDS dataset for training recurrent networks.
|
||||
|
||||
The keys are generated from a meta info txt file.
|
||||
basicsr/data/meta_info/meta_info_REDS_GT.txt
|
||||
|
||||
Each line contains:
|
||||
1. subfolder (clip) name; 2. frame number; 3. image shape, separated by
|
||||
a white space.
|
||||
Examples:
|
||||
000 100 (720,1280,3)
|
||||
001 100 (720,1280,3)
|
||||
...
|
||||
|
||||
Key examples: "000/00000000"
|
||||
GT (gt): Ground-Truth;
|
||||
LQ (lq): Low-Quality, e.g., low-resolution/blurry/noisy/compressed frames.
|
||||
|
||||
Args:
|
||||
opt (dict): Config for train dataset. It contains the following keys:
|
||||
dataroot_gt (str): Data root path for gt.
|
||||
dataroot_lq (str): Data root path for lq.
|
||||
dataroot_flow (str, optional): Data root path for flow.
|
||||
meta_info_file (str): Path for meta information file.
|
||||
val_partition (str): Validation partition types. 'REDS4' or 'official'.
|
||||
io_backend (dict): IO backend type and other kwarg.
|
||||
num_frame (int): Window size for input frames.
|
||||
gt_size (int): Cropped patched size for gt patches.
|
||||
interval_list (list): Interval list for temporal augmentation.
|
||||
random_reverse (bool): Random reverse input frames.
|
||||
use_hflip (bool): Use horizontal flips.
|
||||
use_rot (bool): Use rotation (use vertical flip and transposing h and w for implementation).
|
||||
scale (bool): Scale, which will be added automatically.
|
||||
"""
|
||||
|
||||
def __init__(self, opt):
|
||||
super(REDSRecurrentDataset, self).__init__()
|
||||
self.opt = opt
|
||||
self.gt_root, self.lq_root = Path(opt['dataroot_gt']), Path(opt['dataroot_lq'])
|
||||
self.num_frame = opt['num_frame']
|
||||
|
||||
self.keys = []
|
||||
with open(opt['meta_info_file'], 'r') as fin:
|
||||
for line in fin:
|
||||
folder, frame_num, _ = line.split(' ')
|
||||
self.keys.extend([f'{folder}/{i:08d}' for i in range(int(frame_num))])
|
||||
|
||||
# remove the video clips used in validation
|
||||
if opt['val_partition'] == 'REDS4':
|
||||
val_partition = ['000', '011', '015', '020']
|
||||
elif opt['val_partition'] == 'official':
|
||||
val_partition = [f'{v:03d}' for v in range(240, 270)]
|
||||
else:
|
||||
raise ValueError(f'Wrong validation partition {opt["val_partition"]}.'
|
||||
f"Supported ones are ['official', 'REDS4'].")
|
||||
if opt['test_mode']:
|
||||
self.keys = [v for v in self.keys if v.split('/')[0] in val_partition]
|
||||
else:
|
||||
self.keys = [v for v in self.keys if v.split('/')[0] not in val_partition]
|
||||
|
||||
# file client (io backend)
|
||||
self.file_client = None
|
||||
self.io_backend_opt = opt['io_backend']
|
||||
self.is_lmdb = False
|
||||
if self.io_backend_opt['type'] == 'lmdb':
|
||||
self.is_lmdb = True
|
||||
if hasattr(self, 'flow_root') and self.flow_root is not None:
|
||||
self.io_backend_opt['db_paths'] = [self.lq_root, self.gt_root, self.flow_root]
|
||||
self.io_backend_opt['client_keys'] = ['lq', 'gt', 'flow']
|
||||
else:
|
||||
self.io_backend_opt['db_paths'] = [self.lq_root, self.gt_root]
|
||||
self.io_backend_opt['client_keys'] = ['lq', 'gt']
|
||||
|
||||
# temporal augmentation configs
|
||||
self.interval_list = opt.get('interval_list', [1])
|
||||
self.random_reverse = opt.get('random_reverse', False)
|
||||
interval_str = ','.join(str(x) for x in self.interval_list)
|
||||
logger = get_root_logger()
|
||||
logger.info(f'Temporal augmentation interval list: [{interval_str}]; '
|
||||
f'random reverse is {self.random_reverse}.')
|
||||
|
||||
def __getitem__(self, index):
|
||||
if self.file_client is None:
|
||||
self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
|
||||
|
||||
scale = self.opt['scale']
|
||||
gt_size = self.opt['gt_size']
|
||||
key = self.keys[index]
|
||||
clip_name, frame_name = key.split('/') # key example: 000/00000000
|
||||
|
||||
# determine the neighboring frames
|
||||
interval = random.choice(self.interval_list)
|
||||
|
||||
# ensure not exceeding the borders
|
||||
start_frame_idx = int(frame_name)
|
||||
if start_frame_idx > 100 - self.num_frame * interval:
|
||||
start_frame_idx = random.randint(0, 100 - self.num_frame * interval)
|
||||
end_frame_idx = start_frame_idx + self.num_frame * interval
|
||||
|
||||
neighbor_list = list(range(start_frame_idx, end_frame_idx, interval))
|
||||
|
||||
# random reverse
|
||||
if self.random_reverse and random.random() < 0.5:
|
||||
neighbor_list.reverse()
|
||||
|
||||
# get the neighboring LQ and GT frames
|
||||
img_lqs = []
|
||||
img_gts = []
|
||||
for neighbor in neighbor_list:
|
||||
if self.is_lmdb:
|
||||
img_lq_path = f'{clip_name}/{neighbor:08d}'
|
||||
img_gt_path = f'{clip_name}/{neighbor:08d}'
|
||||
else:
|
||||
img_lq_path = self.lq_root / clip_name / f'{neighbor:08d}.png'
|
||||
img_gt_path = self.gt_root / clip_name / f'{neighbor:08d}.png'
|
||||
|
||||
# get LQ
|
||||
img_bytes = self.file_client.get(img_lq_path, 'lq')
|
||||
img_lq = imfrombytes(img_bytes, float32=True)
|
||||
img_lqs.append(img_lq)
|
||||
|
||||
# get GT
|
||||
img_bytes = self.file_client.get(img_gt_path, 'gt')
|
||||
img_gt = imfrombytes(img_bytes, float32=True)
|
||||
img_gts.append(img_gt)
|
||||
|
||||
# randomly crop
|
||||
img_gts, img_lqs = paired_random_crop(img_gts, img_lqs, gt_size, scale, img_gt_path)
|
||||
|
||||
# augmentation - flip, rotate
|
||||
img_lqs.extend(img_gts)
|
||||
img_results = augment(img_lqs, self.opt['use_hflip'], self.opt['use_rot'])
|
||||
|
||||
img_results = img2tensor(img_results)
|
||||
img_gts = torch.stack(img_results[len(img_lqs) // 2:], dim=0)
|
||||
img_lqs = torch.stack(img_results[:len(img_lqs) // 2], dim=0)
|
||||
|
||||
# img_lqs: (t, c, h, w)
|
||||
# img_gts: (t, c, h, w)
|
||||
# key: str
|
||||
return {'lq': img_lqs, 'gt': img_gts, 'key': key}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.keys)
|
||||
@@ -1,68 +0,0 @@
|
||||
from os import path as osp
|
||||
from torch.utils import data as data
|
||||
from torchvision.transforms.functional import normalize
|
||||
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.data_util import paths_from_lmdb
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils import FileClient, imfrombytes, img2tensor, rgb2ycbcr, scandir
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils.registry import DATASET_REGISTRY
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class SingleImageDataset(data.Dataset):
|
||||
"""Read only lq images in the test phase.
|
||||
|
||||
Read LQ (Low Quality, e.g. LR (Low Resolution), blurry, noisy, etc).
|
||||
|
||||
There are two modes:
|
||||
1. 'meta_info_file': Use meta information file to generate paths.
|
||||
2. 'folder': Scan folders to generate paths.
|
||||
|
||||
Args:
|
||||
opt (dict): Config for train datasets. It contains the following keys:
|
||||
dataroot_lq (str): Data root path for lq.
|
||||
meta_info_file (str): Path for meta information file.
|
||||
io_backend (dict): IO backend type and other kwarg.
|
||||
"""
|
||||
|
||||
def __init__(self, opt):
|
||||
super(SingleImageDataset, self).__init__()
|
||||
self.opt = opt
|
||||
# file client (io backend)
|
||||
self.file_client = None
|
||||
self.io_backend_opt = opt['io_backend']
|
||||
self.mean = opt['mean'] if 'mean' in opt else None
|
||||
self.std = opt['std'] if 'std' in opt else None
|
||||
self.lq_folder = opt['dataroot_lq']
|
||||
|
||||
if self.io_backend_opt['type'] == 'lmdb':
|
||||
self.io_backend_opt['db_paths'] = [self.lq_folder]
|
||||
self.io_backend_opt['client_keys'] = ['lq']
|
||||
self.paths = paths_from_lmdb(self.lq_folder)
|
||||
elif 'meta_info_file' in self.opt:
|
||||
with open(self.opt['meta_info_file'], 'r') as fin:
|
||||
self.paths = [osp.join(self.lq_folder, line.rstrip().split(' ')[0]) for line in fin]
|
||||
else:
|
||||
self.paths = sorted(list(scandir(self.lq_folder, full_path=True)))
|
||||
|
||||
def __getitem__(self, index):
|
||||
if self.file_client is None:
|
||||
self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
|
||||
|
||||
# load lq image
|
||||
lq_path = self.paths[index]
|
||||
img_bytes = self.file_client.get(lq_path, 'lq')
|
||||
img_lq = imfrombytes(img_bytes, float32=True)
|
||||
|
||||
# color space transform
|
||||
if 'color' in self.opt and self.opt['color'] == 'y':
|
||||
img_lq = rgb2ycbcr(img_lq, y_only=True)[..., None]
|
||||
|
||||
# BGR to RGB, HWC to CHW, numpy to tensor
|
||||
img_lq = img2tensor(img_lq, bgr2rgb=True, float32=True)
|
||||
# normalize
|
||||
if self.mean is not None or self.std is not None:
|
||||
normalize(img_lq, self.mean, self.std, inplace=True)
|
||||
return {'lq': img_lq, 'lq_path': lq_path}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.paths)
|
||||
@@ -1,207 +0,0 @@
|
||||
import cv2
|
||||
import random
|
||||
import torch
|
||||
|
||||
def mod_crop(img, scale):
|
||||
"""Mod crop images, used during testing.
|
||||
|
||||
Args:
|
||||
img (ndarray): Input image.
|
||||
scale (int): Scale factor.
|
||||
|
||||
Returns:
|
||||
ndarray: Result image.
|
||||
"""
|
||||
img = img.copy()
|
||||
if img.ndim in (2, 3):
|
||||
h, w = img.shape[0], img.shape[1]
|
||||
h_remainder, w_remainder = h % scale, w % scale
|
||||
img = img[:h - h_remainder, :w - w_remainder, ...]
|
||||
else:
|
||||
raise ValueError(f'Wrong img ndim: {img.ndim}.')
|
||||
return img
|
||||
|
||||
|
||||
def paired_random_crop(img_gts, img_lqs, gt_patch_size, scale, gt_path=None):
|
||||
"""Paired random crop. Support Numpy array and Tensor inputs.
|
||||
|
||||
It crops lists of lq and gt images with corresponding locations.
|
||||
|
||||
Args:
|
||||
img_gts (list[ndarray] | ndarray | list[Tensor] | Tensor): GT images. Note that all images
|
||||
should have the same shape. If the input is an ndarray, it will
|
||||
be transformed to a list containing itself.
|
||||
img_lqs (list[ndarray] | ndarray): LQ images. Note that all images
|
||||
should have the same shape. If the input is an ndarray, it will
|
||||
be transformed to a list containing itself.
|
||||
gt_patch_size (int): GT patch size.
|
||||
scale (int): Scale factor.
|
||||
gt_path (str): Path to ground-truth. Default: None.
|
||||
|
||||
Returns:
|
||||
list[ndarray] | ndarray: GT images and LQ images. If returned results
|
||||
only have one element, just return ndarray.
|
||||
"""
|
||||
|
||||
if not isinstance(img_gts, list):
|
||||
img_gts = [img_gts]
|
||||
if not isinstance(img_lqs, list):
|
||||
img_lqs = [img_lqs]
|
||||
|
||||
# determine input type: Numpy array or Tensor
|
||||
input_type = 'Tensor' if torch.is_tensor(img_gts[0]) else 'Numpy'
|
||||
|
||||
if input_type == 'Tensor':
|
||||
h_lq, w_lq = img_lqs[0].size()[-2:]
|
||||
h_gt, w_gt = img_gts[0].size()[-2:]
|
||||
else:
|
||||
h_lq, w_lq = img_lqs[0].shape[0:2]
|
||||
h_gt, w_gt = img_gts[0].shape[0:2]
|
||||
lq_patch_size = gt_patch_size // scale
|
||||
|
||||
if h_gt != h_lq * scale or w_gt != w_lq * scale:
|
||||
raise ValueError(f'Scale mismatches. GT ({h_gt}, {w_gt}) is not {scale}x ',
|
||||
f'multiplication of LQ ({h_lq}, {w_lq}).')
|
||||
if h_lq < lq_patch_size or w_lq < lq_patch_size:
|
||||
raise ValueError(f'LQ ({h_lq}, {w_lq}) is smaller than patch size '
|
||||
f'({lq_patch_size}, {lq_patch_size}). '
|
||||
f'Please remove {gt_path}.')
|
||||
|
||||
# randomly choose top and left coordinates for lq patch
|
||||
top = random.randint(0, h_lq - lq_patch_size)
|
||||
left = random.randint(0, w_lq - lq_patch_size)
|
||||
|
||||
# crop lq patch
|
||||
if input_type == 'Tensor':
|
||||
img_lqs = [v[:, :, top:top + lq_patch_size, left:left + lq_patch_size] for v in img_lqs]
|
||||
else:
|
||||
img_lqs = [v[top:top + lq_patch_size, left:left + lq_patch_size, ...] for v in img_lqs]
|
||||
|
||||
# crop corresponding gt patch
|
||||
top_gt, left_gt = int(top * scale), int(left * scale)
|
||||
if input_type == 'Tensor':
|
||||
img_gts = [v[:, :, top_gt:top_gt + gt_patch_size, left_gt:left_gt + gt_patch_size] for v in img_gts]
|
||||
else:
|
||||
img_gts = [v[top_gt:top_gt + gt_patch_size, left_gt:left_gt + gt_patch_size, ...] for v in img_gts]
|
||||
if len(img_gts) == 1:
|
||||
img_gts = img_gts[0]
|
||||
if len(img_lqs) == 1:
|
||||
img_lqs = img_lqs[0]
|
||||
return img_gts, img_lqs
|
||||
|
||||
|
||||
def augment(imgs, hflip=True, rotation=True, flows=None, return_status=False):
|
||||
"""Augment: horizontal flips OR rotate (0, 90, 180, 270 degrees).
|
||||
|
||||
We use vertical flip and transpose for rotation implementation.
|
||||
All the images in the list use the same augmentation.
|
||||
|
||||
Args:
|
||||
imgs (list[ndarray] | ndarray): Images to be augmented. If the input
|
||||
is an ndarray, it will be transformed to a list.
|
||||
hflip (bool): Horizontal flip. Default: True.
|
||||
rotation (bool): Ratotation. Default: True.
|
||||
flows (list[ndarray]: Flows to be augmented. If the input is an
|
||||
ndarray, it will be transformed to a list.
|
||||
Dimension is (h, w, 2). Default: None.
|
||||
return_status (bool): Return the status of flip and rotation.
|
||||
Default: False.
|
||||
|
||||
Returns:
|
||||
list[ndarray] | ndarray: Augmented images and flows. If returned
|
||||
results only have one element, just return ndarray.
|
||||
|
||||
"""
|
||||
hflip = hflip and random.random() < 0.5
|
||||
vflip = rotation and random.random() < 0.5
|
||||
rot90 = rotation and random.random() < 0.5
|
||||
|
||||
def _augment(img):
|
||||
if hflip: # horizontal
|
||||
cv2.flip(img, 1, img)
|
||||
if vflip: # vertical
|
||||
cv2.flip(img, 0, img)
|
||||
if rot90:
|
||||
img = img.transpose(1, 0, 2)
|
||||
return img
|
||||
|
||||
def _augment_flow(flow):
|
||||
if hflip: # horizontal
|
||||
cv2.flip(flow, 1, flow)
|
||||
flow[:, :, 0] *= -1
|
||||
if vflip: # vertical
|
||||
cv2.flip(flow, 0, flow)
|
||||
flow[:, :, 1] *= -1
|
||||
if rot90:
|
||||
flow = flow.transpose(1, 0, 2)
|
||||
flow = flow[:, :, [1, 0]]
|
||||
return flow
|
||||
|
||||
if not isinstance(imgs, list):
|
||||
imgs = [imgs]
|
||||
imgs = [_augment(img) for img in imgs]
|
||||
if len(imgs) == 1:
|
||||
imgs = imgs[0]
|
||||
|
||||
if flows is not None:
|
||||
if not isinstance(flows, list):
|
||||
flows = [flows]
|
||||
flows = [_augment_flow(flow) for flow in flows]
|
||||
if len(flows) == 1:
|
||||
flows = flows[0]
|
||||
return imgs, flows
|
||||
else:
|
||||
if return_status:
|
||||
return imgs, (hflip, vflip, rot90)
|
||||
else:
|
||||
return imgs
|
||||
|
||||
|
||||
def img_rotate(img, angle, center=None, scale=1.0):
|
||||
"""Rotate image.
|
||||
|
||||
Args:
|
||||
img (ndarray): Image to be rotated.
|
||||
angle (float): Rotation angle in degrees. Positive values mean
|
||||
counter-clockwise rotation.
|
||||
center (tuple[int]): Rotation center. If the center is None,
|
||||
initialize it as the center of the image. Default: None.
|
||||
scale (float): Isotropic scale factor. Default: 1.0.
|
||||
"""
|
||||
(h, w) = img.shape[:2]
|
||||
|
||||
if center is None:
|
||||
center = (w // 2, h // 2)
|
||||
|
||||
matrix = cv2.getRotationMatrix2D(center, angle, scale)
|
||||
rotated_img = cv2.warpAffine(img, matrix, (w, h))
|
||||
return rotated_img
|
||||
|
||||
def random_crop(im, pch_size):
|
||||
'''
|
||||
Randomly crop a patch from the give image.
|
||||
'''
|
||||
h, w = im.shape[:2]
|
||||
# padding if necessary
|
||||
if h < pch_size or w < pch_size:
|
||||
pad_h = min(max(0, pch_size - h), h)
|
||||
pad_w = min(max(0, pch_size - w), w)
|
||||
im = cv2.copyMakeBorder(im, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT_101)
|
||||
|
||||
h, w = im.shape[:2]
|
||||
if h == pch_size:
|
||||
ind_h = 0
|
||||
elif h > pch_size:
|
||||
ind_h = random.randint(0, h-pch_size)
|
||||
else:
|
||||
raise ValueError('Image height is smaller than the patch size')
|
||||
if w == pch_size:
|
||||
ind_w = 0
|
||||
elif w > pch_size:
|
||||
ind_w = random.randint(0, w-pch_size)
|
||||
else:
|
||||
raise ValueError('Image width is smaller than the patch size')
|
||||
|
||||
im_pch = im[ind_h:ind_h+pch_size, ind_w:ind_w+pch_size,]
|
||||
|
||||
return im_pch
|
||||
@@ -1,283 +0,0 @@
|
||||
import glob
|
||||
import torch
|
||||
from os import path as osp
|
||||
from torch.utils import data as data
|
||||
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.data_util import duf_downsample, generate_frame_indices, read_img_seq
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils import get_root_logger, scandir
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils.registry import DATASET_REGISTRY
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class VideoTestDataset(data.Dataset):
|
||||
"""Video test dataset.
|
||||
|
||||
Supported datasets: Vid4, REDS4, REDSofficial.
|
||||
More generally, it supports testing dataset with following structures:
|
||||
|
||||
::
|
||||
|
||||
dataroot
|
||||
├── subfolder1
|
||||
├── frame000
|
||||
├── frame001
|
||||
├── ...
|
||||
├── subfolder2
|
||||
├── frame000
|
||||
├── frame001
|
||||
├── ...
|
||||
├── ...
|
||||
|
||||
For testing datasets, there is no need to prepare LMDB files.
|
||||
|
||||
Args:
|
||||
opt (dict): Config for train dataset. It contains the following keys:
|
||||
dataroot_gt (str): Data root path for gt.
|
||||
dataroot_lq (str): Data root path for lq.
|
||||
io_backend (dict): IO backend type and other kwarg.
|
||||
cache_data (bool): Whether to cache testing datasets.
|
||||
name (str): Dataset name.
|
||||
meta_info_file (str): The path to the file storing the list of test folders. If not provided, all the folders
|
||||
in the dataroot will be used.
|
||||
num_frame (int): Window size for input frames.
|
||||
padding (str): Padding mode.
|
||||
"""
|
||||
|
||||
def __init__(self, opt):
|
||||
super(VideoTestDataset, self).__init__()
|
||||
self.opt = opt
|
||||
self.cache_data = opt['cache_data']
|
||||
self.gt_root, self.lq_root = opt['dataroot_gt'], opt['dataroot_lq']
|
||||
self.data_info = {'lq_path': [], 'gt_path': [], 'folder': [], 'idx': [], 'border': []}
|
||||
# file client (io backend)
|
||||
self.file_client = None
|
||||
self.io_backend_opt = opt['io_backend']
|
||||
assert self.io_backend_opt['type'] != 'lmdb', 'No need to use lmdb during validation/test.'
|
||||
|
||||
logger = get_root_logger()
|
||||
logger.info(f'Generate data info for VideoTestDataset - {opt["name"]}')
|
||||
self.imgs_lq, self.imgs_gt = {}, {}
|
||||
if 'meta_info_file' in opt:
|
||||
with open(opt['meta_info_file'], 'r') as fin:
|
||||
subfolders = [line.split(' ')[0] for line in fin]
|
||||
subfolders_lq = [osp.join(self.lq_root, key) for key in subfolders]
|
||||
subfolders_gt = [osp.join(self.gt_root, key) for key in subfolders]
|
||||
else:
|
||||
subfolders_lq = sorted(glob.glob(osp.join(self.lq_root, '*')))
|
||||
subfolders_gt = sorted(glob.glob(osp.join(self.gt_root, '*')))
|
||||
|
||||
if opt['name'].lower() in ['vid4', 'reds4', 'redsofficial']:
|
||||
for subfolder_lq, subfolder_gt in zip(subfolders_lq, subfolders_gt):
|
||||
# get frame list for lq and gt
|
||||
subfolder_name = osp.basename(subfolder_lq)
|
||||
img_paths_lq = sorted(list(scandir(subfolder_lq, full_path=True)))
|
||||
img_paths_gt = sorted(list(scandir(subfolder_gt, full_path=True)))
|
||||
|
||||
max_idx = len(img_paths_lq)
|
||||
assert max_idx == len(img_paths_gt), (f'Different number of images in lq ({max_idx})'
|
||||
f' and gt folders ({len(img_paths_gt)})')
|
||||
|
||||
self.data_info['lq_path'].extend(img_paths_lq)
|
||||
self.data_info['gt_path'].extend(img_paths_gt)
|
||||
self.data_info['folder'].extend([subfolder_name] * max_idx)
|
||||
for i in range(max_idx):
|
||||
self.data_info['idx'].append(f'{i}/{max_idx}')
|
||||
border_l = [0] * max_idx
|
||||
for i in range(self.opt['num_frame'] // 2):
|
||||
border_l[i] = 1
|
||||
border_l[max_idx - i - 1] = 1
|
||||
self.data_info['border'].extend(border_l)
|
||||
|
||||
# cache data or save the frame list
|
||||
if self.cache_data:
|
||||
logger.info(f'Cache {subfolder_name} for VideoTestDataset...')
|
||||
self.imgs_lq[subfolder_name] = read_img_seq(img_paths_lq)
|
||||
self.imgs_gt[subfolder_name] = read_img_seq(img_paths_gt)
|
||||
else:
|
||||
self.imgs_lq[subfolder_name] = img_paths_lq
|
||||
self.imgs_gt[subfolder_name] = img_paths_gt
|
||||
else:
|
||||
raise ValueError(f'Non-supported video test dataset: {type(opt["name"])}')
|
||||
|
||||
def __getitem__(self, index):
|
||||
folder = self.data_info['folder'][index]
|
||||
idx, max_idx = self.data_info['idx'][index].split('/')
|
||||
idx, max_idx = int(idx), int(max_idx)
|
||||
border = self.data_info['border'][index]
|
||||
lq_path = self.data_info['lq_path'][index]
|
||||
|
||||
select_idx = generate_frame_indices(idx, max_idx, self.opt['num_frame'], padding=self.opt['padding'])
|
||||
|
||||
if self.cache_data:
|
||||
imgs_lq = self.imgs_lq[folder].index_select(0, torch.LongTensor(select_idx))
|
||||
img_gt = self.imgs_gt[folder][idx]
|
||||
else:
|
||||
img_paths_lq = [self.imgs_lq[folder][i] for i in select_idx]
|
||||
imgs_lq = read_img_seq(img_paths_lq)
|
||||
img_gt = read_img_seq([self.imgs_gt[folder][idx]])
|
||||
img_gt.squeeze_(0)
|
||||
|
||||
return {
|
||||
'lq': imgs_lq, # (t, c, h, w)
|
||||
'gt': img_gt, # (c, h, w)
|
||||
'folder': folder, # folder name
|
||||
'idx': self.data_info['idx'][index], # e.g., 0/99
|
||||
'border': border, # 1 for border, 0 for non-border
|
||||
'lq_path': lq_path # center frame
|
||||
}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_info['gt_path'])
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class VideoTestVimeo90KDataset(data.Dataset):
|
||||
"""Video test dataset for Vimeo90k-Test dataset.
|
||||
|
||||
It only keeps the center frame for testing.
|
||||
For testing datasets, there is no need to prepare LMDB files.
|
||||
|
||||
Args:
|
||||
opt (dict): Config for train dataset. It contains the following keys:
|
||||
dataroot_gt (str): Data root path for gt.
|
||||
dataroot_lq (str): Data root path for lq.
|
||||
io_backend (dict): IO backend type and other kwarg.
|
||||
cache_data (bool): Whether to cache testing datasets.
|
||||
name (str): Dataset name.
|
||||
meta_info_file (str): The path to the file storing the list of test folders. If not provided, all the folders
|
||||
in the dataroot will be used.
|
||||
num_frame (int): Window size for input frames.
|
||||
padding (str): Padding mode.
|
||||
"""
|
||||
|
||||
def __init__(self, opt):
|
||||
super(VideoTestVimeo90KDataset, self).__init__()
|
||||
self.opt = opt
|
||||
self.cache_data = opt['cache_data']
|
||||
if self.cache_data:
|
||||
raise NotImplementedError('cache_data in Vimeo90K-Test dataset is not implemented.')
|
||||
self.gt_root, self.lq_root = opt['dataroot_gt'], opt['dataroot_lq']
|
||||
self.data_info = {'lq_path': [], 'gt_path': [], 'folder': [], 'idx': [], 'border': []}
|
||||
neighbor_list = [i + (9 - opt['num_frame']) // 2 for i in range(opt['num_frame'])]
|
||||
|
||||
# file client (io backend)
|
||||
self.file_client = None
|
||||
self.io_backend_opt = opt['io_backend']
|
||||
assert self.io_backend_opt['type'] != 'lmdb', 'No need to use lmdb during validation/test.'
|
||||
|
||||
logger = get_root_logger()
|
||||
logger.info(f'Generate data info for VideoTestDataset - {opt["name"]}')
|
||||
with open(opt['meta_info_file'], 'r') as fin:
|
||||
subfolders = [line.split(' ')[0] for line in fin]
|
||||
for idx, subfolder in enumerate(subfolders):
|
||||
gt_path = osp.join(self.gt_root, subfolder, 'im4.png')
|
||||
self.data_info['gt_path'].append(gt_path)
|
||||
lq_paths = [osp.join(self.lq_root, subfolder, f'im{i}.png') for i in neighbor_list]
|
||||
self.data_info['lq_path'].append(lq_paths)
|
||||
self.data_info['folder'].append('vimeo90k')
|
||||
self.data_info['idx'].append(f'{idx}/{len(subfolders)}')
|
||||
self.data_info['border'].append(0)
|
||||
|
||||
def __getitem__(self, index):
|
||||
lq_path = self.data_info['lq_path'][index]
|
||||
gt_path = self.data_info['gt_path'][index]
|
||||
imgs_lq = read_img_seq(lq_path)
|
||||
img_gt = read_img_seq([gt_path])
|
||||
img_gt.squeeze_(0)
|
||||
|
||||
return {
|
||||
'lq': imgs_lq, # (t, c, h, w)
|
||||
'gt': img_gt, # (c, h, w)
|
||||
'folder': self.data_info['folder'][index], # folder name
|
||||
'idx': self.data_info['idx'][index], # e.g., 0/843
|
||||
'border': self.data_info['border'][index], # 0 for non-border
|
||||
'lq_path': lq_path[self.opt['num_frame'] // 2] # center frame
|
||||
}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_info['gt_path'])
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class VideoTestDUFDataset(VideoTestDataset):
|
||||
""" Video test dataset for DUF dataset.
|
||||
|
||||
Args:
|
||||
opt (dict): Config for train dataset. Most of keys are the same as VideoTestDataset.
|
||||
It has the following extra keys:
|
||||
use_duf_downsampling (bool): Whether to use duf downsampling to generate low-resolution frames.
|
||||
scale (bool): Scale, which will be added automatically.
|
||||
"""
|
||||
|
||||
def __getitem__(self, index):
|
||||
folder = self.data_info['folder'][index]
|
||||
idx, max_idx = self.data_info['idx'][index].split('/')
|
||||
idx, max_idx = int(idx), int(max_idx)
|
||||
border = self.data_info['border'][index]
|
||||
lq_path = self.data_info['lq_path'][index]
|
||||
|
||||
select_idx = generate_frame_indices(idx, max_idx, self.opt['num_frame'], padding=self.opt['padding'])
|
||||
|
||||
if self.cache_data:
|
||||
if self.opt['use_duf_downsampling']:
|
||||
# read imgs_gt to generate low-resolution frames
|
||||
imgs_lq = self.imgs_gt[folder].index_select(0, torch.LongTensor(select_idx))
|
||||
imgs_lq = duf_downsample(imgs_lq, kernel_size=13, scale=self.opt['scale'])
|
||||
else:
|
||||
imgs_lq = self.imgs_lq[folder].index_select(0, torch.LongTensor(select_idx))
|
||||
img_gt = self.imgs_gt[folder][idx]
|
||||
else:
|
||||
if self.opt['use_duf_downsampling']:
|
||||
img_paths_lq = [self.imgs_gt[folder][i] for i in select_idx]
|
||||
# read imgs_gt to generate low-resolution frames
|
||||
imgs_lq = read_img_seq(img_paths_lq, require_mod_crop=True, scale=self.opt['scale'])
|
||||
imgs_lq = duf_downsample(imgs_lq, kernel_size=13, scale=self.opt['scale'])
|
||||
else:
|
||||
img_paths_lq = [self.imgs_lq[folder][i] for i in select_idx]
|
||||
imgs_lq = read_img_seq(img_paths_lq)
|
||||
img_gt = read_img_seq([self.imgs_gt[folder][idx]], require_mod_crop=True, scale=self.opt['scale'])
|
||||
img_gt.squeeze_(0)
|
||||
|
||||
return {
|
||||
'lq': imgs_lq, # (t, c, h, w)
|
||||
'gt': img_gt, # (c, h, w)
|
||||
'folder': folder, # folder name
|
||||
'idx': self.data_info['idx'][index], # e.g., 0/99
|
||||
'border': border, # 1 for border, 0 for non-border
|
||||
'lq_path': lq_path # center frame
|
||||
}
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class VideoRecurrentTestDataset(VideoTestDataset):
|
||||
"""Video test dataset for recurrent architectures, which takes LR video
|
||||
frames as input and output corresponding HR video frames.
|
||||
|
||||
Args:
|
||||
opt (dict): Same as VideoTestDataset. Unused opt:
|
||||
padding (str): Padding mode.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, opt):
|
||||
super(VideoRecurrentTestDataset, self).__init__(opt)
|
||||
# Find unique folder strings
|
||||
self.folders = sorted(list(set(self.data_info['folder'])))
|
||||
|
||||
def __getitem__(self, index):
|
||||
folder = self.folders[index]
|
||||
|
||||
if self.cache_data:
|
||||
imgs_lq = self.imgs_lq[folder]
|
||||
imgs_gt = self.imgs_gt[folder]
|
||||
else:
|
||||
raise NotImplementedError('Without cache_data is not implemented.')
|
||||
|
||||
return {
|
||||
'lq': imgs_lq,
|
||||
'gt': imgs_gt,
|
||||
'folder': folder,
|
||||
}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.folders)
|
||||
@@ -1,199 +0,0 @@
|
||||
import random
|
||||
import torch
|
||||
from pathlib import Path
|
||||
from torch.utils import data as data
|
||||
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.transforms import augment, paired_random_crop
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils import FileClient, get_root_logger, imfrombytes, img2tensor
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils.registry import DATASET_REGISTRY
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class Vimeo90KDataset(data.Dataset):
|
||||
"""Vimeo90K dataset for training.
|
||||
|
||||
The keys are generated from a meta info txt file.
|
||||
basicsr/data/meta_info/meta_info_Vimeo90K_train_GT.txt
|
||||
|
||||
Each line contains the following items, separated by a white space.
|
||||
|
||||
1. clip name;
|
||||
2. frame number;
|
||||
3. image shape
|
||||
|
||||
Examples:
|
||||
|
||||
::
|
||||
|
||||
00001/0001 7 (256,448,3)
|
||||
00001/0002 7 (256,448,3)
|
||||
|
||||
- Key examples: "00001/0001"
|
||||
- GT (gt): Ground-Truth;
|
||||
- LQ (lq): Low-Quality, e.g., low-resolution/blurry/noisy/compressed frames.
|
||||
|
||||
The neighboring frame list for different num_frame:
|
||||
|
||||
::
|
||||
|
||||
num_frame | frame list
|
||||
1 | 4
|
||||
3 | 3,4,5
|
||||
5 | 2,3,4,5,6
|
||||
7 | 1,2,3,4,5,6,7
|
||||
|
||||
Args:
|
||||
opt (dict): Config for train dataset. It contains the following keys:
|
||||
dataroot_gt (str): Data root path for gt.
|
||||
dataroot_lq (str): Data root path for lq.
|
||||
meta_info_file (str): Path for meta information file.
|
||||
io_backend (dict): IO backend type and other kwarg.
|
||||
num_frame (int): Window size for input frames.
|
||||
gt_size (int): Cropped patched size for gt patches.
|
||||
random_reverse (bool): Random reverse input frames.
|
||||
use_hflip (bool): Use horizontal flips.
|
||||
use_rot (bool): Use rotation (use vertical flip and transposing h and w for implementation).
|
||||
scale (bool): Scale, which will be added automatically.
|
||||
"""
|
||||
|
||||
def __init__(self, opt):
|
||||
super(Vimeo90KDataset, self).__init__()
|
||||
self.opt = opt
|
||||
self.gt_root, self.lq_root = Path(opt['dataroot_gt']), Path(opt['dataroot_lq'])
|
||||
|
||||
with open(opt['meta_info_file'], 'r') as fin:
|
||||
self.keys = [line.split(' ')[0] for line in fin]
|
||||
|
||||
# file client (io backend)
|
||||
self.file_client = None
|
||||
self.io_backend_opt = opt['io_backend']
|
||||
self.is_lmdb = False
|
||||
if self.io_backend_opt['type'] == 'lmdb':
|
||||
self.is_lmdb = True
|
||||
self.io_backend_opt['db_paths'] = [self.lq_root, self.gt_root]
|
||||
self.io_backend_opt['client_keys'] = ['lq', 'gt']
|
||||
|
||||
# indices of input images
|
||||
self.neighbor_list = [i + (9 - opt['num_frame']) // 2 for i in range(opt['num_frame'])]
|
||||
|
||||
# temporal augmentation configs
|
||||
self.random_reverse = opt['random_reverse']
|
||||
logger = get_root_logger()
|
||||
logger.info(f'Random reverse is {self.random_reverse}.')
|
||||
|
||||
def __getitem__(self, index):
|
||||
if self.file_client is None:
|
||||
self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
|
||||
|
||||
# random reverse
|
||||
if self.random_reverse and random.random() < 0.5:
|
||||
self.neighbor_list.reverse()
|
||||
|
||||
scale = self.opt['scale']
|
||||
gt_size = self.opt['gt_size']
|
||||
key = self.keys[index]
|
||||
clip, seq = key.split('/') # key example: 00001/0001
|
||||
|
||||
# get the GT frame (im4.png)
|
||||
if self.is_lmdb:
|
||||
img_gt_path = f'{key}/im4'
|
||||
else:
|
||||
img_gt_path = self.gt_root / clip / seq / 'im4.png'
|
||||
img_bytes = self.file_client.get(img_gt_path, 'gt')
|
||||
img_gt = imfrombytes(img_bytes, float32=True)
|
||||
|
||||
# get the neighboring LQ frames
|
||||
img_lqs = []
|
||||
for neighbor in self.neighbor_list:
|
||||
if self.is_lmdb:
|
||||
img_lq_path = f'{clip}/{seq}/im{neighbor}'
|
||||
else:
|
||||
img_lq_path = self.lq_root / clip / seq / f'im{neighbor}.png'
|
||||
img_bytes = self.file_client.get(img_lq_path, 'lq')
|
||||
img_lq = imfrombytes(img_bytes, float32=True)
|
||||
img_lqs.append(img_lq)
|
||||
|
||||
# randomly crop
|
||||
img_gt, img_lqs = paired_random_crop(img_gt, img_lqs, gt_size, scale, img_gt_path)
|
||||
|
||||
# augmentation - flip, rotate
|
||||
img_lqs.append(img_gt)
|
||||
img_results = augment(img_lqs, self.opt['use_hflip'], self.opt['use_rot'])
|
||||
|
||||
img_results = img2tensor(img_results)
|
||||
img_lqs = torch.stack(img_results[0:-1], dim=0)
|
||||
img_gt = img_results[-1]
|
||||
|
||||
# img_lqs: (t, c, h, w)
|
||||
# img_gt: (c, h, w)
|
||||
# key: str
|
||||
return {'lq': img_lqs, 'gt': img_gt, 'key': key}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.keys)
|
||||
|
||||
|
||||
@DATASET_REGISTRY.register()
|
||||
class Vimeo90KRecurrentDataset(Vimeo90KDataset):
|
||||
|
||||
def __init__(self, opt):
|
||||
super(Vimeo90KRecurrentDataset, self).__init__(opt)
|
||||
|
||||
self.flip_sequence = opt['flip_sequence']
|
||||
self.neighbor_list = [1, 2, 3, 4, 5, 6, 7]
|
||||
|
||||
def __getitem__(self, index):
|
||||
if self.file_client is None:
|
||||
self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
|
||||
|
||||
# random reverse
|
||||
if self.random_reverse and random.random() < 0.5:
|
||||
self.neighbor_list.reverse()
|
||||
|
||||
scale = self.opt['scale']
|
||||
gt_size = self.opt['gt_size']
|
||||
key = self.keys[index]
|
||||
clip, seq = key.split('/') # key example: 00001/0001
|
||||
|
||||
# get the neighboring LQ and GT frames
|
||||
img_lqs = []
|
||||
img_gts = []
|
||||
for neighbor in self.neighbor_list:
|
||||
if self.is_lmdb:
|
||||
img_lq_path = f'{clip}/{seq}/im{neighbor}'
|
||||
img_gt_path = f'{clip}/{seq}/im{neighbor}'
|
||||
else:
|
||||
img_lq_path = self.lq_root / clip / seq / f'im{neighbor}.png'
|
||||
img_gt_path = self.gt_root / clip / seq / f'im{neighbor}.png'
|
||||
# LQ
|
||||
img_bytes = self.file_client.get(img_lq_path, 'lq')
|
||||
img_lq = imfrombytes(img_bytes, float32=True)
|
||||
# GT
|
||||
img_bytes = self.file_client.get(img_gt_path, 'gt')
|
||||
img_gt = imfrombytes(img_bytes, float32=True)
|
||||
|
||||
img_lqs.append(img_lq)
|
||||
img_gts.append(img_gt)
|
||||
|
||||
# randomly crop
|
||||
img_gts, img_lqs = paired_random_crop(img_gts, img_lqs, gt_size, scale, img_gt_path)
|
||||
|
||||
# augmentation - flip, rotate
|
||||
img_lqs.extend(img_gts)
|
||||
img_results = augment(img_lqs, self.opt['use_hflip'], self.opt['use_rot'])
|
||||
|
||||
img_results = img2tensor(img_results)
|
||||
img_lqs = torch.stack(img_results[:7], dim=0)
|
||||
img_gts = torch.stack(img_results[7:], dim=0)
|
||||
|
||||
if self.flip_sequence: # flip the sequence: 7 frames to 14 frames
|
||||
img_lqs = torch.cat([img_lqs, img_lqs.flip(0)], dim=0)
|
||||
img_gts = torch.cat([img_gts, img_gts.flip(0)], dim=0)
|
||||
|
||||
# img_lqs: (t, c, h, w)
|
||||
# img_gt: (c, h, w)
|
||||
# key: str
|
||||
return {'lq': img_lqs, 'gt': img_gts, 'key': key}
|
||||
|
||||
def __len__(self):
|
||||
return len(self.keys)
|
||||
@@ -1,47 +0,0 @@
|
||||
from .color_util import bgr2ycbcr, rgb2ycbcr, rgb2ycbcr_pt, ycbcr2bgr, ycbcr2rgb
|
||||
from .diffjpeg import DiffJPEG
|
||||
from .file_client import FileClient
|
||||
from .img_process_util import USMSharp, usm_sharp
|
||||
from .img_util import crop_border, imfrombytes, img2tensor, imwrite, tensor2img
|
||||
from .logger import AvgTimer, MessageLogger, get_env_info, get_root_logger, init_tb_logger, init_wandb_logger
|
||||
from .misc import check_resume, get_time_str, make_exp_dirs, mkdir_and_rename, scandir, set_random_seed, sizeof_fmt
|
||||
from .options import yaml_load
|
||||
|
||||
__all__ = [
|
||||
# color_util.py
|
||||
'bgr2ycbcr',
|
||||
'rgb2ycbcr',
|
||||
'rgb2ycbcr_pt',
|
||||
'ycbcr2bgr',
|
||||
'ycbcr2rgb',
|
||||
# file_client.py
|
||||
'FileClient',
|
||||
# img_util.py
|
||||
'img2tensor',
|
||||
'tensor2img',
|
||||
'imfrombytes',
|
||||
'imwrite',
|
||||
'crop_border',
|
||||
# logger.py
|
||||
'MessageLogger',
|
||||
'AvgTimer',
|
||||
'init_tb_logger',
|
||||
'init_wandb_logger',
|
||||
'get_root_logger',
|
||||
'get_env_info',
|
||||
# misc.py
|
||||
'set_random_seed',
|
||||
'get_time_str',
|
||||
'mkdir_and_rename',
|
||||
'make_exp_dirs',
|
||||
'scandir',
|
||||
'check_resume',
|
||||
'sizeof_fmt',
|
||||
# diffjpeg
|
||||
'DiffJPEG',
|
||||
# img_process_util
|
||||
'USMSharp',
|
||||
'usm_sharp',
|
||||
# options
|
||||
'yaml_load'
|
||||
]
|
||||
@@ -1,208 +0,0 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def rgb2ycbcr(img, y_only=False):
|
||||
"""Convert a RGB image to YCbCr image.
|
||||
|
||||
This function produces the same results as Matlab's `rgb2ycbcr` function.
|
||||
It implements the ITU-R BT.601 conversion for standard-definition
|
||||
television. See more details in
|
||||
https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
|
||||
|
||||
It differs from a similar function in cv2.cvtColor: `RGB <-> YCrCb`.
|
||||
In OpenCV, it implements a JPEG conversion. See more details in
|
||||
https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
|
||||
|
||||
Args:
|
||||
img (ndarray): The input image. It accepts:
|
||||
1. np.uint8 type with range [0, 255];
|
||||
2. np.float32 type with range [0, 1].
|
||||
y_only (bool): Whether to only return Y channel. Default: False.
|
||||
|
||||
Returns:
|
||||
ndarray: The converted YCbCr image. The output image has the same type
|
||||
and range as input image.
|
||||
"""
|
||||
img_type = img.dtype
|
||||
img = _convert_input_type_range(img)
|
||||
if y_only:
|
||||
out_img = np.dot(img, [65.481, 128.553, 24.966]) + 16.0
|
||||
else:
|
||||
out_img = np.matmul(
|
||||
img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786], [24.966, 112.0, -18.214]]) + [16, 128, 128]
|
||||
out_img = _convert_output_type_range(out_img, img_type)
|
||||
return out_img
|
||||
|
||||
|
||||
def bgr2ycbcr(img, y_only=False):
|
||||
"""Convert a BGR image to YCbCr image.
|
||||
|
||||
The bgr version of rgb2ycbcr.
|
||||
It implements the ITU-R BT.601 conversion for standard-definition
|
||||
television. See more details in
|
||||
https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
|
||||
|
||||
It differs from a similar function in cv2.cvtColor: `BGR <-> YCrCb`.
|
||||
In OpenCV, it implements a JPEG conversion. See more details in
|
||||
https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
|
||||
|
||||
Args:
|
||||
img (ndarray): The input image. It accepts:
|
||||
1. np.uint8 type with range [0, 255];
|
||||
2. np.float32 type with range [0, 1].
|
||||
y_only (bool): Whether to only return Y channel. Default: False.
|
||||
|
||||
Returns:
|
||||
ndarray: The converted YCbCr image. The output image has the same type
|
||||
and range as input image.
|
||||
"""
|
||||
img_type = img.dtype
|
||||
img = _convert_input_type_range(img)
|
||||
if y_only:
|
||||
out_img = np.dot(img, [24.966, 128.553, 65.481]) + 16.0
|
||||
else:
|
||||
out_img = np.matmul(
|
||||
img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786], [65.481, -37.797, 112.0]]) + [16, 128, 128]
|
||||
out_img = _convert_output_type_range(out_img, img_type)
|
||||
return out_img
|
||||
|
||||
|
||||
def ycbcr2rgb(img):
|
||||
"""Convert a YCbCr image to RGB image.
|
||||
|
||||
This function produces the same results as Matlab's ycbcr2rgb function.
|
||||
It implements the ITU-R BT.601 conversion for standard-definition
|
||||
television. See more details in
|
||||
https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
|
||||
|
||||
It differs from a similar function in cv2.cvtColor: `YCrCb <-> RGB`.
|
||||
In OpenCV, it implements a JPEG conversion. See more details in
|
||||
https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
|
||||
|
||||
Args:
|
||||
img (ndarray): The input image. It accepts:
|
||||
1. np.uint8 type with range [0, 255];
|
||||
2. np.float32 type with range [0, 1].
|
||||
|
||||
Returns:
|
||||
ndarray: The converted RGB image. The output image has the same type
|
||||
and range as input image.
|
||||
"""
|
||||
img_type = img.dtype
|
||||
img = _convert_input_type_range(img) * 255
|
||||
out_img = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0, -0.00153632, 0.00791071],
|
||||
[0.00625893, -0.00318811, 0]]) * 255.0 + [-222.921, 135.576, -276.836] # noqa: E126
|
||||
out_img = _convert_output_type_range(out_img, img_type)
|
||||
return out_img
|
||||
|
||||
|
||||
def ycbcr2bgr(img):
|
||||
"""Convert a YCbCr image to BGR image.
|
||||
|
||||
The bgr version of ycbcr2rgb.
|
||||
It implements the ITU-R BT.601 conversion for standard-definition
|
||||
television. See more details in
|
||||
https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
|
||||
|
||||
It differs from a similar function in cv2.cvtColor: `YCrCb <-> BGR`.
|
||||
In OpenCV, it implements a JPEG conversion. See more details in
|
||||
https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
|
||||
|
||||
Args:
|
||||
img (ndarray): The input image. It accepts:
|
||||
1. np.uint8 type with range [0, 255];
|
||||
2. np.float32 type with range [0, 1].
|
||||
|
||||
Returns:
|
||||
ndarray: The converted BGR image. The output image has the same type
|
||||
and range as input image.
|
||||
"""
|
||||
img_type = img.dtype
|
||||
img = _convert_input_type_range(img) * 255
|
||||
out_img = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0.00791071, -0.00153632, 0],
|
||||
[0, -0.00318811, 0.00625893]]) * 255.0 + [-276.836, 135.576, -222.921] # noqa: E126
|
||||
out_img = _convert_output_type_range(out_img, img_type)
|
||||
return out_img
|
||||
|
||||
|
||||
def _convert_input_type_range(img):
|
||||
"""Convert the type and range of the input image.
|
||||
|
||||
It converts the input image to np.float32 type and range of [0, 1].
|
||||
It is mainly used for pre-processing the input image in colorspace
|
||||
conversion functions such as rgb2ycbcr and ycbcr2rgb.
|
||||
|
||||
Args:
|
||||
img (ndarray): The input image. It accepts:
|
||||
1. np.uint8 type with range [0, 255];
|
||||
2. np.float32 type with range [0, 1].
|
||||
|
||||
Returns:
|
||||
(ndarray): The converted image with type of np.float32 and range of
|
||||
[0, 1].
|
||||
"""
|
||||
img_type = img.dtype
|
||||
img = img.astype(np.float32)
|
||||
if img_type == np.float32:
|
||||
pass
|
||||
elif img_type == np.uint8:
|
||||
img /= 255.
|
||||
else:
|
||||
raise TypeError(f'The img type should be np.float32 or np.uint8, but got {img_type}')
|
||||
return img
|
||||
|
||||
|
||||
def _convert_output_type_range(img, dst_type):
|
||||
"""Convert the type and range of the image according to dst_type.
|
||||
|
||||
It converts the image to desired type and range. If `dst_type` is np.uint8,
|
||||
images will be converted to np.uint8 type with range [0, 255]. If
|
||||
`dst_type` is np.float32, it converts the image to np.float32 type with
|
||||
range [0, 1].
|
||||
It is mainly used for post-processing images in colorspace conversion
|
||||
functions such as rgb2ycbcr and ycbcr2rgb.
|
||||
|
||||
Args:
|
||||
img (ndarray): The image to be converted with np.float32 type and
|
||||
range [0, 255].
|
||||
dst_type (np.uint8 | np.float32): If dst_type is np.uint8, it
|
||||
converts the image to np.uint8 type with range [0, 255]. If
|
||||
dst_type is np.float32, it converts the image to np.float32 type
|
||||
with range [0, 1].
|
||||
|
||||
Returns:
|
||||
(ndarray): The converted image with desired type and range.
|
||||
"""
|
||||
if dst_type not in (np.uint8, np.float32):
|
||||
raise TypeError(f'The dst_type should be np.float32 or np.uint8, but got {dst_type}')
|
||||
if dst_type == np.uint8:
|
||||
img = img.round()
|
||||
else:
|
||||
img /= 255.
|
||||
return img.astype(dst_type)
|
||||
|
||||
|
||||
def rgb2ycbcr_pt(img, y_only=False):
|
||||
"""Convert RGB images to YCbCr images (PyTorch version).
|
||||
|
||||
It implements the ITU-R BT.601 conversion for standard-definition television. See more details in
|
||||
https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
|
||||
|
||||
Args:
|
||||
img (Tensor): Images with shape (n, 3, h, w), the range [0, 1], float, RGB format.
|
||||
y_only (bool): Whether to only return Y channel. Default: False.
|
||||
|
||||
Returns:
|
||||
(Tensor): converted images with the shape (n, 3/1, h, w), the range [0, 1], float.
|
||||
"""
|
||||
if y_only:
|
||||
weight = torch.tensor([[65.481], [128.553], [24.966]]).to(img)
|
||||
out_img = torch.matmul(img.permute(0, 2, 3, 1), weight).permute(0, 3, 1, 2) + 16.0
|
||||
else:
|
||||
weight = torch.tensor([[65.481, -37.797, 112.0], [128.553, -74.203, -93.786], [24.966, 112.0, -18.214]]).to(img)
|
||||
bias = torch.tensor([16, 128, 128]).view(1, 3, 1, 1).to(img)
|
||||
out_img = torch.matmul(img.permute(0, 2, 3, 1), weight).permute(0, 3, 1, 2) + bias
|
||||
|
||||
out_img = out_img / 255.
|
||||
return out_img
|
||||
@@ -1,515 +0,0 @@
|
||||
"""
|
||||
Modified from https://github.com/mlomnitz/DiffJPEG
|
||||
|
||||
For images not divisible by 8
|
||||
https://dsp.stackexchange.com/questions/35339/jpeg-dct-padding/35343#35343
|
||||
"""
|
||||
import itertools
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
# ------------------------ utils ------------------------#
|
||||
y_table = np.array(
|
||||
[[16, 11, 10, 16, 24, 40, 51, 61], [12, 12, 14, 19, 26, 58, 60, 55], [14, 13, 16, 24, 40, 57, 69, 56],
|
||||
[14, 17, 22, 29, 51, 87, 80, 62], [18, 22, 37, 56, 68, 109, 103, 77], [24, 35, 55, 64, 81, 104, 113, 92],
|
||||
[49, 64, 78, 87, 103, 121, 120, 101], [72, 92, 95, 98, 112, 100, 103, 99]],
|
||||
dtype=np.float32).T
|
||||
y_table = nn.Parameter(torch.from_numpy(y_table))
|
||||
c_table = np.empty((8, 8), dtype=np.float32)
|
||||
c_table.fill(99)
|
||||
c_table[:4, :4] = np.array([[17, 18, 24, 47], [18, 21, 26, 66], [24, 26, 56, 99], [47, 66, 99, 99]]).T
|
||||
c_table = nn.Parameter(torch.from_numpy(c_table))
|
||||
|
||||
|
||||
def diff_round(x):
|
||||
""" Differentiable rounding function
|
||||
"""
|
||||
return torch.round(x) + (x - torch.round(x))**3
|
||||
|
||||
|
||||
def quality_to_factor(quality):
|
||||
""" Calculate factor corresponding to quality
|
||||
|
||||
Args:
|
||||
quality(float): Quality for jpeg compression.
|
||||
|
||||
Returns:
|
||||
float: Compression factor.
|
||||
"""
|
||||
if quality < 50:
|
||||
quality = 5000. / quality
|
||||
else:
|
||||
quality = 200. - quality * 2
|
||||
return quality / 100.
|
||||
|
||||
|
||||
# ------------------------ compression ------------------------#
|
||||
class RGB2YCbCrJpeg(nn.Module):
|
||||
""" Converts RGB image to YCbCr
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(RGB2YCbCrJpeg, self).__init__()
|
||||
matrix = np.array([[0.299, 0.587, 0.114], [-0.168736, -0.331264, 0.5], [0.5, -0.418688, -0.081312]],
|
||||
dtype=np.float32).T
|
||||
self.shift = nn.Parameter(torch.tensor([0., 128., 128.]))
|
||||
self.matrix = nn.Parameter(torch.from_numpy(matrix))
|
||||
|
||||
def forward(self, image):
|
||||
"""
|
||||
Args:
|
||||
image(Tensor): batch x 3 x height x width
|
||||
|
||||
Returns:
|
||||
Tensor: batch x height x width x 3
|
||||
"""
|
||||
image = image.permute(0, 2, 3, 1)
|
||||
result = torch.tensordot(image, self.matrix, dims=1) + self.shift
|
||||
return result.view(image.shape)
|
||||
|
||||
|
||||
class ChromaSubsampling(nn.Module):
|
||||
""" Chroma subsampling on CbCr channels
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(ChromaSubsampling, self).__init__()
|
||||
|
||||
def forward(self, image):
|
||||
"""
|
||||
Args:
|
||||
image(tensor): batch x height x width x 3
|
||||
|
||||
Returns:
|
||||
y(tensor): batch x height x width
|
||||
cb(tensor): batch x height/2 x width/2
|
||||
cr(tensor): batch x height/2 x width/2
|
||||
"""
|
||||
image_2 = image.permute(0, 3, 1, 2).clone()
|
||||
cb = F.avg_pool2d(image_2[:, 1, :, :].unsqueeze(1), kernel_size=2, stride=(2, 2), count_include_pad=False)
|
||||
cr = F.avg_pool2d(image_2[:, 2, :, :].unsqueeze(1), kernel_size=2, stride=(2, 2), count_include_pad=False)
|
||||
cb = cb.permute(0, 2, 3, 1)
|
||||
cr = cr.permute(0, 2, 3, 1)
|
||||
return image[:, :, :, 0], cb.squeeze(3), cr.squeeze(3)
|
||||
|
||||
|
||||
class BlockSplitting(nn.Module):
|
||||
""" Splitting image into patches
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(BlockSplitting, self).__init__()
|
||||
self.k = 8
|
||||
|
||||
def forward(self, image):
|
||||
"""
|
||||
Args:
|
||||
image(tensor): batch x height x width
|
||||
|
||||
Returns:
|
||||
Tensor: batch x h*w/64 x h x w
|
||||
"""
|
||||
height, _ = image.shape[1:3]
|
||||
batch_size = image.shape[0]
|
||||
image_reshaped = image.view(batch_size, height // self.k, self.k, -1, self.k)
|
||||
image_transposed = image_reshaped.permute(0, 1, 3, 2, 4)
|
||||
return image_transposed.contiguous().view(batch_size, -1, self.k, self.k)
|
||||
|
||||
|
||||
class DCT8x8(nn.Module):
|
||||
""" Discrete Cosine Transformation
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(DCT8x8, self).__init__()
|
||||
tensor = np.zeros((8, 8, 8, 8), dtype=np.float32)
|
||||
for x, y, u, v in itertools.product(range(8), repeat=4):
|
||||
tensor[x, y, u, v] = np.cos((2 * x + 1) * u * np.pi / 16) * np.cos((2 * y + 1) * v * np.pi / 16)
|
||||
alpha = np.array([1. / np.sqrt(2)] + [1] * 7)
|
||||
self.tensor = nn.Parameter(torch.from_numpy(tensor).float())
|
||||
self.scale = nn.Parameter(torch.from_numpy(np.outer(alpha, alpha) * 0.25).float())
|
||||
|
||||
def forward(self, image):
|
||||
"""
|
||||
Args:
|
||||
image(tensor): batch x height x width
|
||||
|
||||
Returns:
|
||||
Tensor: batch x height x width
|
||||
"""
|
||||
image = image - 128
|
||||
result = self.scale * torch.tensordot(image, self.tensor, dims=2)
|
||||
result.view(image.shape)
|
||||
return result
|
||||
|
||||
|
||||
class YQuantize(nn.Module):
|
||||
""" JPEG Quantization for Y channel
|
||||
|
||||
Args:
|
||||
rounding(function): rounding function to use
|
||||
"""
|
||||
|
||||
def __init__(self, rounding):
|
||||
super(YQuantize, self).__init__()
|
||||
self.rounding = rounding
|
||||
self.y_table = y_table
|
||||
|
||||
def forward(self, image, factor=1):
|
||||
"""
|
||||
Args:
|
||||
image(tensor): batch x height x width
|
||||
|
||||
Returns:
|
||||
Tensor: batch x height x width
|
||||
"""
|
||||
if isinstance(factor, (int, float)):
|
||||
image = image.float() / (self.y_table * factor)
|
||||
else:
|
||||
b = factor.size(0)
|
||||
table = self.y_table.expand(b, 1, 8, 8) * factor.view(b, 1, 1, 1)
|
||||
image = image.float() / table
|
||||
image = self.rounding(image)
|
||||
return image
|
||||
|
||||
|
||||
class CQuantize(nn.Module):
|
||||
""" JPEG Quantization for CbCr channels
|
||||
|
||||
Args:
|
||||
rounding(function): rounding function to use
|
||||
"""
|
||||
|
||||
def __init__(self, rounding):
|
||||
super(CQuantize, self).__init__()
|
||||
self.rounding = rounding
|
||||
self.c_table = c_table
|
||||
|
||||
def forward(self, image, factor=1):
|
||||
"""
|
||||
Args:
|
||||
image(tensor): batch x height x width
|
||||
|
||||
Returns:
|
||||
Tensor: batch x height x width
|
||||
"""
|
||||
if isinstance(factor, (int, float)):
|
||||
image = image.float() / (self.c_table * factor)
|
||||
else:
|
||||
b = factor.size(0)
|
||||
table = self.c_table.expand(b, 1, 8, 8) * factor.view(b, 1, 1, 1)
|
||||
image = image.float() / table
|
||||
image = self.rounding(image)
|
||||
return image
|
||||
|
||||
|
||||
class CompressJpeg(nn.Module):
|
||||
"""Full JPEG compression algorithm
|
||||
|
||||
Args:
|
||||
rounding(function): rounding function to use
|
||||
"""
|
||||
|
||||
def __init__(self, rounding=torch.round):
|
||||
super(CompressJpeg, self).__init__()
|
||||
self.l1 = nn.Sequential(RGB2YCbCrJpeg(), ChromaSubsampling())
|
||||
self.l2 = nn.Sequential(BlockSplitting(), DCT8x8())
|
||||
self.c_quantize = CQuantize(rounding=rounding)
|
||||
self.y_quantize = YQuantize(rounding=rounding)
|
||||
|
||||
def forward(self, image, factor=1):
|
||||
"""
|
||||
Args:
|
||||
image(tensor): batch x 3 x height x width
|
||||
|
||||
Returns:
|
||||
dict(tensor): Compressed tensor with batch x h*w/64 x 8 x 8.
|
||||
"""
|
||||
y, cb, cr = self.l1(image * 255)
|
||||
components = {'y': y, 'cb': cb, 'cr': cr}
|
||||
for k in components.keys():
|
||||
comp = self.l2(components[k])
|
||||
if k in ('cb', 'cr'):
|
||||
comp = self.c_quantize(comp, factor=factor)
|
||||
else:
|
||||
comp = self.y_quantize(comp, factor=factor)
|
||||
|
||||
components[k] = comp
|
||||
|
||||
return components['y'], components['cb'], components['cr']
|
||||
|
||||
|
||||
# ------------------------ decompression ------------------------#
|
||||
|
||||
|
||||
class YDequantize(nn.Module):
|
||||
"""Dequantize Y channel
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(YDequantize, self).__init__()
|
||||
self.y_table = y_table
|
||||
|
||||
def forward(self, image, factor=1):
|
||||
"""
|
||||
Args:
|
||||
image(tensor): batch x height x width
|
||||
|
||||
Returns:
|
||||
Tensor: batch x height x width
|
||||
"""
|
||||
if isinstance(factor, (int, float)):
|
||||
out = image * (self.y_table * factor)
|
||||
else:
|
||||
b = factor.size(0)
|
||||
table = self.y_table.expand(b, 1, 8, 8) * factor.view(b, 1, 1, 1)
|
||||
out = image * table
|
||||
return out
|
||||
|
||||
|
||||
class CDequantize(nn.Module):
|
||||
"""Dequantize CbCr channel
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(CDequantize, self).__init__()
|
||||
self.c_table = c_table
|
||||
|
||||
def forward(self, image, factor=1):
|
||||
"""
|
||||
Args:
|
||||
image(tensor): batch x height x width
|
||||
|
||||
Returns:
|
||||
Tensor: batch x height x width
|
||||
"""
|
||||
if isinstance(factor, (int, float)):
|
||||
out = image * (self.c_table * factor)
|
||||
else:
|
||||
b = factor.size(0)
|
||||
table = self.c_table.expand(b, 1, 8, 8) * factor.view(b, 1, 1, 1)
|
||||
out = image * table
|
||||
return out
|
||||
|
||||
|
||||
class iDCT8x8(nn.Module):
|
||||
"""Inverse discrete Cosine Transformation
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(iDCT8x8, self).__init__()
|
||||
alpha = np.array([1. / np.sqrt(2)] + [1] * 7)
|
||||
self.alpha = nn.Parameter(torch.from_numpy(np.outer(alpha, alpha)).float())
|
||||
tensor = np.zeros((8, 8, 8, 8), dtype=np.float32)
|
||||
for x, y, u, v in itertools.product(range(8), repeat=4):
|
||||
tensor[x, y, u, v] = np.cos((2 * u + 1) * x * np.pi / 16) * np.cos((2 * v + 1) * y * np.pi / 16)
|
||||
self.tensor = nn.Parameter(torch.from_numpy(tensor).float())
|
||||
|
||||
def forward(self, image):
|
||||
"""
|
||||
Args:
|
||||
image(tensor): batch x height x width
|
||||
|
||||
Returns:
|
||||
Tensor: batch x height x width
|
||||
"""
|
||||
image = image * self.alpha
|
||||
result = 0.25 * torch.tensordot(image, self.tensor, dims=2) + 128
|
||||
result.view(image.shape)
|
||||
return result
|
||||
|
||||
|
||||
class BlockMerging(nn.Module):
|
||||
"""Merge patches into image
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(BlockMerging, self).__init__()
|
||||
|
||||
def forward(self, patches, height, width):
|
||||
"""
|
||||
Args:
|
||||
patches(tensor) batch x height*width/64, height x width
|
||||
height(int)
|
||||
width(int)
|
||||
|
||||
Returns:
|
||||
Tensor: batch x height x width
|
||||
"""
|
||||
k = 8
|
||||
batch_size = patches.shape[0]
|
||||
image_reshaped = patches.view(batch_size, height // k, width // k, k, k)
|
||||
image_transposed = image_reshaped.permute(0, 1, 3, 2, 4)
|
||||
return image_transposed.contiguous().view(batch_size, height, width)
|
||||
|
||||
|
||||
class ChromaUpsampling(nn.Module):
|
||||
"""Upsample chroma layers
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(ChromaUpsampling, self).__init__()
|
||||
|
||||
def forward(self, y, cb, cr):
|
||||
"""
|
||||
Args:
|
||||
y(tensor): y channel image
|
||||
cb(tensor): cb channel
|
||||
cr(tensor): cr channel
|
||||
|
||||
Returns:
|
||||
Tensor: batch x height x width x 3
|
||||
"""
|
||||
|
||||
def repeat(x, k=2):
|
||||
height, width = x.shape[1:3]
|
||||
x = x.unsqueeze(-1)
|
||||
x = x.repeat(1, 1, k, k)
|
||||
x = x.view(-1, height * k, width * k)
|
||||
return x
|
||||
|
||||
cb = repeat(cb)
|
||||
cr = repeat(cr)
|
||||
return torch.cat([y.unsqueeze(3), cb.unsqueeze(3), cr.unsqueeze(3)], dim=3)
|
||||
|
||||
|
||||
class YCbCr2RGBJpeg(nn.Module):
|
||||
"""Converts YCbCr image to RGB JPEG
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super(YCbCr2RGBJpeg, self).__init__()
|
||||
|
||||
matrix = np.array([[1., 0., 1.402], [1, -0.344136, -0.714136], [1, 1.772, 0]], dtype=np.float32).T
|
||||
self.shift = nn.Parameter(torch.tensor([0, -128., -128.]))
|
||||
self.matrix = nn.Parameter(torch.from_numpy(matrix))
|
||||
|
||||
def forward(self, image):
|
||||
"""
|
||||
Args:
|
||||
image(tensor): batch x height x width x 3
|
||||
|
||||
Returns:
|
||||
Tensor: batch x 3 x height x width
|
||||
"""
|
||||
result = torch.tensordot(image + self.shift, self.matrix, dims=1)
|
||||
return result.view(image.shape).permute(0, 3, 1, 2)
|
||||
|
||||
|
||||
class DeCompressJpeg(nn.Module):
|
||||
"""Full JPEG decompression algorithm
|
||||
|
||||
Args:
|
||||
rounding(function): rounding function to use
|
||||
"""
|
||||
|
||||
def __init__(self, rounding=torch.round):
|
||||
super(DeCompressJpeg, self).__init__()
|
||||
self.c_dequantize = CDequantize()
|
||||
self.y_dequantize = YDequantize()
|
||||
self.idct = iDCT8x8()
|
||||
self.merging = BlockMerging()
|
||||
self.chroma = ChromaUpsampling()
|
||||
self.colors = YCbCr2RGBJpeg()
|
||||
|
||||
def forward(self, y, cb, cr, imgh, imgw, factor=1):
|
||||
"""
|
||||
Args:
|
||||
compressed(dict(tensor)): batch x h*w/64 x 8 x 8
|
||||
imgh(int)
|
||||
imgw(int)
|
||||
factor(float)
|
||||
|
||||
Returns:
|
||||
Tensor: batch x 3 x height x width
|
||||
"""
|
||||
components = {'y': y, 'cb': cb, 'cr': cr}
|
||||
for k in components.keys():
|
||||
if k in ('cb', 'cr'):
|
||||
comp = self.c_dequantize(components[k], factor=factor)
|
||||
height, width = int(imgh / 2), int(imgw / 2)
|
||||
else:
|
||||
comp = self.y_dequantize(components[k], factor=factor)
|
||||
height, width = imgh, imgw
|
||||
comp = self.idct(comp)
|
||||
components[k] = self.merging(comp, height, width)
|
||||
#
|
||||
image = self.chroma(components['y'], components['cb'], components['cr'])
|
||||
image = self.colors(image)
|
||||
|
||||
image = torch.min(255 * torch.ones_like(image), torch.max(torch.zeros_like(image), image))
|
||||
return image / 255
|
||||
|
||||
|
||||
# ------------------------ main DiffJPEG ------------------------ #
|
||||
|
||||
|
||||
class DiffJPEG(nn.Module):
|
||||
"""This JPEG algorithm result is slightly different from cv2.
|
||||
DiffJPEG supports batch processing.
|
||||
|
||||
Args:
|
||||
differentiable(bool): If True, uses custom differentiable rounding function, if False, uses standard torch.round
|
||||
"""
|
||||
|
||||
def __init__(self, differentiable=True):
|
||||
super(DiffJPEG, self).__init__()
|
||||
if differentiable:
|
||||
rounding = diff_round
|
||||
else:
|
||||
rounding = torch.round
|
||||
|
||||
self.compress = CompressJpeg(rounding=rounding)
|
||||
self.decompress = DeCompressJpeg(rounding=rounding)
|
||||
|
||||
def forward(self, x, quality):
|
||||
"""
|
||||
Args:
|
||||
x (Tensor): Input image, bchw, rgb, [0, 1]
|
||||
quality(float): Quality factor for jpeg compression scheme.
|
||||
"""
|
||||
factor = quality
|
||||
if isinstance(factor, (int, float)):
|
||||
factor = quality_to_factor(factor)
|
||||
else:
|
||||
for i in range(factor.size(0)):
|
||||
factor[i] = quality_to_factor(factor[i])
|
||||
h, w = x.size()[-2:]
|
||||
h_pad, w_pad = 0, 0
|
||||
# why should use 16
|
||||
if h % 16 != 0:
|
||||
h_pad = 16 - h % 16
|
||||
if w % 16 != 0:
|
||||
w_pad = 16 - w % 16
|
||||
x = F.pad(x, (0, w_pad, 0, h_pad), mode='constant', value=0)
|
||||
|
||||
y, cb, cr = self.compress(x, factor=factor)
|
||||
recovered = self.decompress(y, cb, cr, (h + h_pad), (w + w_pad), factor=factor)
|
||||
recovered = recovered[:, :, 0:h, 0:w]
|
||||
return recovered
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
import cv2
|
||||
|
||||
from basicsr.utils import img2tensor, tensor2img
|
||||
|
||||
img_gt = cv2.imread('test.png') / 255.
|
||||
|
||||
# -------------- cv2 -------------- #
|
||||
encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 20]
|
||||
_, encimg = cv2.imencode('.jpg', img_gt * 255., encode_param)
|
||||
img_lq = np.float32(cv2.imdecode(encimg, 1))
|
||||
cv2.imwrite('cv2_JPEG_20.png', img_lq)
|
||||
|
||||
# -------------- DiffJPEG -------------- #
|
||||
jpeger = DiffJPEG(differentiable=False).cuda()
|
||||
img_gt = img2tensor(img_gt)
|
||||
img_gt = torch.stack([img_gt, img_gt]).cuda()
|
||||
quality = img_gt.new_tensor([20, 40])
|
||||
out = jpeger(img_gt, quality=quality)
|
||||
|
||||
cv2.imwrite('pt_JPEG_20.png', tensor2img(out[0]))
|
||||
cv2.imwrite('pt_JPEG_40.png', tensor2img(out[1]))
|
||||
@@ -1,82 +0,0 @@
|
||||
# Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/runner/dist_utils.py # noqa: E501
|
||||
import functools
|
||||
import os
|
||||
import subprocess
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
|
||||
def init_dist(launcher, backend='nccl', **kwargs):
|
||||
if mp.get_start_method(allow_none=True) is None:
|
||||
mp.set_start_method('spawn')
|
||||
if launcher == 'pytorch':
|
||||
_init_dist_pytorch(backend, **kwargs)
|
||||
elif launcher == 'slurm':
|
||||
_init_dist_slurm(backend, **kwargs)
|
||||
else:
|
||||
raise ValueError(f'Invalid launcher type: {launcher}')
|
||||
|
||||
|
||||
def _init_dist_pytorch(backend, **kwargs):
|
||||
rank = int(os.environ['RANK'])
|
||||
num_gpus = torch.cuda.device_count()
|
||||
torch.cuda.set_device(rank % num_gpus)
|
||||
dist.init_process_group(backend=backend, **kwargs)
|
||||
|
||||
|
||||
def _init_dist_slurm(backend, port=None):
|
||||
"""Initialize slurm distributed training environment.
|
||||
|
||||
If argument ``port`` is not specified, then the master port will be system
|
||||
environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system
|
||||
environment variable, then a default port ``29500`` will be used.
|
||||
|
||||
Args:
|
||||
backend (str): Backend of torch.distributed.
|
||||
port (int, optional): Master port. Defaults to None.
|
||||
"""
|
||||
proc_id = int(os.environ['SLURM_PROCID'])
|
||||
ntasks = int(os.environ['SLURM_NTASKS'])
|
||||
node_list = os.environ['SLURM_NODELIST']
|
||||
num_gpus = torch.cuda.device_count()
|
||||
torch.cuda.set_device(proc_id % num_gpus)
|
||||
addr = subprocess.getoutput(f'scontrol show hostname {node_list} | head -n1')
|
||||
# specify master port
|
||||
if port is not None:
|
||||
os.environ['MASTER_PORT'] = str(port)
|
||||
elif 'MASTER_PORT' in os.environ:
|
||||
pass # use MASTER_PORT in the environment variable
|
||||
else:
|
||||
# 29500 is torch.distributed default port
|
||||
os.environ['MASTER_PORT'] = '29500'
|
||||
os.environ['MASTER_ADDR'] = addr
|
||||
os.environ['WORLD_SIZE'] = str(ntasks)
|
||||
os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)
|
||||
os.environ['RANK'] = str(proc_id)
|
||||
dist.init_process_group(backend=backend)
|
||||
|
||||
|
||||
def get_dist_info():
|
||||
if dist.is_available():
|
||||
initialized = dist.is_initialized()
|
||||
else:
|
||||
initialized = False
|
||||
if initialized:
|
||||
rank = dist.get_rank()
|
||||
world_size = dist.get_world_size()
|
||||
else:
|
||||
rank = 0
|
||||
world_size = 1
|
||||
return rank, world_size
|
||||
|
||||
|
||||
def master_only(func):
|
||||
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
rank, _ = get_dist_info()
|
||||
if rank == 0:
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
@@ -1,98 +0,0 @@
|
||||
import math
|
||||
import os
|
||||
import requests
|
||||
from torch.hub import download_url_to_file, get_dir
|
||||
from tqdm import tqdm
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from .misc import sizeof_fmt
|
||||
|
||||
|
||||
def download_file_from_google_drive(file_id, save_path):
|
||||
"""Download files from google drive.
|
||||
|
||||
Reference: https://stackoverflow.com/questions/25010369/wget-curl-large-file-from-google-drive
|
||||
|
||||
Args:
|
||||
file_id (str): File id.
|
||||
save_path (str): Save path.
|
||||
"""
|
||||
|
||||
session = requests.Session()
|
||||
URL = 'https://docs.google.com/uc?export=download'
|
||||
params = {'id': file_id}
|
||||
|
||||
response = session.get(URL, params=params, stream=True)
|
||||
token = get_confirm_token(response)
|
||||
if token:
|
||||
params['confirm'] = token
|
||||
response = session.get(URL, params=params, stream=True)
|
||||
|
||||
# get file size
|
||||
response_file_size = session.get(URL, params=params, stream=True, headers={'Range': 'bytes=0-2'})
|
||||
if 'Content-Range' in response_file_size.headers:
|
||||
file_size = int(response_file_size.headers['Content-Range'].split('/')[1])
|
||||
else:
|
||||
file_size = None
|
||||
|
||||
save_response_content(response, save_path, file_size)
|
||||
|
||||
|
||||
def get_confirm_token(response):
|
||||
for key, value in response.cookies.items():
|
||||
if key.startswith('download_warning'):
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def save_response_content(response, destination, file_size=None, chunk_size=32768):
|
||||
if file_size is not None:
|
||||
pbar = tqdm(total=math.ceil(file_size / chunk_size), unit='chunk')
|
||||
|
||||
readable_file_size = sizeof_fmt(file_size)
|
||||
else:
|
||||
pbar = None
|
||||
|
||||
with open(destination, 'wb') as f:
|
||||
downloaded_size = 0
|
||||
for chunk in response.iter_content(chunk_size):
|
||||
downloaded_size += chunk_size
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
pbar.set_description(f'Download {sizeof_fmt(downloaded_size)} / {readable_file_size}')
|
||||
if chunk: # filter out keep-alive new chunks
|
||||
f.write(chunk)
|
||||
if pbar is not None:
|
||||
pbar.close()
|
||||
|
||||
|
||||
def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
|
||||
"""Load file form http url, will download models if necessary.
|
||||
|
||||
Reference: https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
|
||||
|
||||
Args:
|
||||
url (str): URL to be downloaded.
|
||||
model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir.
|
||||
Default: None.
|
||||
progress (bool): Whether to show the download progress. Default: True.
|
||||
file_name (str): The downloaded file name. If None, use the file name in the url. Default: None.
|
||||
|
||||
Returns:
|
||||
str: The path to the downloaded file.
|
||||
"""
|
||||
if model_dir is None: # use the pytorch hub_dir
|
||||
hub_dir = get_dir()
|
||||
model_dir = os.path.join(hub_dir, 'checkpoints')
|
||||
|
||||
os.makedirs(model_dir, exist_ok=True)
|
||||
|
||||
parts = urlparse(url)
|
||||
filename = os.path.basename(parts.path)
|
||||
if file_name is not None:
|
||||
filename = file_name
|
||||
cached_file = os.path.abspath(os.path.join(model_dir, filename))
|
||||
if not os.path.exists(cached_file):
|
||||
print(f'Downloading: "{url}" to {cached_file}\n')
|
||||
download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
|
||||
return cached_file
|
||||
@@ -1,167 +0,0 @@
|
||||
# Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py # noqa: E501
|
||||
from abc import ABCMeta, abstractmethod
|
||||
|
||||
|
||||
class BaseStorageBackend(metaclass=ABCMeta):
|
||||
"""Abstract class of storage backends.
|
||||
|
||||
All backends need to implement two apis: ``get()`` and ``get_text()``.
|
||||
``get()`` reads the file as a byte stream and ``get_text()`` reads the file
|
||||
as texts.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get(self, filepath):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_text(self, filepath):
|
||||
pass
|
||||
|
||||
|
||||
class MemcachedBackend(BaseStorageBackend):
|
||||
"""Memcached storage backend.
|
||||
|
||||
Attributes:
|
||||
server_list_cfg (str): Config file for memcached server list.
|
||||
client_cfg (str): Config file for memcached client.
|
||||
sys_path (str | None): Additional path to be appended to `sys.path`.
|
||||
Default: None.
|
||||
"""
|
||||
|
||||
def __init__(self, server_list_cfg, client_cfg, sys_path=None):
|
||||
if sys_path is not None:
|
||||
import sys
|
||||
sys.path.append(sys_path)
|
||||
try:
|
||||
import mc
|
||||
except ImportError:
|
||||
raise ImportError('Please install memcached to enable MemcachedBackend.')
|
||||
|
||||
self.server_list_cfg = server_list_cfg
|
||||
self.client_cfg = client_cfg
|
||||
self._client = mc.MemcachedClient.GetInstance(self.server_list_cfg, self.client_cfg)
|
||||
# mc.pyvector servers as a point which points to a memory cache
|
||||
self._mc_buffer = mc.pyvector()
|
||||
|
||||
def get(self, filepath):
|
||||
filepath = str(filepath)
|
||||
import mc
|
||||
self._client.Get(filepath, self._mc_buffer)
|
||||
value_buf = mc.ConvertBuffer(self._mc_buffer)
|
||||
return value_buf
|
||||
|
||||
def get_text(self, filepath):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class HardDiskBackend(BaseStorageBackend):
|
||||
"""Raw hard disks storage backend."""
|
||||
|
||||
def get(self, filepath):
|
||||
filepath = str(filepath)
|
||||
with open(filepath, 'rb') as f:
|
||||
value_buf = f.read()
|
||||
return value_buf
|
||||
|
||||
def get_text(self, filepath):
|
||||
filepath = str(filepath)
|
||||
with open(filepath, 'r') as f:
|
||||
value_buf = f.read()
|
||||
return value_buf
|
||||
|
||||
|
||||
class LmdbBackend(BaseStorageBackend):
|
||||
"""Lmdb storage backend.
|
||||
|
||||
Args:
|
||||
db_paths (str | list[str]): Lmdb database paths.
|
||||
client_keys (str | list[str]): Lmdb client keys. Default: 'default'.
|
||||
readonly (bool, optional): Lmdb environment parameter. If True,
|
||||
disallow any write operations. Default: True.
|
||||
lock (bool, optional): Lmdb environment parameter. If False, when
|
||||
concurrent access occurs, do not lock the database. Default: False.
|
||||
readahead (bool, optional): Lmdb environment parameter. If False,
|
||||
disable the OS filesystem readahead mechanism, which may improve
|
||||
random read performance when a database is larger than RAM.
|
||||
Default: False.
|
||||
|
||||
Attributes:
|
||||
db_paths (list): Lmdb database path.
|
||||
_client (list): A list of several lmdb envs.
|
||||
"""
|
||||
|
||||
def __init__(self, db_paths, client_keys='default', readonly=True, lock=False, readahead=False, **kwargs):
|
||||
try:
|
||||
import lmdb
|
||||
except ImportError:
|
||||
raise ImportError('Please install lmdb to enable LmdbBackend.')
|
||||
|
||||
if isinstance(client_keys, str):
|
||||
client_keys = [client_keys]
|
||||
|
||||
if isinstance(db_paths, list):
|
||||
self.db_paths = [str(v) for v in db_paths]
|
||||
elif isinstance(db_paths, str):
|
||||
self.db_paths = [str(db_paths)]
|
||||
assert len(client_keys) == len(self.db_paths), ('client_keys and db_paths should have the same length, '
|
||||
f'but received {len(client_keys)} and {len(self.db_paths)}.')
|
||||
|
||||
self._client = {}
|
||||
for client, path in zip(client_keys, self.db_paths):
|
||||
self._client[client] = lmdb.open(path, readonly=readonly, lock=lock, readahead=readahead, **kwargs)
|
||||
|
||||
def get(self, filepath, client_key):
|
||||
"""Get values according to the filepath from one lmdb named client_key.
|
||||
|
||||
Args:
|
||||
filepath (str | obj:`Path`): Here, filepath is the lmdb key.
|
||||
client_key (str): Used for distinguishing different lmdb envs.
|
||||
"""
|
||||
filepath = str(filepath)
|
||||
assert client_key in self._client, (f'client_key {client_key} is not in lmdb clients.')
|
||||
client = self._client[client_key]
|
||||
with client.begin(write=False) as txn:
|
||||
value_buf = txn.get(filepath.encode('ascii'))
|
||||
return value_buf
|
||||
|
||||
def get_text(self, filepath):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class FileClient(object):
|
||||
"""A general file client to access files in different backend.
|
||||
|
||||
The client loads a file or text in a specified backend from its path
|
||||
and return it as a binary file. it can also register other backend
|
||||
accessor with a given name and backend class.
|
||||
|
||||
Attributes:
|
||||
backend (str): The storage backend type. Options are "disk",
|
||||
"memcached" and "lmdb".
|
||||
client (:obj:`BaseStorageBackend`): The backend object.
|
||||
"""
|
||||
|
||||
_backends = {
|
||||
'disk': HardDiskBackend,
|
||||
'memcached': MemcachedBackend,
|
||||
'lmdb': LmdbBackend,
|
||||
}
|
||||
|
||||
def __init__(self, backend='disk', **kwargs):
|
||||
if backend not in self._backends:
|
||||
raise ValueError(f'Backend {backend} is not supported. Currently supported ones'
|
||||
f' are {list(self._backends.keys())}')
|
||||
self.backend = backend
|
||||
self.client = self._backends[backend](**kwargs)
|
||||
|
||||
def get(self, filepath, client_key='default'):
|
||||
# client_key is used only for lmdb, where different fileclients have
|
||||
# different lmdb environments.
|
||||
if self.backend == 'lmdb':
|
||||
return self.client.get(filepath, client_key)
|
||||
else:
|
||||
return self.client.get(filepath)
|
||||
|
||||
def get_text(self, filepath):
|
||||
return self.client.get_text(filepath)
|
||||
@@ -1,170 +0,0 @@
|
||||
# Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/video/optflow.py # noqa: E501
|
||||
import cv2
|
||||
import numpy as np
|
||||
import os
|
||||
|
||||
|
||||
def flowread(flow_path, quantize=False, concat_axis=0, *args, **kwargs):
|
||||
"""Read an optical flow map.
|
||||
|
||||
Args:
|
||||
flow_path (ndarray or str): Flow path.
|
||||
quantize (bool): whether to read quantized pair, if set to True,
|
||||
remaining args will be passed to :func:`dequantize_flow`.
|
||||
concat_axis (int): The axis that dx and dy are concatenated,
|
||||
can be either 0 or 1. Ignored if quantize is False.
|
||||
|
||||
Returns:
|
||||
ndarray: Optical flow represented as a (h, w, 2) numpy array
|
||||
"""
|
||||
if quantize:
|
||||
assert concat_axis in [0, 1]
|
||||
cat_flow = cv2.imread(flow_path, cv2.IMREAD_UNCHANGED)
|
||||
if cat_flow.ndim != 2:
|
||||
raise IOError(f'{flow_path} is not a valid quantized flow file, its dimension is {cat_flow.ndim}.')
|
||||
assert cat_flow.shape[concat_axis] % 2 == 0
|
||||
dx, dy = np.split(cat_flow, 2, axis=concat_axis)
|
||||
flow = dequantize_flow(dx, dy, *args, **kwargs)
|
||||
else:
|
||||
with open(flow_path, 'rb') as f:
|
||||
try:
|
||||
header = f.read(4).decode('utf-8')
|
||||
except Exception:
|
||||
raise IOError(f'Invalid flow file: {flow_path}')
|
||||
else:
|
||||
if header != 'PIEH':
|
||||
raise IOError(f'Invalid flow file: {flow_path}, header does not contain PIEH')
|
||||
|
||||
w = np.fromfile(f, np.int32, 1).squeeze()
|
||||
h = np.fromfile(f, np.int32, 1).squeeze()
|
||||
flow = np.fromfile(f, np.float32, w * h * 2).reshape((h, w, 2))
|
||||
|
||||
return flow.astype(np.float32)
|
||||
|
||||
|
||||
def flowwrite(flow, filename, quantize=False, concat_axis=0, *args, **kwargs):
|
||||
"""Write optical flow to file.
|
||||
|
||||
If the flow is not quantized, it will be saved as a .flo file losslessly,
|
||||
otherwise a jpeg image which is lossy but of much smaller size. (dx and dy
|
||||
will be concatenated horizontally into a single image if quantize is True.)
|
||||
|
||||
Args:
|
||||
flow (ndarray): (h, w, 2) array of optical flow.
|
||||
filename (str): Output filepath.
|
||||
quantize (bool): Whether to quantize the flow and save it to 2 jpeg
|
||||
images. If set to True, remaining args will be passed to
|
||||
:func:`quantize_flow`.
|
||||
concat_axis (int): The axis that dx and dy are concatenated,
|
||||
can be either 0 or 1. Ignored if quantize is False.
|
||||
"""
|
||||
if not quantize:
|
||||
with open(filename, 'wb') as f:
|
||||
f.write('PIEH'.encode('utf-8'))
|
||||
np.array([flow.shape[1], flow.shape[0]], dtype=np.int32).tofile(f)
|
||||
flow = flow.astype(np.float32)
|
||||
flow.tofile(f)
|
||||
f.flush()
|
||||
else:
|
||||
assert concat_axis in [0, 1]
|
||||
dx, dy = quantize_flow(flow, *args, **kwargs)
|
||||
dxdy = np.concatenate((dx, dy), axis=concat_axis)
|
||||
os.makedirs(os.path.dirname(filename), exist_ok=True)
|
||||
cv2.imwrite(filename, dxdy)
|
||||
|
||||
|
||||
def quantize_flow(flow, max_val=0.02, norm=True):
|
||||
"""Quantize flow to [0, 255].
|
||||
|
||||
After this step, the size of flow will be much smaller, and can be
|
||||
dumped as jpeg images.
|
||||
|
||||
Args:
|
||||
flow (ndarray): (h, w, 2) array of optical flow.
|
||||
max_val (float): Maximum value of flow, values beyond
|
||||
[-max_val, max_val] will be truncated.
|
||||
norm (bool): Whether to divide flow values by image width/height.
|
||||
|
||||
Returns:
|
||||
tuple[ndarray]: Quantized dx and dy.
|
||||
"""
|
||||
h, w, _ = flow.shape
|
||||
dx = flow[..., 0]
|
||||
dy = flow[..., 1]
|
||||
if norm:
|
||||
dx = dx / w # avoid inplace operations
|
||||
dy = dy / h
|
||||
# use 255 levels instead of 256 to make sure 0 is 0 after dequantization.
|
||||
flow_comps = [quantize(d, -max_val, max_val, 255, np.uint8) for d in [dx, dy]]
|
||||
return tuple(flow_comps)
|
||||
|
||||
|
||||
def dequantize_flow(dx, dy, max_val=0.02, denorm=True):
|
||||
"""Recover from quantized flow.
|
||||
|
||||
Args:
|
||||
dx (ndarray): Quantized dx.
|
||||
dy (ndarray): Quantized dy.
|
||||
max_val (float): Maximum value used when quantizing.
|
||||
denorm (bool): Whether to multiply flow values with width/height.
|
||||
|
||||
Returns:
|
||||
ndarray: Dequantized flow.
|
||||
"""
|
||||
assert dx.shape == dy.shape
|
||||
assert dx.ndim == 2 or (dx.ndim == 3 and dx.shape[-1] == 1)
|
||||
|
||||
dx, dy = [dequantize(d, -max_val, max_val, 255) for d in [dx, dy]]
|
||||
|
||||
if denorm:
|
||||
dx *= dx.shape[1]
|
||||
dy *= dx.shape[0]
|
||||
flow = np.dstack((dx, dy))
|
||||
return flow
|
||||
|
||||
|
||||
def quantize(arr, min_val, max_val, levels, dtype=np.int64):
|
||||
"""Quantize an array of (-inf, inf) to [0, levels-1].
|
||||
|
||||
Args:
|
||||
arr (ndarray): Input array.
|
||||
min_val (scalar): Minimum value to be clipped.
|
||||
max_val (scalar): Maximum value to be clipped.
|
||||
levels (int): Quantization levels.
|
||||
dtype (np.type): The type of the quantized array.
|
||||
|
||||
Returns:
|
||||
tuple: Quantized array.
|
||||
"""
|
||||
if not (isinstance(levels, int) and levels > 1):
|
||||
raise ValueError(f'levels must be a positive integer, but got {levels}')
|
||||
if min_val >= max_val:
|
||||
raise ValueError(f'min_val ({min_val}) must be smaller than max_val ({max_val})')
|
||||
|
||||
arr = np.clip(arr, min_val, max_val) - min_val
|
||||
quantized_arr = np.minimum(np.floor(levels * arr / (max_val - min_val)).astype(dtype), levels - 1)
|
||||
|
||||
return quantized_arr
|
||||
|
||||
|
||||
def dequantize(arr, min_val, max_val, levels, dtype=np.float64):
|
||||
"""Dequantize an array.
|
||||
|
||||
Args:
|
||||
arr (ndarray): Input array.
|
||||
min_val (scalar): Minimum value to be clipped.
|
||||
max_val (scalar): Maximum value to be clipped.
|
||||
levels (int): Quantization levels.
|
||||
dtype (np.type): The type of the dequantized array.
|
||||
|
||||
Returns:
|
||||
tuple: Dequantized array.
|
||||
"""
|
||||
if not (isinstance(levels, int) and levels > 1):
|
||||
raise ValueError(f'levels must be a positive integer, but got {levels}')
|
||||
if min_val >= max_val:
|
||||
raise ValueError(f'min_val ({min_val}) must be smaller than max_val ({max_val})')
|
||||
|
||||
dequantized_arr = (arr + 0.5).astype(dtype) * (max_val - min_val) / levels + min_val
|
||||
|
||||
return dequantized_arr
|
||||
@@ -1,83 +0,0 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
def filter2D(img, kernel):
|
||||
"""PyTorch version of cv2.filter2D
|
||||
|
||||
Args:
|
||||
img (Tensor): (b, c, h, w)
|
||||
kernel (Tensor): (b, k, k)
|
||||
"""
|
||||
k = kernel.size(-1)
|
||||
b, c, h, w = img.size()
|
||||
if k % 2 == 1:
|
||||
img = F.pad(img, (k // 2, k // 2, k // 2, k // 2), mode='reflect')
|
||||
else:
|
||||
raise ValueError('Wrong kernel size')
|
||||
|
||||
ph, pw = img.size()[-2:]
|
||||
|
||||
if kernel.size(0) == 1:
|
||||
# apply the same kernel to all batch images
|
||||
img = img.view(b * c, 1, ph, pw)
|
||||
kernel = kernel.view(1, 1, k, k)
|
||||
return F.conv2d(img, kernel, padding=0).view(b, c, h, w)
|
||||
else:
|
||||
img = img.view(1, b * c, ph, pw)
|
||||
kernel = kernel.view(b, 1, k, k).repeat(1, c, 1, 1).view(b * c, 1, k, k)
|
||||
return F.conv2d(img, kernel, groups=b * c).view(b, c, h, w)
|
||||
|
||||
|
||||
def usm_sharp(img, weight=0.5, radius=50, threshold=10):
|
||||
"""USM sharpening.
|
||||
|
||||
Input image: I; Blurry image: B.
|
||||
1. sharp = I + weight * (I - B)
|
||||
2. Mask = 1 if abs(I - B) > threshold, else: 0
|
||||
3. Blur mask:
|
||||
4. Out = Mask * sharp + (1 - Mask) * I
|
||||
|
||||
|
||||
Args:
|
||||
img (Numpy array): Input image, HWC, BGR; float32, [0, 1].
|
||||
weight (float): Sharp weight. Default: 1.
|
||||
radius (float): Kernel size of Gaussian blur. Default: 50.
|
||||
threshold (int):
|
||||
"""
|
||||
if radius % 2 == 0:
|
||||
radius += 1
|
||||
blur = cv2.GaussianBlur(img, (radius, radius), 0)
|
||||
residual = img - blur
|
||||
mask = np.abs(residual) * 255 > threshold
|
||||
mask = mask.astype('float32')
|
||||
soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0)
|
||||
|
||||
sharp = img + weight * residual
|
||||
sharp = np.clip(sharp, 0, 1)
|
||||
return soft_mask * sharp + (1 - soft_mask) * img
|
||||
|
||||
|
||||
class USMSharp(torch.nn.Module):
|
||||
|
||||
def __init__(self, radius=50, sigma=0):
|
||||
super(USMSharp, self).__init__()
|
||||
if radius % 2 == 0:
|
||||
radius += 1
|
||||
self.radius = radius
|
||||
kernel = cv2.getGaussianKernel(radius, sigma)
|
||||
kernel = torch.FloatTensor(np.dot(kernel, kernel.transpose())).unsqueeze_(0)
|
||||
self.register_buffer('kernel', kernel)
|
||||
|
||||
def forward(self, img, weight=0.5, threshold=10):
|
||||
blur = filter2D(img, self.kernel)
|
||||
residual = img - blur
|
||||
|
||||
mask = torch.abs(residual) * 255 > threshold
|
||||
mask = mask.float()
|
||||
soft_mask = filter2D(mask, self.kernel)
|
||||
sharp = img + weight * residual
|
||||
sharp = torch.clip(sharp, 0, 1)
|
||||
return soft_mask * sharp + (1 - soft_mask) * img
|
||||
@@ -1,172 +0,0 @@
|
||||
import cv2
|
||||
import math
|
||||
import numpy as np
|
||||
import os
|
||||
import torch
|
||||
from torchvision.utils import make_grid
|
||||
|
||||
|
||||
def img2tensor(imgs, bgr2rgb=True, float32=True):
|
||||
"""Numpy array to tensor.
|
||||
|
||||
Args:
|
||||
imgs (list[ndarray] | ndarray): Input images.
|
||||
bgr2rgb (bool): Whether to change bgr to rgb.
|
||||
float32 (bool): Whether to change to float32.
|
||||
|
||||
Returns:
|
||||
list[tensor] | tensor: Tensor images. If returned results only have
|
||||
one element, just return tensor.
|
||||
"""
|
||||
|
||||
def _totensor(img, bgr2rgb, float32):
|
||||
if img.shape[2] == 3 and bgr2rgb:
|
||||
if img.dtype == 'float64':
|
||||
img = img.astype('float32')
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
img = torch.from_numpy(img.transpose(2, 0, 1))
|
||||
if float32:
|
||||
img = img.float()
|
||||
return img
|
||||
|
||||
if isinstance(imgs, list):
|
||||
return [_totensor(img, bgr2rgb, float32) for img in imgs]
|
||||
else:
|
||||
return _totensor(imgs, bgr2rgb, float32)
|
||||
|
||||
|
||||
def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)):
|
||||
"""Convert torch Tensors into image numpy arrays.
|
||||
|
||||
After clamping to [min, max], values will be normalized to [0, 1].
|
||||
|
||||
Args:
|
||||
tensor (Tensor or list[Tensor]): Accept shapes:
|
||||
1) 4D mini-batch Tensor of shape (B x 3/1 x H x W);
|
||||
2) 3D Tensor of shape (3/1 x H x W);
|
||||
3) 2D Tensor of shape (H x W).
|
||||
Tensor channel should be in RGB order.
|
||||
rgb2bgr (bool): Whether to change rgb to bgr.
|
||||
out_type (numpy type): output types. If ``np.uint8``, transform outputs
|
||||
to uint8 type with range [0, 255]; otherwise, float type with
|
||||
range [0, 1]. Default: ``np.uint8``.
|
||||
min_max (tuple[int]): min and max values for clamp.
|
||||
|
||||
Returns:
|
||||
(Tensor or list): 3D ndarray of shape (H x W x C) OR 2D ndarray of
|
||||
shape (H x W). The channel order is BGR.
|
||||
"""
|
||||
if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):
|
||||
raise TypeError(f'tensor or list of tensors expected, got {type(tensor)}')
|
||||
|
||||
if torch.is_tensor(tensor):
|
||||
tensor = [tensor]
|
||||
result = []
|
||||
for _tensor in tensor:
|
||||
_tensor = _tensor.squeeze(0).float().detach().cpu().clamp_(*min_max)
|
||||
_tensor = (_tensor - min_max[0]) / (min_max[1] - min_max[0])
|
||||
|
||||
n_dim = _tensor.dim()
|
||||
if n_dim == 4:
|
||||
img_np = make_grid(_tensor, nrow=int(math.sqrt(_tensor.size(0))), normalize=False).numpy()
|
||||
img_np = img_np.transpose(1, 2, 0)
|
||||
if rgb2bgr:
|
||||
img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
|
||||
elif n_dim == 3:
|
||||
img_np = _tensor.numpy()
|
||||
img_np = img_np.transpose(1, 2, 0)
|
||||
if img_np.shape[2] == 1: # gray image
|
||||
img_np = np.squeeze(img_np, axis=2)
|
||||
else:
|
||||
if rgb2bgr:
|
||||
img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
|
||||
elif n_dim == 2:
|
||||
img_np = _tensor.numpy()
|
||||
else:
|
||||
raise TypeError(f'Only support 4D, 3D or 2D tensor. But received with dimension: {n_dim}')
|
||||
if out_type == np.uint8:
|
||||
# Unlike MATLAB, numpy.unit8() WILL NOT round by default.
|
||||
img_np = (img_np * 255.0).round()
|
||||
img_np = img_np.astype(out_type)
|
||||
result.append(img_np)
|
||||
if len(result) == 1 and torch.is_tensor(tensor):
|
||||
result = result[0]
|
||||
return result
|
||||
|
||||
|
||||
def tensor2img_fast(tensor, rgb2bgr=True, min_max=(0, 1)):
|
||||
"""This implementation is slightly faster than tensor2img.
|
||||
It now only supports torch tensor with shape (1, c, h, w).
|
||||
|
||||
Args:
|
||||
tensor (Tensor): Now only support torch tensor with (1, c, h, w).
|
||||
rgb2bgr (bool): Whether to change rgb to bgr. Default: True.
|
||||
min_max (tuple[int]): min and max values for clamp.
|
||||
"""
|
||||
output = tensor.squeeze(0).detach().clamp_(*min_max).permute(1, 2, 0)
|
||||
output = (output - min_max[0]) / (min_max[1] - min_max[0]) * 255
|
||||
output = output.type(torch.uint8).cpu().numpy()
|
||||
if rgb2bgr:
|
||||
output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
|
||||
return output
|
||||
|
||||
|
||||
def imfrombytes(content, flag='color', float32=False):
|
||||
"""Read an image from bytes.
|
||||
|
||||
Args:
|
||||
content (bytes): Image bytes got from files or other streams.
|
||||
flag (str): Flags specifying the color type of a loaded image,
|
||||
candidates are `color`, `grayscale` and `unchanged`.
|
||||
float32 (bool): Whether to change to float32., If True, will also norm
|
||||
to [0, 1]. Default: False.
|
||||
|
||||
Returns:
|
||||
ndarray: Loaded image array.
|
||||
"""
|
||||
img_np = np.frombuffer(content, np.uint8)
|
||||
imread_flags = {'color': cv2.IMREAD_COLOR, 'grayscale': cv2.IMREAD_GRAYSCALE, 'unchanged': cv2.IMREAD_UNCHANGED}
|
||||
img = cv2.imdecode(img_np, imread_flags[flag])
|
||||
if float32:
|
||||
img = img.astype(np.float32) / 255.
|
||||
return img
|
||||
|
||||
|
||||
def imwrite(img, file_path, params=None, auto_mkdir=True):
|
||||
"""Write image to file.
|
||||
|
||||
Args:
|
||||
img (ndarray): Image array to be written.
|
||||
file_path (str): Image file path.
|
||||
params (None or list): Same as opencv's :func:`imwrite` interface.
|
||||
auto_mkdir (bool): If the parent folder of `file_path` does not exist,
|
||||
whether to create it automatically.
|
||||
|
||||
Returns:
|
||||
bool: Successful or not.
|
||||
"""
|
||||
if auto_mkdir:
|
||||
dir_name = os.path.abspath(os.path.dirname(file_path))
|
||||
os.makedirs(dir_name, exist_ok=True)
|
||||
ok = cv2.imwrite(file_path, img, params)
|
||||
if not ok:
|
||||
raise IOError('Failed in writing images.')
|
||||
|
||||
|
||||
def crop_border(imgs, crop_border):
|
||||
"""Crop borders of images.
|
||||
|
||||
Args:
|
||||
imgs (list[ndarray] | ndarray): Images with shape (h, w, c).
|
||||
crop_border (int): Crop border for each end of height and weight.
|
||||
|
||||
Returns:
|
||||
list[ndarray]: Cropped images.
|
||||
"""
|
||||
if crop_border == 0:
|
||||
return imgs
|
||||
else:
|
||||
if isinstance(imgs, list):
|
||||
return [v[crop_border:-crop_border, crop_border:-crop_border, ...] for v in imgs]
|
||||
else:
|
||||
return imgs[crop_border:-crop_border, crop_border:-crop_border, ...]
|
||||
@@ -1,199 +0,0 @@
|
||||
import cv2
|
||||
import lmdb
|
||||
import sys
|
||||
from multiprocessing import Pool
|
||||
from os import path as osp
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
def make_lmdb_from_imgs(data_path,
|
||||
lmdb_path,
|
||||
img_path_list,
|
||||
keys,
|
||||
batch=5000,
|
||||
compress_level=1,
|
||||
multiprocessing_read=False,
|
||||
n_thread=40,
|
||||
map_size=None):
|
||||
"""Make lmdb from images.
|
||||
|
||||
Contents of lmdb. The file structure is:
|
||||
|
||||
::
|
||||
|
||||
example.lmdb
|
||||
├── data.mdb
|
||||
├── lock.mdb
|
||||
├── meta_info.txt
|
||||
|
||||
The data.mdb and lock.mdb are standard lmdb files and you can refer to
|
||||
https://lmdb.readthedocs.io/en/release/ for more details.
|
||||
|
||||
The meta_info.txt is a specified txt file to record the meta information
|
||||
of our datasets. It will be automatically created when preparing
|
||||
datasets by our provided dataset tools.
|
||||
Each line in the txt file records 1)image name (with extension),
|
||||
2)image shape, and 3)compression level, separated by a white space.
|
||||
|
||||
For example, the meta information could be:
|
||||
`000_00000000.png (720,1280,3) 1`, which means:
|
||||
1) image name (with extension): 000_00000000.png;
|
||||
2) image shape: (720,1280,3);
|
||||
3) compression level: 1
|
||||
|
||||
We use the image name without extension as the lmdb key.
|
||||
|
||||
If `multiprocessing_read` is True, it will read all the images to memory
|
||||
using multiprocessing. Thus, your server needs to have enough memory.
|
||||
|
||||
Args:
|
||||
data_path (str): Data path for reading images.
|
||||
lmdb_path (str): Lmdb save path.
|
||||
img_path_list (str): Image path list.
|
||||
keys (str): Used for lmdb keys.
|
||||
batch (int): After processing batch images, lmdb commits.
|
||||
Default: 5000.
|
||||
compress_level (int): Compress level when encoding images. Default: 1.
|
||||
multiprocessing_read (bool): Whether use multiprocessing to read all
|
||||
the images to memory. Default: False.
|
||||
n_thread (int): For multiprocessing.
|
||||
map_size (int | None): Map size for lmdb env. If None, use the
|
||||
estimated size from images. Default: None
|
||||
"""
|
||||
|
||||
assert len(img_path_list) == len(keys), ('img_path_list and keys should have the same length, '
|
||||
f'but got {len(img_path_list)} and {len(keys)}')
|
||||
print(f'Create lmdb for {data_path}, save to {lmdb_path}...')
|
||||
print(f'Totoal images: {len(img_path_list)}')
|
||||
if not lmdb_path.endswith('.lmdb'):
|
||||
raise ValueError("lmdb_path must end with '.lmdb'.")
|
||||
if osp.exists(lmdb_path):
|
||||
print(f'Folder {lmdb_path} already exists. Exit.')
|
||||
sys.exit(1)
|
||||
|
||||
if multiprocessing_read:
|
||||
# read all the images to memory (multiprocessing)
|
||||
dataset = {} # use dict to keep the order for multiprocessing
|
||||
shapes = {}
|
||||
print(f'Read images with multiprocessing, #thread: {n_thread} ...')
|
||||
pbar = tqdm(total=len(img_path_list), unit='image')
|
||||
|
||||
def callback(arg):
|
||||
"""get the image data and update pbar."""
|
||||
key, dataset[key], shapes[key] = arg
|
||||
pbar.update(1)
|
||||
pbar.set_description(f'Read {key}')
|
||||
|
||||
pool = Pool(n_thread)
|
||||
for path, key in zip(img_path_list, keys):
|
||||
pool.apply_async(read_img_worker, args=(osp.join(data_path, path), key, compress_level), callback=callback)
|
||||
pool.close()
|
||||
pool.join()
|
||||
pbar.close()
|
||||
print(f'Finish reading {len(img_path_list)} images.')
|
||||
|
||||
# create lmdb environment
|
||||
if map_size is None:
|
||||
# obtain data size for one image
|
||||
img = cv2.imread(osp.join(data_path, img_path_list[0]), cv2.IMREAD_UNCHANGED)
|
||||
_, img_byte = cv2.imencode('.png', img, [cv2.IMWRITE_PNG_COMPRESSION, compress_level])
|
||||
data_size_per_img = img_byte.nbytes
|
||||
print('Data size per image is: ', data_size_per_img)
|
||||
data_size = data_size_per_img * len(img_path_list)
|
||||
map_size = data_size * 10
|
||||
|
||||
env = lmdb.open(lmdb_path, map_size=map_size)
|
||||
|
||||
# write data to lmdb
|
||||
pbar = tqdm(total=len(img_path_list), unit='chunk')
|
||||
txn = env.begin(write=True)
|
||||
txt_file = open(osp.join(lmdb_path, 'meta_info.txt'), 'w')
|
||||
for idx, (path, key) in enumerate(zip(img_path_list, keys)):
|
||||
pbar.update(1)
|
||||
pbar.set_description(f'Write {key}')
|
||||
key_byte = key.encode('ascii')
|
||||
if multiprocessing_read:
|
||||
img_byte = dataset[key]
|
||||
h, w, c = shapes[key]
|
||||
else:
|
||||
_, img_byte, img_shape = read_img_worker(osp.join(data_path, path), key, compress_level)
|
||||
h, w, c = img_shape
|
||||
|
||||
txn.put(key_byte, img_byte)
|
||||
# write meta information
|
||||
txt_file.write(f'{key}.png ({h},{w},{c}) {compress_level}\n')
|
||||
if idx % batch == 0:
|
||||
txn.commit()
|
||||
txn = env.begin(write=True)
|
||||
pbar.close()
|
||||
txn.commit()
|
||||
env.close()
|
||||
txt_file.close()
|
||||
print('\nFinish writing lmdb.')
|
||||
|
||||
|
||||
def read_img_worker(path, key, compress_level):
|
||||
"""Read image worker.
|
||||
|
||||
Args:
|
||||
path (str): Image path.
|
||||
key (str): Image key.
|
||||
compress_level (int): Compress level when encoding images.
|
||||
|
||||
Returns:
|
||||
str: Image key.
|
||||
byte: Image byte.
|
||||
tuple[int]: Image shape.
|
||||
"""
|
||||
|
||||
img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
|
||||
if img.ndim == 2:
|
||||
h, w = img.shape
|
||||
c = 1
|
||||
else:
|
||||
h, w, c = img.shape
|
||||
_, img_byte = cv2.imencode('.png', img, [cv2.IMWRITE_PNG_COMPRESSION, compress_level])
|
||||
return (key, img_byte, (h, w, c))
|
||||
|
||||
|
||||
class LmdbMaker():
|
||||
"""LMDB Maker.
|
||||
|
||||
Args:
|
||||
lmdb_path (str): Lmdb save path.
|
||||
map_size (int): Map size for lmdb env. Default: 1024 ** 4, 1TB.
|
||||
batch (int): After processing batch images, lmdb commits.
|
||||
Default: 5000.
|
||||
compress_level (int): Compress level when encoding images. Default: 1.
|
||||
"""
|
||||
|
||||
def __init__(self, lmdb_path, map_size=1024**4, batch=5000, compress_level=1):
|
||||
if not lmdb_path.endswith('.lmdb'):
|
||||
raise ValueError("lmdb_path must end with '.lmdb'.")
|
||||
if osp.exists(lmdb_path):
|
||||
print(f'Folder {lmdb_path} already exists. Exit.')
|
||||
sys.exit(1)
|
||||
|
||||
self.lmdb_path = lmdb_path
|
||||
self.batch = batch
|
||||
self.compress_level = compress_level
|
||||
self.env = lmdb.open(lmdb_path, map_size=map_size)
|
||||
self.txn = self.env.begin(write=True)
|
||||
self.txt_file = open(osp.join(lmdb_path, 'meta_info.txt'), 'w')
|
||||
self.counter = 0
|
||||
|
||||
def put(self, img_byte, key, img_shape):
|
||||
self.counter += 1
|
||||
key_byte = key.encode('ascii')
|
||||
self.txn.put(key_byte, img_byte)
|
||||
# write meta information
|
||||
h, w, c = img_shape
|
||||
self.txt_file.write(f'{key}.png ({h},{w},{c}) {self.compress_level}\n')
|
||||
if self.counter % self.batch == 0:
|
||||
self.txn.commit()
|
||||
self.txn = self.env.begin(write=True)
|
||||
|
||||
def close(self):
|
||||
self.txn.commit()
|
||||
self.env.close()
|
||||
self.txt_file.close()
|
||||
@@ -1,213 +0,0 @@
|
||||
import datetime
|
||||
import logging
|
||||
import time
|
||||
|
||||
from .dist_util import get_dist_info, master_only
|
||||
|
||||
initialized_logger = {}
|
||||
|
||||
|
||||
class AvgTimer():
|
||||
|
||||
def __init__(self, window=200):
|
||||
self.window = window # average window
|
||||
self.current_time = 0
|
||||
self.total_time = 0
|
||||
self.count = 0
|
||||
self.avg_time = 0
|
||||
self.start()
|
||||
|
||||
def start(self):
|
||||
self.start_time = self.tic = time.time()
|
||||
|
||||
def record(self):
|
||||
self.count += 1
|
||||
self.toc = time.time()
|
||||
self.current_time = self.toc - self.tic
|
||||
self.total_time += self.current_time
|
||||
# calculate average time
|
||||
self.avg_time = self.total_time / self.count
|
||||
|
||||
# reset
|
||||
if self.count > self.window:
|
||||
self.count = 0
|
||||
self.total_time = 0
|
||||
|
||||
self.tic = time.time()
|
||||
|
||||
def get_current_time(self):
|
||||
return self.current_time
|
||||
|
||||
def get_avg_time(self):
|
||||
return self.avg_time
|
||||
|
||||
|
||||
class MessageLogger():
|
||||
"""Message logger for printing.
|
||||
|
||||
Args:
|
||||
opt (dict): Config. It contains the following keys:
|
||||
name (str): Exp name.
|
||||
logger (dict): Contains 'print_freq' (str) for logger interval.
|
||||
train (dict): Contains 'total_iter' (int) for total iters.
|
||||
use_tb_logger (bool): Use tensorboard logger.
|
||||
start_iter (int): Start iter. Default: 1.
|
||||
tb_logger (obj:`tb_logger`): Tensorboard logger. Default： None.
|
||||
"""
|
||||
|
||||
def __init__(self, opt, start_iter=1, tb_logger=None):
|
||||
self.exp_name = opt['name']
|
||||
self.interval = opt['logger']['print_freq']
|
||||
self.start_iter = start_iter
|
||||
self.max_iters = opt['train']['total_iter']
|
||||
self.use_tb_logger = opt['logger']['use_tb_logger']
|
||||
self.tb_logger = tb_logger
|
||||
self.start_time = time.time()
|
||||
self.logger = get_root_logger()
|
||||
|
||||
def reset_start_time(self):
|
||||
self.start_time = time.time()
|
||||
|
||||
@master_only
|
||||
def __call__(self, log_vars):
|
||||
"""Format logging message.
|
||||
|
||||
Args:
|
||||
log_vars (dict): It contains the following keys:
|
||||
epoch (int): Epoch number.
|
||||
iter (int): Current iter.
|
||||
lrs (list): List for learning rates.
|
||||
|
||||
time (float): Iter time.
|
||||
data_time (float): Data time for each iter.
|
||||
"""
|
||||
# epoch, iter, learning rates
|
||||
epoch = log_vars.pop('epoch')
|
||||
current_iter = log_vars.pop('iter')
|
||||
lrs = log_vars.pop('lrs')
|
||||
|
||||
message = (f'[{self.exp_name[:5]}..][epoch:{epoch:3d}, iter:{current_iter:8,d}, lr:(')
|
||||
for v in lrs:
|
||||
message += f'{v:.3e},'
|
||||
message += ')] '
|
||||
|
||||
# time and estimated time
|
||||
if 'time' in log_vars.keys():
|
||||
iter_time = log_vars.pop('time')
|
||||
data_time = log_vars.pop('data_time')
|
||||
|
||||
total_time = time.time() - self.start_time
|
||||
time_sec_avg = total_time / (current_iter - self.start_iter + 1)
|
||||
eta_sec = time_sec_avg * (self.max_iters - current_iter - 1)
|
||||
eta_str = str(datetime.timedelta(seconds=int(eta_sec)))
|
||||
message += f'[eta: {eta_str}, '
|
||||
message += f'time (data): {iter_time:.3f} ({data_time:.3f})] '
|
||||
|
||||
# other items, especially losses
|
||||
for k, v in log_vars.items():
|
||||
message += f'{k}: {v:.4e} '
|
||||
# tensorboard logger
|
||||
if self.use_tb_logger and 'debug' not in self.exp_name:
|
||||
if k.startswith('l_'):
|
||||
self.tb_logger.add_scalar(f'losses/{k}', v, current_iter)
|
||||
else:
|
||||
self.tb_logger.add_scalar(k, v, current_iter)
|
||||
self.logger.info(message)
|
||||
|
||||
|
||||
@master_only
|
||||
def init_tb_logger(log_dir):
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
tb_logger = SummaryWriter(log_dir=log_dir)
|
||||
return tb_logger
|
||||
|
||||
|
||||
@master_only
|
||||
def init_wandb_logger(opt):
|
||||
"""We now only use wandb to sync tensorboard log."""
|
||||
import wandb
|
||||
logger = get_root_logger()
|
||||
|
||||
project = opt['logger']['wandb']['project']
|
||||
resume_id = opt['logger']['wandb'].get('resume_id')
|
||||
if resume_id:
|
||||
wandb_id = resume_id
|
||||
resume = 'allow'
|
||||
logger.warning(f'Resume wandb logger with id={wandb_id}.')
|
||||
else:
|
||||
wandb_id = wandb.util.generate_id()
|
||||
resume = 'never'
|
||||
|
||||
wandb.init(id=wandb_id, resume=resume, name=opt['name'], config=opt, project=project, sync_tensorboard=True)
|
||||
|
||||
logger.info(f'Use wandb logger with id={wandb_id}; project={project}.')
|
||||
|
||||
|
||||
def get_root_logger(logger_name='basicsr', log_level=logging.INFO, log_file=None):
|
||||
"""Get the root logger.
|
||||
|
||||
The logger will be initialized if it has not been initialized. By default a
|
||||
StreamHandler will be added. If `log_file` is specified, a FileHandler will
|
||||
also be added.
|
||||
|
||||
Args:
|
||||
logger_name (str): root logger name. Default: 'basicsr'.
|
||||
log_file (str | None): The log filename. If specified, a FileHandler
|
||||
will be added to the root logger.
|
||||
log_level (int): The root logger level. Note that only the process of
|
||||
rank 0 is affected, while other processes will set the level to
|
||||
"Error" and be silent most of the time.
|
||||
|
||||
Returns:
|
||||
logging.Logger: The root logger.
|
||||
"""
|
||||
logger = logging.getLogger(logger_name)
|
||||
# if the logger has been initialized, just return it
|
||||
if logger_name in initialized_logger:
|
||||
return logger
|
||||
|
||||
format_str = '%(asctime)s %(levelname)s: %(message)s'
|
||||
stream_handler = logging.StreamHandler()
|
||||
stream_handler.setFormatter(logging.Formatter(format_str))
|
||||
logger.addHandler(stream_handler)
|
||||
logger.propagate = False
|
||||
rank, _ = get_dist_info()
|
||||
if rank != 0:
|
||||
logger.setLevel('ERROR')
|
||||
elif log_file is not None:
|
||||
logger.setLevel(log_level)
|
||||
# add file handler
|
||||
file_handler = logging.FileHandler(log_file, 'w')
|
||||
file_handler.setFormatter(logging.Formatter(format_str))
|
||||
file_handler.setLevel(log_level)
|
||||
logger.addHandler(file_handler)
|
||||
initialized_logger[logger_name] = True
|
||||
return logger
|
||||
|
||||
|
||||
def get_env_info():
|
||||
"""Get environment information.
|
||||
|
||||
Currently, only log the software version.
|
||||
"""
|
||||
import torch
|
||||
import torchvision
|
||||
|
||||
from basicsr.version import __version__
|
||||
msg = r"""
|
||||
____ _ _____ ____
|
||||
/ __ ) ____ _ _____ (_)_____/ ___/ / __ \
|
||||
/ __ |/ __ `// ___// // ___/\__ \ / /_/ /
|
||||
/ /_/ // /_/ /(__ )/ // /__ ___/ // _, _/
|
||||
/_____/ \__,_//____//_/ \___//____//_/ |_|
|
||||
______ __ __ __ __
|
||||
/ ____/____ ____ ____/ / / / __ __ _____ / /__ / /
|
||||
/ / __ / __ \ / __ \ / __ / / / / / / // ___// //_/ / /
|
||||
/ /_/ // /_/ // /_/ // /_/ / / /___/ /_/ // /__ / /< /_/
|
||||
\____/ \____/ \____/ \____/ /_____/\____/ \___//_/|_| (_)
|
||||
"""
|
||||
msg += ('\nVersion Information: '
|
||||
f'\n\tBasicSR: {__version__}'
|
||||
f'\n\tPyTorch: {torch.__version__}'
|
||||
f'\n\tTorchVision: {torchvision.__version__}')
|
||||
return msg
|
||||
@@ -1,178 +0,0 @@
|
||||
import math
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def cubic(x):
|
||||
"""cubic function used for calculate_weights_indices."""
|
||||
absx = torch.abs(x)
|
||||
absx2 = absx**2
|
||||
absx3 = absx**3
|
||||
return (1.5 * absx3 - 2.5 * absx2 + 1) * (
|
||||
(absx <= 1).type_as(absx)) + (-0.5 * absx3 + 2.5 * absx2 - 4 * absx + 2) * (((absx > 1) *
|
||||
(absx <= 2)).type_as(absx))
|
||||
|
||||
|
||||
def calculate_weights_indices(in_length, out_length, scale, kernel, kernel_width, antialiasing):
|
||||
"""Calculate weights and indices, used for imresize function.
|
||||
|
||||
Args:
|
||||
in_length (int): Input length.
|
||||
out_length (int): Output length.
|
||||
scale (float): Scale factor.
|
||||
kernel_width (int): Kernel width.
|
||||
antialisaing (bool): Whether to apply anti-aliasing when downsampling.
|
||||
"""
|
||||
|
||||
if (scale < 1) and antialiasing:
|
||||
# Use a modified kernel (larger kernel width) to simultaneously
|
||||
# interpolate and antialias
|
||||
kernel_width = kernel_width / scale
|
||||
|
||||
# Output-space coordinates
|
||||
x = torch.linspace(1, out_length, out_length)
|
||||
|
||||
# Input-space coordinates. Calculate the inverse mapping such that 0.5
|
||||
# in output space maps to 0.5 in input space, and 0.5 + scale in output
|
||||
# space maps to 1.5 in input space.
|
||||
u = x / scale + 0.5 * (1 - 1 / scale)
|
||||
|
||||
# What is the left-most pixel that can be involved in the computation?
|
||||
left = torch.floor(u - kernel_width / 2)
|
||||
|
||||
# What is the maximum number of pixels that can be involved in the
|
||||
# computation? Note: it's OK to use an extra pixel here; if the
|
||||
# corresponding weights are all zero, it will be eliminated at the end
|
||||
# of this function.
|
||||
p = math.ceil(kernel_width) + 2
|
||||
|
||||
# The indices of the input pixels involved in computing the k-th output
|
||||
# pixel are in row k of the indices matrix.
|
||||
indices = left.view(out_length, 1).expand(out_length, p) + torch.linspace(0, p - 1, p).view(1, p).expand(
|
||||
out_length, p)
|
||||
|
||||
# The weights used to compute the k-th output pixel are in row k of the
|
||||
# weights matrix.
|
||||
distance_to_center = u.view(out_length, 1).expand(out_length, p) - indices
|
||||
|
||||
# apply cubic kernel
|
||||
if (scale < 1) and antialiasing:
|
||||
weights = scale * cubic(distance_to_center * scale)
|
||||
else:
|
||||
weights = cubic(distance_to_center)
|
||||
|
||||
# Normalize the weights matrix so that each row sums to 1.
|
||||
weights_sum = torch.sum(weights, 1).view(out_length, 1)
|
||||
weights = weights / weights_sum.expand(out_length, p)
|
||||
|
||||
# If a column in weights is all zero, get rid of it. only consider the
|
||||
# first and last column.
|
||||
weights_zero_tmp = torch.sum((weights == 0), 0)
|
||||
if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):
|
||||
indices = indices.narrow(1, 1, p - 2)
|
||||
weights = weights.narrow(1, 1, p - 2)
|
||||
if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):
|
||||
indices = indices.narrow(1, 0, p - 2)
|
||||
weights = weights.narrow(1, 0, p - 2)
|
||||
weights = weights.contiguous()
|
||||
indices = indices.contiguous()
|
||||
sym_len_s = -indices.min() + 1
|
||||
sym_len_e = indices.max() - in_length
|
||||
indices = indices + sym_len_s - 1
|
||||
return weights, indices, int(sym_len_s), int(sym_len_e)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def imresize(img, scale, antialiasing=True):
|
||||
"""imresize function same as MATLAB.
|
||||
|
||||
It now only supports bicubic.
|
||||
The same scale applies for both height and width.
|
||||
|
||||
Args:
|
||||
img (Tensor | Numpy array):
|
||||
Tensor: Input image with shape (c, h, w), [0, 1] range.
|
||||
Numpy: Input image with shape (h, w, c), [0, 1] range.
|
||||
scale (float): Scale factor. The same scale applies for both height
|
||||
and width.
|
||||
antialisaing (bool): Whether to apply anti-aliasing when downsampling.
|
||||
Default: True.
|
||||
|
||||
Returns:
|
||||
Tensor: Output image with shape (c, h, w), [0, 1] range, w/o round.
|
||||
"""
|
||||
squeeze_flag = False
|
||||
if type(img).__module__ == np.__name__: # numpy type
|
||||
numpy_type = True
|
||||
if img.ndim == 2:
|
||||
img = img[:, :, None]
|
||||
squeeze_flag = True
|
||||
img = torch.from_numpy(img.transpose(2, 0, 1)).float()
|
||||
else:
|
||||
numpy_type = False
|
||||
if img.ndim == 2:
|
||||
img = img.unsqueeze(0)
|
||||
squeeze_flag = True
|
||||
|
||||
in_c, in_h, in_w = img.size()
|
||||
out_h, out_w = math.ceil(in_h * scale), math.ceil(in_w * scale)
|
||||
kernel_width = 4
|
||||
kernel = 'cubic'
|
||||
|
||||
# get weights and indices
|
||||
weights_h, indices_h, sym_len_hs, sym_len_he = calculate_weights_indices(in_h, out_h, scale, kernel, kernel_width,
|
||||
antialiasing)
|
||||
weights_w, indices_w, sym_len_ws, sym_len_we = calculate_weights_indices(in_w, out_w, scale, kernel, kernel_width,
|
||||
antialiasing)
|
||||
# process H dimension
|
||||
# symmetric copying
|
||||
img_aug = torch.FloatTensor(in_c, in_h + sym_len_hs + sym_len_he, in_w)
|
||||
img_aug.narrow(1, sym_len_hs, in_h).copy_(img)
|
||||
|
||||
sym_patch = img[:, :sym_len_hs, :]
|
||||
inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
|
||||
sym_patch_inv = sym_patch.index_select(1, inv_idx)
|
||||
img_aug.narrow(1, 0, sym_len_hs).copy_(sym_patch_inv)
|
||||
|
||||
sym_patch = img[:, -sym_len_he:, :]
|
||||
inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
|
||||
sym_patch_inv = sym_patch.index_select(1, inv_idx)
|
||||
img_aug.narrow(1, sym_len_hs + in_h, sym_len_he).copy_(sym_patch_inv)
|
||||
|
||||
out_1 = torch.FloatTensor(in_c, out_h, in_w)
|
||||
kernel_width = weights_h.size(1)
|
||||
for i in range(out_h):
|
||||
idx = int(indices_h[i][0])
|
||||
for j in range(in_c):
|
||||
out_1[j, i, :] = img_aug[j, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_h[i])
|
||||
|
||||
# process W dimension
|
||||
# symmetric copying
|
||||
out_1_aug = torch.FloatTensor(in_c, out_h, in_w + sym_len_ws + sym_len_we)
|
||||
out_1_aug.narrow(2, sym_len_ws, in_w).copy_(out_1)
|
||||
|
||||
sym_patch = out_1[:, :, :sym_len_ws]
|
||||
inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
|
||||
sym_patch_inv = sym_patch.index_select(2, inv_idx)
|
||||
out_1_aug.narrow(2, 0, sym_len_ws).copy_(sym_patch_inv)
|
||||
|
||||
sym_patch = out_1[:, :, -sym_len_we:]
|
||||
inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
|
||||
sym_patch_inv = sym_patch.index_select(2, inv_idx)
|
||||
out_1_aug.narrow(2, sym_len_ws + in_w, sym_len_we).copy_(sym_patch_inv)
|
||||
|
||||
out_2 = torch.FloatTensor(in_c, out_h, out_w)
|
||||
kernel_width = weights_w.size(1)
|
||||
for i in range(out_w):
|
||||
idx = int(indices_w[i][0])
|
||||
for j in range(in_c):
|
||||
out_2[j, :, i] = out_1_aug[j, :, idx:idx + kernel_width].mv(weights_w[i])
|
||||
|
||||
if squeeze_flag:
|
||||
out_2 = out_2.squeeze(0)
|
||||
if numpy_type:
|
||||
out_2 = out_2.numpy()
|
||||
if not squeeze_flag:
|
||||
out_2 = out_2.transpose(1, 2, 0)
|
||||
|
||||
return out_2
|
||||
@@ -1,141 +0,0 @@
|
||||
import numpy as np
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
import torch
|
||||
from os import path as osp
|
||||
|
||||
from .dist_util import master_only
|
||||
|
||||
|
||||
def set_random_seed(seed):
|
||||
"""Set random seeds."""
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
def get_time_str():
|
||||
return time.strftime('%Y%m%d_%H%M%S', time.localtime())
|
||||
|
||||
|
||||
def mkdir_and_rename(path):
|
||||
"""mkdirs. If path exists, rename it with timestamp and create a new one.
|
||||
|
||||
Args:
|
||||
path (str): Folder path.
|
||||
"""
|
||||
if osp.exists(path):
|
||||
new_name = path + '_archived_' + get_time_str()
|
||||
print(f'Path already exists. Rename it to {new_name}', flush=True)
|
||||
os.rename(path, new_name)
|
||||
os.makedirs(path, exist_ok=True)
|
||||
|
||||
|
||||
@master_only
|
||||
def make_exp_dirs(opt):
|
||||
"""Make dirs for experiments."""
|
||||
path_opt = opt['path'].copy()
|
||||
if opt['is_train']:
|
||||
mkdir_and_rename(path_opt.pop('experiments_root'))
|
||||
else:
|
||||
mkdir_and_rename(path_opt.pop('results_root'))
|
||||
for key, path in path_opt.items():
|
||||
if ('strict_load' in key) or ('pretrain_network' in key) or ('resume' in key) or ('param_key' in key):
|
||||
continue
|
||||
else:
|
||||
os.makedirs(path, exist_ok=True)
|
||||
|
||||
|
||||
def scandir(dir_path, suffix=None, recursive=False, full_path=False):
|
||||
"""Scan a directory to find the interested files.
|
||||
|
||||
Args:
|
||||
dir_path (str): Path of the directory.
|
||||
suffix (str | tuple(str), optional): File suffix that we are
|
||||
interested in. Default: None.
|
||||
recursive (bool, optional): If set to True, recursively scan the
|
||||
directory. Default: False.
|
||||
full_path (bool, optional): If set to True, include the dir_path.
|
||||
Default: False.
|
||||
|
||||
Returns:
|
||||
A generator for all the interested files with relative paths.
|
||||
"""
|
||||
|
||||
if (suffix is not None) and not isinstance(suffix, (str, tuple)):
|
||||
raise TypeError('"suffix" must be a string or tuple of strings')
|
||||
|
||||
root = dir_path
|
||||
|
||||
def _scandir(dir_path, suffix, recursive):
|
||||
for entry in os.scandir(dir_path):
|
||||
if not entry.name.startswith('.') and entry.is_file():
|
||||
if full_path:
|
||||
return_path = entry.path
|
||||
else:
|
||||
return_path = osp.relpath(entry.path, root)
|
||||
|
||||
if suffix is None:
|
||||
yield return_path
|
||||
elif return_path.endswith(suffix):
|
||||
yield return_path
|
||||
else:
|
||||
if recursive:
|
||||
yield from _scandir(entry.path, suffix=suffix, recursive=recursive)
|
||||
else:
|
||||
continue
|
||||
|
||||
return _scandir(dir_path, suffix=suffix, recursive=recursive)
|
||||
|
||||
|
||||
def check_resume(opt, resume_iter):
|
||||
"""Check resume states and pretrain_network paths.
|
||||
|
||||
Args:
|
||||
opt (dict): Options.
|
||||
resume_iter (int): Resume iteration.
|
||||
"""
|
||||
if opt['path']['resume_state']:
|
||||
# get all the networks
|
||||
networks = [key for key in opt.keys() if key.startswith('network_')]
|
||||
flag_pretrain = False
|
||||
for network in networks:
|
||||
if opt['path'].get(f'pretrain_{network}') is not None:
|
||||
flag_pretrain = True
|
||||
if flag_pretrain:
|
||||
print('pretrain_network path will be ignored during resuming.')
|
||||
# set pretrained model paths
|
||||
for network in networks:
|
||||
name = f'pretrain_{network}'
|
||||
basename = network.replace('network_', '')
|
||||
if opt['path'].get('ignore_resume_networks') is None or (network
|
||||
not in opt['path']['ignore_resume_networks']):
|
||||
opt['path'][name] = osp.join(opt['path']['models'], f'net_{basename}_{resume_iter}.pth')
|
||||
print(f"Set {name} to {opt['path'][name]}")
|
||||
|
||||
# change param_key to params in resume
|
||||
param_keys = [key for key in opt['path'].keys() if key.startswith('param_key')]
|
||||
for param_key in param_keys:
|
||||
if opt['path'][param_key] == 'params_ema':
|
||||
opt['path'][param_key] = 'params'
|
||||
print(f'Set {param_key} to params')
|
||||
|
||||
|
||||
def sizeof_fmt(size, suffix='B'):
|
||||
"""Get human readable file size.
|
||||
|
||||
Args:
|
||||
size (int): File size.
|
||||
suffix (str): Suffix. Default: 'B'.
|
||||
|
||||
Return:
|
||||
str: Formatted file size.
|
||||
"""
|
||||
for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
|
||||
if abs(size) < 1024.0:
|
||||
return f'{size:3.1f} {unit}{suffix}'
|
||||
size /= 1024.0
|
||||
return f'{size:3.1f} Y{suffix}'
|
||||
@@ -1,210 +0,0 @@
|
||||
import argparse
|
||||
import os
|
||||
import random
|
||||
import torch
|
||||
import yaml
|
||||
from collections import OrderedDict
|
||||
from os import path as osp
|
||||
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils import set_random_seed
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils.dist_util import get_dist_info, init_dist, master_only
|
||||
|
||||
|
||||
def ordered_yaml():
|
||||
"""Support OrderedDict for yaml.
|
||||
|
||||
Returns:
|
||||
tuple: yaml Loader and Dumper.
|
||||
"""
|
||||
try:
|
||||
from yaml import CDumper as Dumper
|
||||
from yaml import CLoader as Loader
|
||||
except ImportError:
|
||||
from yaml import Dumper, Loader
|
||||
|
||||
_mapping_tag = yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG
|
||||
|
||||
def dict_representer(dumper, data):
|
||||
return dumper.represent_dict(data.items())
|
||||
|
||||
def dict_constructor(loader, node):
|
||||
return OrderedDict(loader.construct_pairs(node))
|
||||
|
||||
Dumper.add_representer(OrderedDict, dict_representer)
|
||||
Loader.add_constructor(_mapping_tag, dict_constructor)
|
||||
return Loader, Dumper
|
||||
|
||||
|
||||
def yaml_load(f):
|
||||
"""Load yaml file or string.
|
||||
|
||||
Args:
|
||||
f (str): File path or a python string.
|
||||
|
||||
Returns:
|
||||
dict: Loaded dict.
|
||||
"""
|
||||
if os.path.isfile(f):
|
||||
with open(f, 'r') as f:
|
||||
return yaml.load(f, Loader=ordered_yaml()[0])
|
||||
else:
|
||||
return yaml.load(f, Loader=ordered_yaml()[0])
|
||||
|
||||
|
||||
def dict2str(opt, indent_level=1):
|
||||
"""dict to string for printing options.
|
||||
|
||||
Args:
|
||||
opt (dict): Option dict.
|
||||
indent_level (int): Indent level. Default: 1.
|
||||
|
||||
Return:
|
||||
(str): Option string for printing.
|
||||
"""
|
||||
msg = '\n'
|
||||
for k, v in opt.items():
|
||||
if isinstance(v, dict):
|
||||
msg += ' ' * (indent_level * 2) + k + ':['
|
||||
msg += dict2str(v, indent_level + 1)
|
||||
msg += ' ' * (indent_level * 2) + ']\n'
|
||||
else:
|
||||
msg += ' ' * (indent_level * 2) + k + ': ' + str(v) + '\n'
|
||||
return msg
|
||||
|
||||
|
||||
def _postprocess_yml_value(value):
|
||||
# None
|
||||
if value == '~' or value.lower() == 'none':
|
||||
return None
|
||||
# bool
|
||||
if value.lower() == 'true':
|
||||
return True
|
||||
elif value.lower() == 'false':
|
||||
return False
|
||||
# !!float number
|
||||
if value.startswith('!!float'):
|
||||
return float(value.replace('!!float', ''))
|
||||
# number
|
||||
if value.isdigit():
|
||||
return int(value)
|
||||
elif value.replace('.', '', 1).isdigit() and value.count('.') < 2:
|
||||
return float(value)
|
||||
# list
|
||||
if value.startswith('['):
|
||||
return eval(value)
|
||||
# str
|
||||
return value
|
||||
|
||||
|
||||
def parse_options(root_path, is_train=True):
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('-opt', type=str, required=True, help='Path to option YAML file.')
|
||||
parser.add_argument('--launcher', choices=['none', 'pytorch', 'slurm'], default='none', help='job launcher')
|
||||
parser.add_argument('--auto_resume', action='store_true')
|
||||
parser.add_argument('--debug', action='store_true')
|
||||
parser.add_argument('--local_rank', type=int, default=0)
|
||||
parser.add_argument(
|
||||
'--force_yml', nargs='+', default=None, help='Force to update yml files. Examples: train:ema_decay=0.999')
|
||||
args = parser.parse_args()
|
||||
|
||||
# parse yml to dict
|
||||
opt = yaml_load(args.opt)
|
||||
|
||||
# distributed settings
|
||||
if args.launcher == 'none':
|
||||
opt['dist'] = False
|
||||
print('Disable distributed.', flush=True)
|
||||
else:
|
||||
opt['dist'] = True
|
||||
if args.launcher == 'slurm' and 'dist_params' in opt:
|
||||
init_dist(args.launcher, **opt['dist_params'])
|
||||
else:
|
||||
init_dist(args.launcher)
|
||||
opt['rank'], opt['world_size'] = get_dist_info()
|
||||
|
||||
# random seed
|
||||
seed = opt.get('manual_seed')
|
||||
if seed is None:
|
||||
seed = random.randint(1, 10000)
|
||||
opt['manual_seed'] = seed
|
||||
set_random_seed(seed + opt['rank'])
|
||||
|
||||
# force to update yml options
|
||||
if args.force_yml is not None:
|
||||
for entry in args.force_yml:
|
||||
# now do not support creating new keys
|
||||
keys, value = entry.split('=')
|
||||
keys, value = keys.strip(), value.strip()
|
||||
value = _postprocess_yml_value(value)
|
||||
eval_str = 'opt'
|
||||
for key in keys.split(':'):
|
||||
eval_str += f'["{key}"]'
|
||||
eval_str += '=value'
|
||||
# using exec function
|
||||
exec(eval_str)
|
||||
|
||||
opt['auto_resume'] = args.auto_resume
|
||||
opt['is_train'] = is_train
|
||||
|
||||
# debug setting
|
||||
if args.debug and not opt['name'].startswith('debug'):
|
||||
opt['name'] = 'debug_' + opt['name']
|
||||
|
||||
if opt['num_gpu'] == 'auto':
|
||||
opt['num_gpu'] = torch.cuda.device_count()
|
||||
|
||||
# datasets
|
||||
for phase, dataset in opt['datasets'].items():
|
||||
# for multiple datasets, e.g., val_1, val_2; test_1, test_2
|
||||
phase = phase.split('_')[0]
|
||||
dataset['phase'] = phase
|
||||
if 'scale' in opt:
|
||||
dataset['scale'] = opt['scale']
|
||||
if dataset.get('dataroot_gt') is not None:
|
||||
dataset['dataroot_gt'] = osp.expanduser(dataset['dataroot_gt'])
|
||||
if dataset.get('dataroot_lq') is not None:
|
||||
dataset['dataroot_lq'] = osp.expanduser(dataset['dataroot_lq'])
|
||||
|
||||
# paths
|
||||
for key, val in opt['path'].items():
|
||||
if (val is not None) and ('resume_state' in key or 'pretrain_network' in key):
|
||||
opt['path'][key] = osp.expanduser(val)
|
||||
|
||||
if is_train:
|
||||
experiments_root = osp.join(root_path, 'experiments', opt['name'])
|
||||
opt['path']['experiments_root'] = experiments_root
|
||||
opt['path']['models'] = osp.join(experiments_root, 'models')
|
||||
opt['path']['training_states'] = osp.join(experiments_root, 'training_states')
|
||||
opt['path']['log'] = experiments_root
|
||||
opt['path']['visualization'] = osp.join(experiments_root, 'visualization')
|
||||
|
||||
# change some options for debug mode
|
||||
if 'debug' in opt['name']:
|
||||
if 'val' in opt:
|
||||
opt['val']['val_freq'] = 8
|
||||
opt['logger']['print_freq'] = 1
|
||||
opt['logger']['save_checkpoint_freq'] = 8
|
||||
else: # test
|
||||
results_root = osp.join(root_path, 'results', opt['name'])
|
||||
opt['path']['results_root'] = results_root
|
||||
opt['path']['log'] = results_root
|
||||
opt['path']['visualization'] = osp.join(results_root, 'visualization')
|
||||
|
||||
return opt, args
|
||||
|
||||
|
||||
@master_only
|
||||
def copy_opt_file(opt_file, experiments_root):
|
||||
# copy the yml file to the experiment root
|
||||
import sys
|
||||
import time
|
||||
from shutil import copyfile
|
||||
cmd = ' '.join(sys.argv)
|
||||
filename = osp.join(experiments_root, osp.basename(opt_file))
|
||||
copyfile(opt_file, filename)
|
||||
|
||||
with open(filename, 'r+') as f:
|
||||
lines = f.readlines()
|
||||
lines.insert(0, f'# GENERATE TIME: {time.asctime()}\n# CMD:\n# {cmd}\n\n')
|
||||
f.seek(0)
|
||||
f.writelines(lines)
|
||||
@@ -1,83 +0,0 @@
|
||||
import re
|
||||
|
||||
|
||||
def read_data_from_tensorboard(log_path, tag):
|
||||
"""Get raw data (steps and values) from tensorboard events.
|
||||
|
||||
Args:
|
||||
log_path (str): Path to the tensorboard log.
|
||||
tag (str): tag to be read.
|
||||
"""
|
||||
from tensorboard.backend.event_processing.event_accumulator import EventAccumulator
|
||||
|
||||
# tensorboard event
|
||||
event_acc = EventAccumulator(log_path)
|
||||
event_acc.Reload()
|
||||
scalar_list = event_acc.Tags()['scalars']
|
||||
print('tag list: ', scalar_list)
|
||||
steps = [int(s.step) for s in event_acc.Scalars(tag)]
|
||||
values = [s.value for s in event_acc.Scalars(tag)]
|
||||
return steps, values
|
||||
|
||||
|
||||
def read_data_from_txt_2v(path, pattern, step_one=False):
|
||||
"""Read data from txt with 2 returned values (usually [step, value]).
|
||||
|
||||
Args:
|
||||
path (str): path to the txt file.
|
||||
pattern (str): re (regular expression) pattern.
|
||||
step_one (bool): add 1 to steps. Default: False.
|
||||
"""
|
||||
with open(path) as f:
|
||||
lines = f.readlines()
|
||||
lines = [line.strip() for line in lines]
|
||||
steps = []
|
||||
values = []
|
||||
|
||||
pattern = re.compile(pattern)
|
||||
for line in lines:
|
||||
match = pattern.match(line)
|
||||
if match:
|
||||
steps.append(int(match.group(1)))
|
||||
values.append(float(match.group(2)))
|
||||
if step_one:
|
||||
steps = [v + 1 for v in steps]
|
||||
return steps, values
|
||||
|
||||
|
||||
def read_data_from_txt_1v(path, pattern):
|
||||
"""Read data from txt with 1 returned values.
|
||||
|
||||
Args:
|
||||
path (str): path to the txt file.
|
||||
pattern (str): re (regular expression) pattern.
|
||||
"""
|
||||
with open(path) as f:
|
||||
lines = f.readlines()
|
||||
lines = [line.strip() for line in lines]
|
||||
data = []
|
||||
|
||||
pattern = re.compile(pattern)
|
||||
for line in lines:
|
||||
match = pattern.match(line)
|
||||
if match:
|
||||
data.append(float(match.group(1)))
|
||||
return data
|
||||
|
||||
|
||||
def smooth_data(values, smooth_weight):
|
||||
""" Smooth data using 1st-order IIR low-pass filter (what tensorflow does).
|
||||
|
||||
Reference: https://github.com/tensorflow/tensorboard/blob/f801ebf1f9fbfe2baee1ddd65714d0bccc640fb1/tensorboard/plugins/scalar/vz_line_chart/vz-line-chart.ts#L704 # noqa: E501
|
||||
|
||||
Args:
|
||||
values (list): A list of values to be smoothed.
|
||||
smooth_weight (float): Smooth weight.
|
||||
"""
|
||||
values_sm = []
|
||||
last_sm_value = values[0]
|
||||
for value in values:
|
||||
value_sm = last_sm_value * smooth_weight + (1 - smooth_weight) * value
|
||||
values_sm.append(value_sm)
|
||||
last_sm_value = value_sm
|
||||
return values_sm
|
||||
@@ -1,293 +0,0 @@
|
||||
import cv2
|
||||
import math
|
||||
import numpy as np
|
||||
import os
|
||||
import queue
|
||||
import threading
|
||||
import torch
|
||||
from basicsr.utils.download_util import load_file_from_url
|
||||
from torch.nn import functional as F
|
||||
|
||||
# ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
|
||||
class RealESRGANer():
|
||||
"""A helper class for upsampling images with RealESRGAN.
|
||||
|
||||
Args:
|
||||
scale (int): Upsampling scale factor used in the networks. It is usually 2 or 4.
|
||||
model_path (str): The path to the pretrained model. It can be urls (will first download it automatically).
|
||||
model (nn.Module): The defined network. Default: None.
|
||||
tile (int): As too large images result in the out of GPU memory issue, so this tile option will first crop
|
||||
input images into tiles, and then process each of them. Finally, they will be merged into one image.
|
||||
0 denotes for do not use tile. Default: 0.
|
||||
tile_pad (int): The pad size for each tile, to remove border artifacts. Default: 10.
|
||||
pre_pad (int): Pad the input images to avoid border artifacts. Default: 10.
|
||||
half (float): Whether to use half precision during inference. Default: False.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
scale,
|
||||
model_path,
|
||||
model=None,
|
||||
tile=0,
|
||||
tile_pad=10,
|
||||
pre_pad=10,
|
||||
half=False,
|
||||
device=None,
|
||||
gpu_id=None):
|
||||
self.scale = scale
|
||||
self.tile_size = tile
|
||||
self.tile_pad = tile_pad
|
||||
self.pre_pad = pre_pad
|
||||
self.mod_scale = None
|
||||
self.half = half
|
||||
|
||||
# initialize model
|
||||
if gpu_id:
|
||||
self.device = torch.device(
|
||||
f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu') if device is None else device
|
||||
else:
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device
|
||||
# if the model_path starts with https, it will first download models to the folder: realesrgan/weights
|
||||
if model_path.startswith('https://'):
|
||||
model_path = load_file_from_url(
|
||||
url=model_path, model_dir=os.path.join('weights/realesrgan'), progress=True, file_name=None)
|
||||
loadnet = torch.load(model_path, map_location=torch.device('cpu'))
|
||||
# prefer to use params_ema
|
||||
if 'params_ema' in loadnet:
|
||||
keyname = 'params_ema'
|
||||
else:
|
||||
keyname = 'params'
|
||||
model.load_state_dict(loadnet[keyname], strict=True)
|
||||
model.eval()
|
||||
self.model = model.to(self.device)
|
||||
if self.half:
|
||||
self.model = self.model.half()
|
||||
|
||||
def pre_process(self, img):
|
||||
"""Pre-process, such as pre-pad and mod pad, so that the images can be divisible
|
||||
"""
|
||||
img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()
|
||||
self.img = img.unsqueeze(0).to(self.device)
|
||||
if self.half:
|
||||
self.img = self.img.half()
|
||||
|
||||
# pre_pad
|
||||
if self.pre_pad != 0:
|
||||
self.img = F.pad(self.img, (0, self.pre_pad, 0, self.pre_pad), 'reflect')
|
||||
# mod pad for divisible borders
|
||||
if self.scale == 2:
|
||||
self.mod_scale = 2
|
||||
elif self.scale == 1:
|
||||
self.mod_scale = 4
|
||||
if self.mod_scale is not None:
|
||||
self.mod_pad_h, self.mod_pad_w = 0, 0
|
||||
_, _, h, w = self.img.size()
|
||||
if (h % self.mod_scale != 0):
|
||||
self.mod_pad_h = (self.mod_scale - h % self.mod_scale)
|
||||
if (w % self.mod_scale != 0):
|
||||
self.mod_pad_w = (self.mod_scale - w % self.mod_scale)
|
||||
self.img = F.pad(self.img, (0, self.mod_pad_w, 0, self.mod_pad_h), 'reflect')
|
||||
|
||||
def process(self):
|
||||
# model inference
|
||||
self.output = self.model(self.img)
|
||||
|
||||
def tile_process(self):
|
||||
"""It will first crop input images to tiles, and then process each tile.
|
||||
Finally, all the processed tiles are merged into one images.
|
||||
|
||||
Modified from: https://github.com/ata4/esrgan-launcher
|
||||
"""
|
||||
batch, channel, height, width = self.img.shape
|
||||
output_height = height * self.scale
|
||||
output_width = width * self.scale
|
||||
output_shape = (batch, channel, output_height, output_width)
|
||||
|
||||
# start with black image
|
||||
self.output = self.img.new_zeros(output_shape)
|
||||
tiles_x = math.ceil(width / self.tile_size)
|
||||
tiles_y = math.ceil(height / self.tile_size)
|
||||
|
||||
# loop over all tiles
|
||||
for y in range(tiles_y):
|
||||
for x in range(tiles_x):
|
||||
# extract tile from input image
|
||||
ofs_x = x * self.tile_size
|
||||
ofs_y = y * self.tile_size
|
||||
# input tile area on total image
|
||||
input_start_x = ofs_x
|
||||
input_end_x = min(ofs_x + self.tile_size, width)
|
||||
input_start_y = ofs_y
|
||||
input_end_y = min(ofs_y + self.tile_size, height)
|
||||
|
||||
# input tile area on total image with padding
|
||||
input_start_x_pad = max(input_start_x - self.tile_pad, 0)
|
||||
input_end_x_pad = min(input_end_x + self.tile_pad, width)
|
||||
input_start_y_pad = max(input_start_y - self.tile_pad, 0)
|
||||
input_end_y_pad = min(input_end_y + self.tile_pad, height)
|
||||
|
||||
# input tile dimensions
|
||||
input_tile_width = input_end_x - input_start_x
|
||||
input_tile_height = input_end_y - input_start_y
|
||||
tile_idx = y * tiles_x + x + 1
|
||||
input_tile = self.img[:, :, input_start_y_pad:input_end_y_pad, input_start_x_pad:input_end_x_pad]
|
||||
|
||||
# upscale tile
|
||||
try:
|
||||
with torch.no_grad():
|
||||
output_tile = self.model(input_tile)
|
||||
except RuntimeError as error:
|
||||
print('Error', error)
|
||||
# print(f'\tTile {tile_idx}/{tiles_x * tiles_y}')
|
||||
|
||||
# output tile area on total image
|
||||
output_start_x = input_start_x * self.scale
|
||||
output_end_x = input_end_x * self.scale
|
||||
output_start_y = input_start_y * self.scale
|
||||
output_end_y = input_end_y * self.scale
|
||||
|
||||
# output tile area without padding
|
||||
output_start_x_tile = (input_start_x - input_start_x_pad) * self.scale
|
||||
output_end_x_tile = output_start_x_tile + input_tile_width * self.scale
|
||||
output_start_y_tile = (input_start_y - input_start_y_pad) * self.scale
|
||||
output_end_y_tile = output_start_y_tile + input_tile_height * self.scale
|
||||
|
||||
# put tile into output image
|
||||
self.output[:, :, output_start_y:output_end_y,
|
||||
output_start_x:output_end_x] = output_tile[:, :, output_start_y_tile:output_end_y_tile,
|
||||
output_start_x_tile:output_end_x_tile]
|
||||
|
||||
def post_process(self):
|
||||
# remove extra pad
|
||||
if self.mod_scale is not None:
|
||||
_, _, h, w = self.output.size()
|
||||
self.output = self.output[:, :, 0:h - self.mod_pad_h * self.scale, 0:w - self.mod_pad_w * self.scale]
|
||||
# remove prepad
|
||||
if self.pre_pad != 0:
|
||||
_, _, h, w = self.output.size()
|
||||
self.output = self.output[:, :, 0:h - self.pre_pad * self.scale, 0:w - self.pre_pad * self.scale]
|
||||
return self.output
|
||||
|
||||
@torch.no_grad()
|
||||
def enhance(self, img, outscale=None, alpha_upsampler='realesrgan'):
|
||||
h_input, w_input = img.shape[0:2]
|
||||
# img: numpy
|
||||
img = img.astype(np.float32)
|
||||
if np.max(img) > 256: # 16-bit image
|
||||
max_range = 65535
|
||||
print('\tInput is a 16-bit image')
|
||||
else:
|
||||
max_range = 255
|
||||
img = img / max_range
|
||||
if len(img.shape) == 2: # gray image
|
||||
img_mode = 'L'
|
||||
img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
|
||||
elif img.shape[2] == 4: # RGBA image with alpha channel
|
||||
img_mode = 'RGBA'
|
||||
alpha = img[:, :, 3]
|
||||
img = img[:, :, 0:3]
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
if alpha_upsampler == 'realesrgan':
|
||||
alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2RGB)
|
||||
else:
|
||||
img_mode = 'RGB'
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
|
||||
# ------------------- process image (without the alpha channel) ------------------- #
|
||||
self.pre_process(img)
|
||||
if self.tile_size > 0:
|
||||
self.tile_process()
|
||||
else:
|
||||
self.process()
|
||||
output_img = self.post_process()
|
||||
output_img = output_img.data.squeeze().float().cpu().clamp_(0, 1).numpy()
|
||||
output_img = np.transpose(output_img[[2, 1, 0], :, :], (1, 2, 0))
|
||||
if img_mode == 'L':
|
||||
output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
# ------------------- process the alpha channel if necessary ------------------- #
|
||||
if img_mode == 'RGBA':
|
||||
if alpha_upsampler == 'realesrgan':
|
||||
self.pre_process(alpha)
|
||||
if self.tile_size > 0:
|
||||
self.tile_process()
|
||||
else:
|
||||
self.process()
|
||||
output_alpha = self.post_process()
|
||||
output_alpha = output_alpha.data.squeeze().float().cpu().clamp_(0, 1).numpy()
|
||||
output_alpha = np.transpose(output_alpha[[2, 1, 0], :, :], (1, 2, 0))
|
||||
output_alpha = cv2.cvtColor(output_alpha, cv2.COLOR_BGR2GRAY)
|
||||
else: # use the cv2 resize for alpha channel
|
||||
h, w = alpha.shape[0:2]
|
||||
output_alpha = cv2.resize(alpha, (w * self.scale, h * self.scale), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# merge the alpha channel
|
||||
output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2BGRA)
|
||||
output_img[:, :, 3] = output_alpha
|
||||
|
||||
# ------------------------------ return ------------------------------ #
|
||||
if max_range == 65535: # 16-bit image
|
||||
output = (output_img * 65535.0).round().astype(np.uint16)
|
||||
else:
|
||||
output = (output_img * 255.0).round().astype(np.uint8)
|
||||
|
||||
if outscale is not None and outscale != float(self.scale):
|
||||
output = cv2.resize(
|
||||
output, (
|
||||
int(w_input * outscale),
|
||||
int(h_input * outscale),
|
||||
), interpolation=cv2.INTER_LANCZOS4)
|
||||
|
||||
return output, img_mode
|
||||
|
||||
|
||||
class PrefetchReader(threading.Thread):
|
||||
"""Prefetch images.
|
||||
|
||||
Args:
|
||||
img_list (list[str]): A image list of image paths to be read.
|
||||
num_prefetch_queue (int): Number of prefetch queue.
|
||||
"""
|
||||
|
||||
def __init__(self, img_list, num_prefetch_queue):
|
||||
super().__init__()
|
||||
self.que = queue.Queue(num_prefetch_queue)
|
||||
self.img_list = img_list
|
||||
|
||||
def run(self):
|
||||
for img_path in self.img_list:
|
||||
img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
|
||||
self.que.put(img)
|
||||
|
||||
self.que.put(None)
|
||||
|
||||
def __next__(self):
|
||||
next_item = self.que.get()
|
||||
if next_item is None:
|
||||
raise StopIteration
|
||||
return next_item
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
|
||||
class IOConsumer(threading.Thread):
|
||||
|
||||
def __init__(self, opt, que, qid):
|
||||
super().__init__()
|
||||
self._queue = que
|
||||
self.qid = qid
|
||||
self.opt = opt
|
||||
|
||||
def run(self):
|
||||
while True:
|
||||
msg = self._queue.get()
|
||||
if isinstance(msg, str) and msg == 'quit':
|
||||
break
|
||||
|
||||
output = msg['output']
|
||||
save_path = msg['save_path']
|
||||
cv2.imwrite(save_path, output)
|
||||
print(f'IO worker {self.qid} is done.')
|
||||
@@ -1,88 +0,0 @@
|
||||
# Modified from: https://github.com/facebookresearch/fvcore/blob/master/fvcore/common/registry.py # noqa: E501
|
||||
|
||||
|
||||
class Registry():
|
||||
"""
|
||||
The registry that provides name -> object mapping, to support third-party
|
||||
users' custom modules.
|
||||
|
||||
To create a registry (e.g. a backbone registry):
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
BACKBONE_REGISTRY = Registry('BACKBONE')
|
||||
|
||||
To register an object:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@BACKBONE_REGISTRY.register()
|
||||
class MyBackbone():
|
||||
...
|
||||
|
||||
Or:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
BACKBONE_REGISTRY.register(MyBackbone)
|
||||
"""
|
||||
|
||||
def __init__(self, name):
|
||||
"""
|
||||
Args:
|
||||
name (str): the name of this registry
|
||||
"""
|
||||
self._name = name
|
||||
self._obj_map = {}
|
||||
|
||||
def _do_register(self, name, obj, suffix=None):
|
||||
if isinstance(suffix, str):
|
||||
name = name + '_' + suffix
|
||||
|
||||
assert (name not in self._obj_map), (f"An object named '{name}' was already registered "
|
||||
f"in '{self._name}' registry!")
|
||||
self._obj_map[name] = obj
|
||||
|
||||
def register(self, obj=None, suffix=None):
|
||||
"""
|
||||
Register the given object under the the name `obj.__name__`.
|
||||
Can be used as either a decorator or not.
|
||||
See docstring of this class for usage.
|
||||
"""
|
||||
if obj is None:
|
||||
# used as a decorator
|
||||
def deco(func_or_class):
|
||||
name = func_or_class.__name__
|
||||
self._do_register(name, func_or_class, suffix)
|
||||
return func_or_class
|
||||
|
||||
return deco
|
||||
|
||||
# used as a function call
|
||||
name = obj.__name__
|
||||
self._do_register(name, obj, suffix)
|
||||
|
||||
def get(self, name, suffix='basicsr'):
|
||||
ret = self._obj_map.get(name)
|
||||
if ret is None:
|
||||
ret = self._obj_map.get(name + '_' + suffix)
|
||||
print(f'Name {name} is not found, use name: {name}_{suffix}!')
|
||||
if ret is None:
|
||||
raise KeyError(f"No object named '{name}' found in '{self._name}' registry!")
|
||||
return ret
|
||||
|
||||
def __contains__(self, name):
|
||||
return name in self._obj_map
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._obj_map.items())
|
||||
|
||||
def keys(self):
|
||||
return self._obj_map.keys()
|
||||
|
||||
|
||||
DATASET_REGISTRY = Registry('dataset')
|
||||
ARCH_REGISTRY = Registry('arch')
|
||||
MODEL_REGISTRY = Registry('model')
|
||||
LOSS_REGISTRY = Registry('loss')
|
||||
METRIC_REGISTRY = Registry('metric')
|
||||
@@ -1,43 +0,0 @@
|
||||
# Configuration for Cog ⚙️
|
||||
# Reference: https://cog.run/yaml
|
||||
|
||||
build:
|
||||
# set to true if your model requires a GPU
|
||||
gpu: true
|
||||
|
||||
# a list of ubuntu apt packages to install
|
||||
system_packages:
|
||||
- "libgl1-mesa-glx"
|
||||
- "libglib2.0-0"
|
||||
|
||||
# python version in the form '3.11' or '3.11.4'
|
||||
python_version: "3.10"
|
||||
|
||||
# a list of packages in the format <package-name>==<version>
|
||||
python_packages:
|
||||
- "torch==2.4.0"
|
||||
- "torchvision==0.19.0"
|
||||
- "torchaudio==2.4.0"
|
||||
- "git+https://github.com/huggingface/transformers.git@v4.46.0"
|
||||
- "xformers==0.0.27.post2"
|
||||
- "git+https://github.com/zsyOAOA/InvSR"
|
||||
- "scikit-image==0.24.0"
|
||||
- "albumentations==1.4.3"
|
||||
- "opencv-python==4.10.0.84"
|
||||
- "bitsandbytes==0.45.0"
|
||||
- "sentencepiece==0.2.0"
|
||||
- "protobuf==5.29.1"
|
||||
- "python-box==7.3.0"
|
||||
- "omegaconf==2.3.0"
|
||||
- "loguru==0.7.3"
|
||||
- "einops==0.8.0"
|
||||
- "pydantic==1.10.11"
|
||||
- "accelerate==0.34.2"
|
||||
|
||||
# commands run after the environment is setup
|
||||
run:
|
||||
#- "pip install git+https://github.com/zsyOAOA/InvSR"
|
||||
- "echo env is ready!"
|
||||
|
||||
# predict.py defines how predictions are run on your model
|
||||
predict: "predict.py:Predictor"
|
||||
@@ -1,34 +0,0 @@
|
||||
services:
|
||||
gradio:
|
||||
container_name: invsr-gradio
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
ports:
|
||||
- 7860:7860 # PORT by default is 7860
|
||||
volumes:
|
||||
- invsr_weights:/invsr/weights/
|
||||
ipc: host # shm
|
||||
stdin_open: true
|
||||
tty: true
|
||||
deploy:
|
||||
resources:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: all
|
||||
capabilities:
|
||||
- gpu
|
||||
|
||||
volumes:
|
||||
invsr_weights:
|
||||
name: invsr_weights
|
||||
driver: local
|
||||
driver_opts:
|
||||
type: none
|
||||
o: bind
|
||||
device: ./weights
|
||||
|
||||
networks:
|
||||
default:
|
||||
name: invsr_net
|
||||
@@ -1,101 +0,0 @@
|
||||
{
|
||||
"seed": 12345,
|
||||
"basesr": {
|
||||
"sf": 4,
|
||||
"chopping": {
|
||||
"pch_size": 128,
|
||||
"weight_type": "Gaussian",
|
||||
"extra_bs": 8
|
||||
}
|
||||
},
|
||||
"tiled_vae": true,
|
||||
"latent_tiled_size": 128,
|
||||
"sample_tiled_size": 1024,
|
||||
"gradient_checkpointing_vae": true,
|
||||
"sliced_vae": false,
|
||||
"cfg_scale": 1.0,
|
||||
"start_timesteps": 200,
|
||||
"color_fix": "",
|
||||
"base_model": "sd-turbo",
|
||||
"sd_pipe": {
|
||||
"target": "diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline",
|
||||
"enable_grad_checkpoint": true,
|
||||
"params": {
|
||||
"pretrained_model_name_or_path": "stabilityai/sd-turbo",
|
||||
"cache_dir": "./weights",
|
||||
"use_safetensors": true,
|
||||
"torch_dtype": "torch.float16"
|
||||
}
|
||||
},
|
||||
"model_start": {
|
||||
"target": "diffusers.models.autoencoders.NoisePredictor",
|
||||
"ckpt_path": "weights/noise_predictor_sd_turbo_v5.pth",
|
||||
"params": {
|
||||
"in_channels": 3,
|
||||
"down_block_types": [
|
||||
"AttnDownBlock2D",
|
||||
"AttnDownBlock2D"
|
||||
],
|
||||
"up_block_types": [
|
||||
"AttnUpBlock2D",
|
||||
"AttnUpBlock2D"
|
||||
],
|
||||
"block_out_channels": [
|
||||
256,
|
||||
512
|
||||
],
|
||||
"layers_per_block": [
|
||||
3,
|
||||
3
|
||||
],
|
||||
"act_fn": "silu",
|
||||
"latent_channels": 4,
|
||||
"norm_num_groups": 32,
|
||||
"sample_size": 128,
|
||||
"mid_block_add_attention": true,
|
||||
"resnet_time_scale_shift": "default",
|
||||
"temb_channels": 512,
|
||||
"attention_head_dim": 64,
|
||||
"freq_shift": 0,
|
||||
"flip_sin_to_cos": true,
|
||||
"double_z": true
|
||||
}
|
||||
},
|
||||
"model_middle": {
|
||||
"target": "diffusers.models.autoencoders.NoisePredictor",
|
||||
"params": {
|
||||
"in_channels": 3,
|
||||
"down_block_types": [
|
||||
"AttnDownBlock2D",
|
||||
"AttnDownBlock2D"
|
||||
],
|
||||
"up_block_types": [
|
||||
"AttnUpBlock2D",
|
||||
"AttnUpBlock2D"
|
||||
],
|
||||
"block_out_channels": [
|
||||
256,
|
||||
512
|
||||
],
|
||||
"layers_per_block": [
|
||||
3,
|
||||
3
|
||||
],
|
||||
"act_fn": "silu",
|
||||
"latent_channels": 4,
|
||||
"norm_num_groups": 32,
|
||||
"sample_size": 128,
|
||||
"mid_block_add_attention": true,
|
||||
"resnet_time_scale_shift": "default",
|
||||
"temb_channels": 512,
|
||||
"attention_head_dim": 64,
|
||||
"freq_shift": 0,
|
||||
"flip_sin_to_cos": true,
|
||||
"double_z": true
|
||||
}
|
||||
},
|
||||
"timesteps": [
|
||||
200
|
||||
],
|
||||
"bs": 1
|
||||
}
|
||||
@@ -1,4 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding:utf-8 -*-
|
||||
# Power by Zongsheng Yue 2022-06-07 17:27:22
|
||||
|
||||
@@ -1,215 +0,0 @@
|
||||
import random
|
||||
import json
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from typing import Iterable
|
||||
from omegaconf import ListConfig
|
||||
|
||||
import cv2
|
||||
import torch
|
||||
from functools import partial
|
||||
import torchvision as thv
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from ..utils import util_sisr
|
||||
from ..utils import util_image
|
||||
from ..utils import util_common
|
||||
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.transforms import augment
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.data.realesrgan_dataset import RealESRGANDataset
|
||||
|
||||
def get_transforms(transform_type, kwargs):
|
||||
'''
|
||||
Accepted optins in kwargs.
|
||||
mean: scaler or sequence, for nornmalization
|
||||
std: scaler or sequence, for nornmalization
|
||||
crop_size: int or sequence, random or center cropping
|
||||
scale, out_shape: for Bicubic
|
||||
min_max: tuple or list with length 2, for cliping
|
||||
'''
|
||||
if transform_type == 'default':
|
||||
transform = thv.transforms.Compose([
|
||||
thv.transforms.ToTensor(),
|
||||
thv.transforms.Normalize(mean=kwargs.get('mean', 0.5), std=kwargs.get('std', 0.5)),
|
||||
])
|
||||
elif transform_type == 'resize_ccrop_norm':
|
||||
transform = thv.transforms.Compose([
|
||||
util_image.SmallestMaxSize(
|
||||
max_size=kwargs.get('size'),
|
||||
interpolation=kwargs.get('interpolation'),
|
||||
),
|
||||
thv.transforms.ToTensor(),
|
||||
thv.transforms.CenterCrop(size=kwargs.get('size', None)),
|
||||
thv.transforms.Normalize(mean=kwargs.get('mean', 0.5), std=kwargs.get('std', 0.5)),
|
||||
])
|
||||
elif transform_type == 'ccrop_norm':
|
||||
transform = thv.transforms.Compose([
|
||||
thv.transforms.ToTensor(),
|
||||
thv.transforms.CenterCrop(size=kwargs.get('size', None)),
|
||||
thv.transforms.Normalize(mean=kwargs.get('mean', 0.5), std=kwargs.get('std', 0.5)),
|
||||
])
|
||||
elif transform_type == 'rcrop_aug_norm':
|
||||
transform = thv.transforms.Compose([
|
||||
util_image.RandomCrop(pch_size=kwargs.get('pch_size', 256)),
|
||||
util_image.SpatialAug(
|
||||
only_hflip=kwargs.get('only_hflip', False),
|
||||
only_vflip=kwargs.get('only_vflip', False),
|
||||
only_hvflip=kwargs.get('only_hvflip', False),
|
||||
),
|
||||
util_image.ToTensor(max_value=kwargs.get('max_value')), # (ndarray, hwc) --> (Tensor, chw)
|
||||
thv.transforms.Normalize(mean=kwargs.get('mean', 0.5), std=kwargs.get('std', 0.5)),
|
||||
])
|
||||
elif transform_type == 'aug_norm':
|
||||
transform = thv.transforms.Compose([
|
||||
util_image.SpatialAug(
|
||||
only_hflip=kwargs.get('only_hflip', False),
|
||||
only_vflip=kwargs.get('only_vflip', False),
|
||||
only_hvflip=kwargs.get('only_hvflip', False),
|
||||
),
|
||||
util_image.ToTensor(), # hwc --> chw
|
||||
thv.transforms.Normalize(mean=kwargs.get('mean', 0.5), std=kwargs.get('std', 0.5)),
|
||||
])
|
||||
else:
|
||||
raise ValueError(f'Unexpected transform_variant {transform_variant}')
|
||||
return transform
|
||||
|
||||
def create_dataset(dataset_config):
|
||||
if dataset_config['type'] == 'base':
|
||||
dataset = BaseData(**dataset_config['params'])
|
||||
elif dataset_config['type'] == 'base_meta':
|
||||
dataset = BaseDataMetaCond(**dataset_config['params'])
|
||||
elif dataset_config['type'] == 'realesrgan':
|
||||
dataset = RealESRGANDataset(dataset_config['params'])
|
||||
else:
|
||||
raise NotImplementedError(f"{dataset_config['type']}")
|
||||
|
||||
return dataset
|
||||
|
||||
class BaseData(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
dir_path,
|
||||
txt_path=None,
|
||||
transform_type='default',
|
||||
transform_kwargs={'mean':0.0, 'std':1.0},
|
||||
extra_dir_path=None,
|
||||
extra_transform_type=None,
|
||||
extra_transform_kwargs=None,
|
||||
length=None,
|
||||
need_path=False,
|
||||
im_exts=['png', 'jpg', 'jpeg', 'JPEG', 'bmp'],
|
||||
recursive=False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
file_paths_all = []
|
||||
if dir_path is not None:
|
||||
file_paths_all.extend(util_common.scan_files_from_folder(dir_path, im_exts, recursive))
|
||||
if txt_path is not None:
|
||||
file_paths_all.extend(util_common.readline_txt(txt_path))
|
||||
|
||||
self.file_paths = file_paths_all if length is None else random.sample(file_paths_all, length)
|
||||
self.file_paths_all = file_paths_all
|
||||
|
||||
self.length = length
|
||||
self.need_path = need_path
|
||||
self.transform = get_transforms(transform_type, transform_kwargs)
|
||||
|
||||
self.extra_dir_path = extra_dir_path
|
||||
if extra_dir_path is not None:
|
||||
assert extra_transform_type is not None
|
||||
self.extra_transform = get_transforms(extra_transform_type, extra_transform_kwargs)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.file_paths)
|
||||
|
||||
def __getitem__(self, index):
|
||||
im_path_base = self.file_paths[index]
|
||||
im_base = util_image.imread(im_path_base, chn='rgb', dtype='float32')
|
||||
|
||||
im_target = self.transform(im_base)
|
||||
out = {'image':im_target, 'lq':im_target}
|
||||
|
||||
if self.extra_dir_path is not None:
|
||||
im_path_extra = Path(self.extra_dir_path) / Path(im_path_base).name
|
||||
im_extra = util_image.imread(im_path_extra, chn='rgb', dtype='float32')
|
||||
im_extra = self.extra_transform(im_extra)
|
||||
out['gt'] = im_extra
|
||||
|
||||
if self.need_path:
|
||||
out['path'] = im_path_base
|
||||
|
||||
return out
|
||||
|
||||
def reset_dataset(self):
|
||||
self.file_paths = random.sample(self.file_paths_all, self.length)
|
||||
|
||||
class BaseDataMetaCond(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
meta_dir,
|
||||
transform_type='default',
|
||||
transform_kwargs={'mean':0.5, 'std':0.5},
|
||||
length=None,
|
||||
need_path=False,
|
||||
cond_key='canny',
|
||||
cond_transform_type='default',
|
||||
cond_transform_kwargs={'mean':0.5, 'std':0.5},
|
||||
):
|
||||
super().__init__()
|
||||
if not isinstance(meta_dir, ListConfig):
|
||||
meta_dir = [meta_dir,]
|
||||
meta_list = []
|
||||
# for current_dir in meta_dir:
|
||||
# for json_path in Path(current_dir).glob("*.json"):
|
||||
# with open(json_path, 'r') as json_file:
|
||||
# meta_info = json.load(json_file)
|
||||
# meta_list.append(meta_info)
|
||||
for current_dir in meta_dir:
|
||||
meta_list.extend(sorted([str(x) for x in Path(current_dir).glob("*.json")]))
|
||||
self.meta_list = meta_list if length is None else meta_list[:length]
|
||||
|
||||
self.cond_key = cond_key
|
||||
self.length = length
|
||||
self.need_path = need_path
|
||||
self.transform = get_transforms(transform_type, transform_kwargs)
|
||||
self.cond_trasform = get_transforms(cond_transform_type, cond_transform_kwargs)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.meta_list)
|
||||
|
||||
def __getitem__(self, index):
|
||||
# meta_info = self.meta_list[index]
|
||||
json_path = self.meta_list[index]
|
||||
with open(json_path, 'r') as json_file:
|
||||
meta_info = json.load(json_file)
|
||||
|
||||
# images
|
||||
im_path = meta_info['source']
|
||||
im_source = util_image.imread(im_path, chn='rgb', dtype='uint8')
|
||||
im_source = self.transform(im_source)
|
||||
out = {'image': im_source,}
|
||||
if self.need_path:
|
||||
out['path'] = im_path
|
||||
|
||||
# latent
|
||||
if 'latent' in meta_info:
|
||||
latent_path = meta_info['latent']
|
||||
out['latent'] = np.load(latent_path)
|
||||
|
||||
# prompt
|
||||
out['txt'] = meta_info['prompt']
|
||||
|
||||
# condition
|
||||
cond_key = self.cond_key
|
||||
cond_path = meta_info[cond_key]
|
||||
if cond_key == 'canny':
|
||||
cond = util_image.imread(cond_path, chn='gray', dtype='uint8')[:, :, None]
|
||||
elif cond_key == 'seg':
|
||||
cond = util_image.imread(cond_path, chn='rgb', dtype='uint8')
|
||||
else:
|
||||
raise ValueError(f"Unexpected cond key: {cond_key}")
|
||||
cond = self.cond_trasform(cond)
|
||||
out['cond'] = cond
|
||||
|
||||
return out
|
||||
@@ -1,151 +0,0 @@
|
||||
name: invsr
|
||||
channels:
|
||||
- conda-forge
|
||||
- defaults
|
||||
dependencies:
|
||||
- _libgcc_mutex=0.1=conda_forge
|
||||
- _openmp_mutex=4.5=2_gnu
|
||||
- bzip2=1.0.8=h4bc722e_7
|
||||
- ca-certificates=2024.8.30=hbcca054_0
|
||||
- ld_impl_linux-64=2.43=h712a8e2_1
|
||||
- libffi=3.4.2=h7f98852_5
|
||||
- libgcc=14.1.0=h77fa898_1
|
||||
- libgcc-ng=14.1.0=h69a702a_1
|
||||
- libgomp=14.1.0=h77fa898_1
|
||||
- libnsl=2.0.1=hd590300_0
|
||||
- libsqlite=3.46.1=hadc24fc_0
|
||||
- libuuid=2.38.1=h0b41bf4_0
|
||||
- libzlib=1.3.1=hb9d3cd8_2
|
||||
- ncurses=6.5=he02047a_1
|
||||
- openssl=3.3.2=hb9d3cd8_0
|
||||
- pip=24.2=pyh8b19718_1
|
||||
- python=3.10.0=h543edf9_3_cpython
|
||||
- readline=8.2=h8228510_1
|
||||
- setuptools=75.1.0=pyhd8ed1ab_0
|
||||
- sqlite=3.46.1=h9eae976_0
|
||||
- tk=8.6.13=noxft_h4845f30_101
|
||||
- wheel=0.44.0=pyhd8ed1ab_0
|
||||
- xz=5.2.6=h166bdaf_0
|
||||
- pip:
|
||||
- absl-py==2.1.0
|
||||
- accelerate==0.34.2
|
||||
- addict==2.4.0
|
||||
- albucore==0.0.17
|
||||
- albumentations==1.4.18
|
||||
- annotated-types==0.7.0
|
||||
- antlr4-python3-runtime==4.9.3
|
||||
- asttokens==2.4.1
|
||||
- bitsandbytes==0.44.1
|
||||
- certifi==2022.12.7
|
||||
- charset-normalizer==2.1.1
|
||||
- colorama==0.4.6
|
||||
- contourpy==1.3.0
|
||||
- cycler==0.12.1
|
||||
- decorator==5.1.1
|
||||
- einops==0.8.0
|
||||
- eval-type-backport==0.2.0
|
||||
- exceptiongroup==1.2.2
|
||||
- executing==2.1.0
|
||||
- facexlib==0.3.0
|
||||
- filelock==3.13.1
|
||||
- filterpy==1.4.5
|
||||
- fonttools==4.54.1
|
||||
- fsspec==2024.2.0
|
||||
- ftfy==6.2.3
|
||||
- future==1.0.0
|
||||
- grpcio==1.66.1
|
||||
- huggingface-hub==0.25.1
|
||||
- icecream==2.1.3
|
||||
- idna==3.4
|
||||
- imageio==2.35.1
|
||||
- imgaug==0.4.0
|
||||
- importlib-metadata==8.5.0
|
||||
- ipdb==0.13.13
|
||||
- ipython==8.28.0
|
||||
- jedi==0.19.1
|
||||
- jinja2==3.1.3
|
||||
- joblib==1.4.2
|
||||
- kiwisolver==1.4.7
|
||||
- lazy-loader==0.4
|
||||
- llvmlite==0.43.0
|
||||
- lmdb==1.5.1
|
||||
- loguru==0.7.2
|
||||
- lpips==0.1.4
|
||||
- markdown==3.7
|
||||
- markupsafe==2.1.5
|
||||
- matplotlib==3.9.2
|
||||
- matplotlib-inline==0.1.7
|
||||
- mpmath==1.3.0
|
||||
- networkx==3.2.1
|
||||
- ninja==1.11.1.1
|
||||
- numba==0.60.0
|
||||
- numpy==1.26.3
|
||||
- nvidia-cublas-cu12==12.1.3.1
|
||||
- nvidia-cuda-cupti-cu12==12.1.105
|
||||
- nvidia-cuda-nvrtc-cu12==12.1.105
|
||||
- nvidia-cuda-runtime-cu12==12.1.105
|
||||
- nvidia-cudnn-cu12==9.1.0.70
|
||||
- nvidia-cufft-cu12==11.0.2.54
|
||||
- nvidia-curand-cu12==10.3.2.106
|
||||
- nvidia-cusolver-cu12==11.4.5.107
|
||||
- nvidia-cusparse-cu12==12.1.0.106
|
||||
- nvidia-nccl-cu12==2.20.5
|
||||
- nvidia-nvjitlink-cu12==12.1.105
|
||||
- nvidia-nvtx-cu12==12.1.105
|
||||
- omegaconf==2.3.0
|
||||
- openai-clip==1.0.1
|
||||
- opencv-python==4.10.0.84
|
||||
- opencv-python-headless==4.10.0.84
|
||||
- packaging==24.1
|
||||
- pandas==2.2.3
|
||||
- parso==0.8.4
|
||||
- pexpect==4.9.0
|
||||
- pillow==10.2.0
|
||||
- platformdirs==4.3.6
|
||||
- prompt-toolkit==3.0.48
|
||||
- protobuf==5.28.2
|
||||
- psutil==6.0.0
|
||||
- ptyprocess==0.7.0
|
||||
- pure-eval==0.2.3
|
||||
- pydantic==2.9.2
|
||||
- pydantic-core==2.23.4
|
||||
- pygments==2.18.0
|
||||
- pyiqa==0.1.12
|
||||
- pyparsing==3.1.4
|
||||
- python-box==7.2.0
|
||||
- python-dateutil==2.9.0.post0
|
||||
- pytz==2024.2
|
||||
- pyyaml==6.0.2
|
||||
- regex==2024.9.11
|
||||
- requests==2.28.1
|
||||
- safetensors==0.4.5
|
||||
- scikit-image==0.24.0
|
||||
- scikit-learn==1.5.2
|
||||
- scipy==1.14.1
|
||||
- sentencepiece==0.2.0
|
||||
- shapely==2.0.6
|
||||
- six==1.16.0
|
||||
- stack-data==0.6.3
|
||||
- sympy==1.12
|
||||
- tensorboard==2.18.0
|
||||
- tensorboard-data-server==0.7.2
|
||||
- threadpoolctl==3.5.0
|
||||
- tifffile==2024.9.20
|
||||
- timm==1.0.9
|
||||
- tokenizers==0.20.1
|
||||
- tomli==2.0.1
|
||||
- torch==2.4.0+cu121
|
||||
- torchaudio==2.4.0+cu121
|
||||
- torchvision==0.19.0+cu121
|
||||
- tqdm==4.66.5
|
||||
- traitlets==5.14.3
|
||||
- transformers==4.46.0
|
||||
- triton==3.0.0
|
||||
- typing-extensions==4.9.0
|
||||
- tzdata==2024.2
|
||||
- urllib3==1.26.13
|
||||
- wcwidth==0.2.13
|
||||
- werkzeug==3.0.4
|
||||
- xformers==0.0.27.post2
|
||||
- yapf==0.40.2
|
||||
- zipp==3.20.2
|
||||
@@ -1,133 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding:utf-8 -*-
|
||||
# Power by Zongsheng Yue 2023-03-11 17:17:41
|
||||
|
||||
import warnings
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
import argparse
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from omegaconf import OmegaConf
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.sampler_invsr import InvSamplerSR, BaseSampler
|
||||
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.utils import util_common
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.utils.util_opts import str2bool
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.basicsr.utils.download_util import load_file_from_url
|
||||
|
||||
class Namespace:
|
||||
def __init__(self, **kwargs):
|
||||
for key, value in kwargs.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
def __repr__(self):
|
||||
items = [f"{key}={repr(value)}" for key, value in vars(self).items()]
|
||||
return f"Namespace({', '.join(items)})"
|
||||
|
||||
def get_parser(**parser_kwargs):
|
||||
parser = argparse.ArgumentParser(**parser_kwargs)
|
||||
parser.add_argument("-i", "--in_path", type=str, default="", help="Input path")
|
||||
parser.add_argument("-o", "--out_path", type=str, default="", help="Output path")
|
||||
parser.add_argument("--bs", type=int, default=1, help="Batchsize for loading image")
|
||||
parser.add_argument("--chopping_bs", type=int, default=8, help="Batchsize for chopped patch")
|
||||
parser.add_argument("-t", "--timesteps", type=int, nargs="+", help="The inversed timesteps")
|
||||
parser.add_argument("-n", "--num_steps", type=int, default=1, help="Number of inference steps")
|
||||
parser.add_argument(
|
||||
"--cfg_path", type=str, default="./configs/sample-sd-turbo.yaml", help="Configuration path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sd_path", type=str, default="", help="Path for Stable Diffusion Model",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--started_ckpt_path", type=str, default="", help="Checkpoint path for noise predictor"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tiled_vae", type=str2bool, default='true', help="Enabled tiled VAE.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--color_fix", type=str, default='', choices=['wavelet', 'ycbcr'], help="Fix the color shift",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--chopping_size", type=int, default=128, help="Chopping size when dealing large images"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
return args
|
||||
|
||||
def get_configs(args, log=False):
|
||||
configs = OmegaConf.load(args.cfg_path)
|
||||
|
||||
if args.timesteps is not None:
|
||||
assert len(args.timesteps) == args.num_steps
|
||||
configs.timesteps = sorted(args.timesteps, reverse=True)
|
||||
else:
|
||||
if args.num_steps == 1:
|
||||
configs.timesteps = [200,]
|
||||
elif args.num_steps == 2:
|
||||
configs.timesteps = [200, 100]
|
||||
elif args.num_steps == 3:
|
||||
configs.timesteps = [200, 100, 50]
|
||||
elif args.num_steps == 4:
|
||||
configs.timesteps = [200, 150, 100, 50]
|
||||
elif args.num_steps == 5:
|
||||
configs.timesteps = [250, 200, 150, 100, 50]
|
||||
else:
|
||||
assert args.num_steps <= 250
|
||||
configs.timesteps = np.linspace(
|
||||
start=args.started_step, stop=0, num=args.num_steps, endpoint=False, dtype=np.int64()
|
||||
).tolist()
|
||||
if log:
|
||||
print(f'[InvSR] - Setting timesteps for inference: {configs.timesteps}')
|
||||
|
||||
# path to save Stable Diffusion
|
||||
sd_path = args.sd_path if args.sd_path else "./weights"
|
||||
util_common.mkdir(sd_path, delete=False, parents=True)
|
||||
configs.sd_pipe.params.cache_dir = sd_path
|
||||
|
||||
# path to save noise predictor
|
||||
started_ckpt_name = "noise_predictor_sd_turbo_v5.pth"
|
||||
if args.started_ckpt_path:
|
||||
started_ckpt_dir = "custom_nodes/ComfyUI_InvSR/weights"
|
||||
started_ckpt_path = args.started_ckpt_path
|
||||
else:
|
||||
started_ckpt_dir = "./weights"
|
||||
util_common.mkdir(started_ckpt_dir, delete=False, parents=True)
|
||||
|
||||
started_ckpt_path = Path(started_ckpt_dir) / started_ckpt_name
|
||||
if not Path(started_ckpt_path).exists():
|
||||
load_file_from_url(
|
||||
url="https://huggingface.co/OAOA/InvSR/resolve/main/noise_predictor_sd_turbo_v5.pth",
|
||||
model_dir=started_ckpt_dir,
|
||||
progress=True,
|
||||
file_name=started_ckpt_name,
|
||||
)
|
||||
configs.model_start.ckpt_path = str(started_ckpt_path)
|
||||
|
||||
configs.bs = args.bs
|
||||
configs.tiled_vae = args.tiled_vae
|
||||
configs.color_fix = args.color_fix
|
||||
configs.basesr.chopping.pch_size = args.chopping_size
|
||||
if args.bs > 1:
|
||||
configs.basesr.chopping.extra_bs = 1
|
||||
else:
|
||||
configs.basesr.chopping.extra_bs = args.chopping_bs
|
||||
|
||||
return configs
|
||||
|
||||
def main():
|
||||
# args = get_parser()
|
||||
|
||||
# Load model node
|
||||
args = Namespace(in_path='./testdata/RealSet80/9.jpg', out_path='outputs', bs=1, chopping_bs=8, timesteps=None, num_steps=1, cfg_path='./configs/sample-sd-turbo.yaml', sd_path='', started_ckpt_path='', tiled_vae=True, color_fix='', chopping_size=128)
|
||||
configs = get_configs(args)
|
||||
base_sampler = BaseSampler(configs)
|
||||
|
||||
# Sampling node
|
||||
args = Namespace(in_path='/workspace/ComfyUI/custom_nodes/ComfyUI_InvSR/InvSR/testdata/RealSet80/23.jpg', out_path='/workspace', bs=1, chopping_bs=8, timesteps=None, num_steps=1, cfg_path='./configs/sample-sd-turbo.yaml', sd_path='', started_ckpt_path='/workspace/ComfyUI/custom_nodes/ComfyUI_InvSR/weights/noise_predictor_sd_turbo_v5.pth', tiled_vae=True, color_fix='', chopping_size=256)
|
||||
sampler = InvSamplerSR(base_sampler, configs)
|
||||
sampler.configs = get_configs(args)
|
||||
sampler.setup_seed(3423423)
|
||||
sampler.inference(args.in_path, out_path=args.out_path, bs=args.bs)
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -1,82 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding:utf-8 -*-
|
||||
# Power by Zongsheng Yue 2023-10-26 20:20:36
|
||||
|
||||
import warnings
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
import argparse
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from utils.util_common import get_obj_from_str
|
||||
from utils.util_opts import str2bool
|
||||
|
||||
def get_parser(**parser_kwargs):
|
||||
parser = argparse.ArgumentParser(**parser_kwargs)
|
||||
parser.add_argument(
|
||||
"--save_dir",
|
||||
type=str,
|
||||
default="./save_dir",
|
||||
help="Folder to save the checkpoints and training log",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--resume",
|
||||
type=str,
|
||||
const=True,
|
||||
default="",
|
||||
nargs="?",
|
||||
help="resume from the save_dir or checkpoint",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg_path",
|
||||
type=str,
|
||||
default="./configs/sd-turbo-sr-ldis.yaml",
|
||||
help="Configs of yaml file",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ldif",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="Loss coefficient for diffsuion in latent space",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--llpips",
|
||||
type=float,
|
||||
default=2.0,
|
||||
help="Loss coefficient for latent lpips",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ldis",
|
||||
type=float,
|
||||
default=0.1,
|
||||
help="Loss coefficient for latent discriminator",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_text",
|
||||
type=str2bool,
|
||||
default='False',
|
||||
help="Text Prompt",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
return args
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = get_parser()
|
||||
|
||||
configs = OmegaConf.load(args.cfg_path)
|
||||
if args.ldif > 0:
|
||||
configs.train.loss_coef.ldif = args.ldif
|
||||
if args.ldis > 0:
|
||||
configs.train.loss_coef.ldis = args.ldis
|
||||
if args.llpips > 0:
|
||||
configs.train.loss_coef.llpips = args.llpips
|
||||
configs.train.use_text = args.use_text
|
||||
|
||||
# merge args to config
|
||||
for key in vars(args):
|
||||
if key in ['cfg_path', 'save_dir', 'resume', ]:
|
||||
configs[key] = getattr(args, key)
|
||||
|
||||
trainer = get_obj_from_str(configs.trainer.target)(configs)
|
||||
trainer.train()
|
||||
@@ -1,91 +0,0 @@
|
||||
# Prediction interface for Cog ⚙️
|
||||
# https://cog.run/python
|
||||
|
||||
|
||||
import shutil, os
|
||||
from omegaconf import OmegaConf
|
||||
from cog import BasePredictor, Input, Path
|
||||
|
||||
import numpy as np
|
||||
from utils import util_common
|
||||
from sampler_invsr import InvSamplerSR
|
||||
from basicsr.utils.download_util import load_file_from_url
|
||||
|
||||
class Predictor(BasePredictor):
|
||||
def setup(self) -> None:
|
||||
self.configs = OmegaConf.load("./configs/sample-sd-turbo.yaml")
|
||||
|
||||
def set_configs(self, num_steps=1, chopping_size=128, seed=12345):
|
||||
if num_steps == 1:
|
||||
self.configs.timesteps = [200,]
|
||||
elif num_steps == 2:
|
||||
self.configs.timesteps = [200, 100]
|
||||
elif num_steps == 3:
|
||||
self.configs.timesteps = [200, 100, 50]
|
||||
elif num_steps == 4:
|
||||
self.configs.timesteps = [200, 150, 100, 50]
|
||||
elif num_steps == 5:
|
||||
self.configs.timesteps = [250, 200, 150, 100, 50]
|
||||
else:
|
||||
assert num_steps <= 250
|
||||
self.configs.timesteps = np.linspace(
|
||||
start=250, stop=0, num=num_steps, endpoint=False, dtype=np.int64()
|
||||
).tolist()
|
||||
print(f'Setting timesteps for inference: {self.configs.timesteps}')
|
||||
|
||||
# path to save Stable Diffusion
|
||||
sd_path = "./weights"
|
||||
util_common.mkdir(sd_path, delete=False, parents=True)
|
||||
self.configs.sd_pipe.params.cache_dir = sd_path
|
||||
|
||||
# path to save noise predictor
|
||||
started_ckpt_name = "noise_predictor_sd_turbo_v5.pth"
|
||||
started_ckpt_dir = "./weights"
|
||||
util_common.mkdir(started_ckpt_dir, delete=False, parents=True)
|
||||
started_ckpt_path = os.path.join(started_ckpt_dir, started_ckpt_name)
|
||||
if not os.path.exists(started_ckpt_path):
|
||||
load_file_from_url(
|
||||
url="https://huggingface.co/OAOA/InvSR/resolve/main/noise_predictor_sd_turbo_v5.pth",
|
||||
model_dir=started_ckpt_dir,
|
||||
progress=True,
|
||||
file_name=started_ckpt_name,
|
||||
)
|
||||
self.configs.model_start.ckpt_path = started_ckpt_path
|
||||
|
||||
self.configs.bs = 1
|
||||
self.configs.seed = 12345
|
||||
self.configs.basesr.chopping.pch_size = chopping_size
|
||||
if chopping_size == 128:
|
||||
self.configs.basesr.chopping.extra_bs = 4
|
||||
elif chopping_size == 256:
|
||||
self.configs.basesr.chopping.extra_bs = 2
|
||||
else:
|
||||
self.configs.basesr.chopping.extra_bs = 1
|
||||
|
||||
def predict(
|
||||
self,
|
||||
in_path: Path = Input(description="Input low-quality image"),
|
||||
num_steps: int = Input(
|
||||
choices=[1,2,3,4,5], description="Number of sampling steps.", default=1
|
||||
),
|
||||
chopping_size: int = Input(
|
||||
choices=[128, 256, 512], description="Chopping resolution", default=128
|
||||
),
|
||||
seed: int = Input(
|
||||
description="Random seed. Leave blank to randomize the seed.", default=12345
|
||||
),
|
||||
) -> Path:
|
||||
# setting configurations
|
||||
self.set_configs(num_steps, chopping_size, seed)
|
||||
|
||||
sampler = InvSamplerSR(self.configs)
|
||||
|
||||
out_dir = 'invsr_output'
|
||||
if os.path.exists(out_dir):
|
||||
shutil.rmtree(out_dir)
|
||||
sampler.inference(in_path, out_path=out_dir, bs=1)
|
||||
|
||||
out = "/tmp/out.png"
|
||||
shutil.copy(os.path.join(out_dir, os.listdir(out_dir)[0]), out)
|
||||
|
||||
return Path(out)
|
||||
@@ -1,12 +0,0 @@
|
||||
pyiqa==0.1.12
|
||||
opencv-python
|
||||
albumentations==1.4.18
|
||||
gradio==5.8.0
|
||||
bitsandbytes
|
||||
sentencepiece
|
||||
protobuf
|
||||
python-box
|
||||
omegaconf
|
||||
loguru
|
||||
einops
|
||||
transformers
|
||||
@@ -1,149 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding:utf-8 -*-
|
||||
# Power by Zongsheng Yue 2022-08-13 21:37:58
|
||||
|
||||
'''
|
||||
Calculate PSNR, SSIM, LPIPS, and NIQE.
|
||||
'''
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import os
|
||||
import math
|
||||
import pyiqa
|
||||
import torch
|
||||
import argparse
|
||||
from einops import rearrange
|
||||
from loguru import logger as base_logger
|
||||
|
||||
sys.path.append(str(Path(__file__).resolve().parents[1]))
|
||||
from utils import util_image
|
||||
from utils.util_opts import str2bool
|
||||
from datapipe.datasets import BaseData
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--bs", type=int, default=1, help="Batch size")
|
||||
parser.add_argument("-i", "--indir", type=str, default="", help="Path to save the testing images")
|
||||
parser.add_argument("-r", "--refdir", type=str, default="", help="Reference images for fid")
|
||||
parser.add_argument("-t", "--tocpu", type=str2bool, default='false')
|
||||
parser.add_argument("--pi", type=str2bool, default='false', help="PI metric")
|
||||
parser.add_argument("--niqe", type=str2bool, default='false', help="NIQE metric")
|
||||
parser.add_argument("--maniqa", type=str2bool, default='false', help="MANIQA metric")
|
||||
parser.add_argument("--tres", type=str2bool, default='false', help="TReS metric")
|
||||
parser.add_argument("--dbcnn", type=str2bool, default='false', help="DBCNN metric")
|
||||
args = parser.parse_args()
|
||||
|
||||
# setting logger
|
||||
log_path = str(Path(args.indir).parent / 'metrics.log')
|
||||
logger = base_logger
|
||||
logger.remove()
|
||||
logger.add(log_path, format="{time:YYYY-MM-DD(HH:mm:ss)}: {message}", mode='w', level='INFO')
|
||||
logger.add(sys.stderr, format="{message}", level='INFO')
|
||||
logger.info(f"Image Floder: {args.indir}")
|
||||
|
||||
if args.pi:
|
||||
pi_metric = pyiqa.create_metric('pi')
|
||||
if args.niqe:
|
||||
niqe_metric = pyiqa.create_metric('niqe')
|
||||
if args.maniqa:
|
||||
maniqa_metric = pyiqa.create_metric('maniqa')
|
||||
if args.tres:
|
||||
tres_metric = pyiqa.create_metric('tres')
|
||||
if args.dbcnn:
|
||||
dbcnn_metric = pyiqa.create_metric('dbcnn')
|
||||
if args.refdir:
|
||||
fid_metric = pyiqa.create_metric('fid')
|
||||
if args.tocpu:
|
||||
clipiqa_metric = pyiqa.create_metric('clipiqa').to('cpu')
|
||||
musiq_metric = pyiqa.create_metric('musiq').to('cpu')
|
||||
else:
|
||||
clipiqa_metric = pyiqa.create_metric('clipiqa')
|
||||
musiq_metric = pyiqa.create_metric('musiq')
|
||||
|
||||
dataset = BaseData(
|
||||
dir_path=args.indir,
|
||||
transform_type='default',
|
||||
transform_kwargs={'mean': 0.0, 'std': 1.0},
|
||||
need_path=True,
|
||||
im_exts=['png', 'jpeg', 'jpg', ],
|
||||
recursive=False,
|
||||
)
|
||||
dataloader = torch.utils.data.DataLoader(
|
||||
dataset,
|
||||
batch_size=args.bs,
|
||||
shuffle=False,
|
||||
drop_last=False,
|
||||
num_workers=0
|
||||
)
|
||||
logger.info(f'Number of images: {len(dataset)}')
|
||||
|
||||
metrics = {
|
||||
'PI': 0,
|
||||
'CLIPIQA': 0,
|
||||
'MUSIQ': 0,
|
||||
'MANIQA': 0,
|
||||
'TRES': 0,
|
||||
'DBCNN': 0,
|
||||
}
|
||||
if args.niqe:
|
||||
metrics['NIQE'] = 0
|
||||
for ii, data in enumerate(dataloader):
|
||||
im = data['image'].cuda() # N x h x w x 3, [0,1]
|
||||
current_bs = im.shape[0]
|
||||
|
||||
if args.pi:
|
||||
current_pi = pi_metric(im).sum().item()
|
||||
if args.niqe:
|
||||
current_niqe = niqe_metric(im).sum().item()
|
||||
if args.maniqa:
|
||||
current_maniqa = maniqa_metric(im).sum().item()
|
||||
if args.tres:
|
||||
current_tres = tres_metric(im).sum().item()
|
||||
if args.dbcnn:
|
||||
current_dbcnn = dbcnn_metric(im).sum().item()
|
||||
if args.tocpu:
|
||||
current_clipiqa = clipiqa_metric(im.cpu()).sum().item()
|
||||
current_musiq = musiq_metric(im.cpu()).sum().item()
|
||||
else:
|
||||
current_clipiqa = clipiqa_metric(im).sum().item()
|
||||
current_musiq = musiq_metric(im).sum().item()
|
||||
|
||||
if (ii+1) % 10 == 0:
|
||||
log_str = ('Processing: {:03d}/{:03d}'.format(ii+1, math.ceil(len(dataset) / args.bs)))
|
||||
logger.info(log_str)
|
||||
|
||||
metrics['CLIPIQA'] += current_clipiqa
|
||||
metrics['MUSIQ'] += current_musiq
|
||||
if args.pi:
|
||||
metrics['PI'] += current_pi
|
||||
if args.niqe:
|
||||
metrics['NIQE'] += current_niqe
|
||||
if args.maniqa:
|
||||
metrics['MANIQA'] += current_maniqa
|
||||
if args.tres:
|
||||
metrics['TRES'] += current_tres
|
||||
if args.dbcnn:
|
||||
metrics['DBCNN'] += current_dbcnn
|
||||
|
||||
for key in metrics.keys():
|
||||
metrics[key] /= len(dataset)
|
||||
|
||||
if args.refdir:
|
||||
metrics['FID'] = fid_metric(args.indir, args.refdir, mode='legacy_pytorch')
|
||||
|
||||
logger.info(f"MEAN CLIPIQA: {metrics['CLIPIQA']:6.4f}")
|
||||
logger.info(f"MEAN MUSIQ: {metrics['MUSIQ']:6.4f}")
|
||||
if args.pi:
|
||||
logger.info(f"MEAN PI: {metrics['PI']:6.4f}")
|
||||
if args.niqe:
|
||||
logger.info(f"MEAN NIQE: {metrics['NIQE']:6.4f}")
|
||||
if args.maniqa:
|
||||
logger.info(f"MEAN MANIQA: {metrics['MANIQA']:6.4f}")
|
||||
if args.tres:
|
||||
logger.info(f"MEAN TRES: {metrics['TRES']:6.4f}")
|
||||
if args.dbcnn:
|
||||
logger.info(f"MEAN DBCNN: {metrics['DBCNN']:6.4f}")
|
||||
if args.refdir:
|
||||
logger.info(f"MEAN FID: {metrics['FID']:6.4f}")
|
||||
|
||||
@@ -1,195 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding:utf-8 -*-
|
||||
# Power by Zongsheng Yue 2022-08-13 21:37:58
|
||||
|
||||
'''
|
||||
Calculate PSNR, SSIM, LPIPS, and NIQE.
|
||||
'''
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import os
|
||||
import math
|
||||
import lpips
|
||||
import pyiqa
|
||||
import torch
|
||||
import argparse
|
||||
from einops import rearrange
|
||||
from loguru import logger as base_logger
|
||||
|
||||
sys.path.append(str(Path(__file__).resolve().parents[1]))
|
||||
from utils import util_image
|
||||
from utils.util_opts import str2bool
|
||||
from datapipe.datasets import BaseData
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--bs", type=int, default=16, help="Batch size")
|
||||
parser.add_argument("--gt_dir", type=str, default="", help="Path to save the HQ images")
|
||||
parser.add_argument("--sr_dir", type=str, default="", help="Path to save the SR images")
|
||||
parser.add_argument("--log_name", type=str, default='metrics.log', help="Logging path")
|
||||
parser.add_argument("--test_y_channel", type=str2bool, default='true', help="Y channel for PSNR and SSIM")
|
||||
parser.add_argument("--fid", type=str2bool, default='false', help="Calculating FID")
|
||||
parser.add_argument("--niqe", type=str2bool, default='false', help="Calculating NIQE")
|
||||
parser.add_argument("--dists", type=str2bool, default='false', help="Calculating DISTS")
|
||||
parser.add_argument("--maniqa", type=str2bool, default='false', help="Calculating MANIQA")
|
||||
parser.add_argument("--pi", type=str2bool, default='false', help="Calculating PI")
|
||||
parser.add_argument("--tocpu", type=str2bool, default='false', help="Moving model to CPU")
|
||||
args = parser.parse_args()
|
||||
|
||||
# setting logger
|
||||
log_path = str(Path(args.sr_dir).parent / f'{args.log_name}')
|
||||
logger = base_logger
|
||||
logger.remove()
|
||||
logger.add(log_path, format="{time:YYYY-MM-DD(HH:mm:ss)}: {message}", mode='w', level='INFO')
|
||||
logger.add(sys.stderr, format="{message}", level='INFO')
|
||||
logger.info(f"Ground truth: {args.gt_dir}")
|
||||
logger.info(f"SR result: {args.sr_dir}")
|
||||
|
||||
if args.test_y_channel:
|
||||
psnr_metric = pyiqa.create_metric('psnr', test_y_channel=True, color_space='ycbcr')
|
||||
ssim_metric = pyiqa.create_metric('ssim', test_y_channel=True, color_space='ycbcr')
|
||||
else:
|
||||
psnr_metric = pyiqa.create_metric('psnr', test_y_channel=False, color_space='rgb')
|
||||
ssim_metric = pyiqa.create_metric('ssim', test_y_channel=False, color_space='rgb')
|
||||
if args.fid:
|
||||
fid_metric = pyiqa.create_metric('fid')
|
||||
if args.niqe:
|
||||
niqe_metric = pyiqa.create_metric('niqe')
|
||||
if args.dists:
|
||||
dists_metric = pyiqa.create_metric('dists')
|
||||
if args.maniqa:
|
||||
maniqa_metric = pyiqa.create_metric('maniqa')
|
||||
if args.pi:
|
||||
pi_metric = pyiqa.create_metric('pi')
|
||||
loss_fn_vgg = lpips.LPIPS(net='vgg').cuda()
|
||||
loss_fn_alex = lpips.LPIPS(net='alex').cuda()
|
||||
if args.tocpu:
|
||||
clipiqa_metric = pyiqa.create_metric('clipiqa').to('cpu')
|
||||
musiq_metric = pyiqa.create_metric('musiq').to('cpu')
|
||||
else:
|
||||
clipiqa_metric = pyiqa.create_metric('clipiqa')
|
||||
musiq_metric = pyiqa.create_metric('musiq')
|
||||
|
||||
dataset = BaseData(
|
||||
dir_path=args.sr_dir,
|
||||
transform_type='default',
|
||||
transform_kwargs={'mean': 0.0, 'std': 1.0},
|
||||
extra_dir_path=args.gt_dir,
|
||||
extra_transform_type='default',
|
||||
extra_transform_kwargs={'mean': 0.0, 'std': 1.0},
|
||||
need_path=True,
|
||||
im_exts=['png', 'jpg'],
|
||||
recursive=False,
|
||||
)
|
||||
dataloader = torch.utils.data.DataLoader(
|
||||
dataset,
|
||||
batch_size=args.bs,
|
||||
shuffle=False,
|
||||
drop_last=False,
|
||||
num_workers=0
|
||||
)
|
||||
logger.info(f'Number of images: {len(dataset)}')
|
||||
|
||||
metrics = {
|
||||
'PSNR': 0,
|
||||
'SSIM': 0,
|
||||
'LPIPS_VGG': 0,
|
||||
'LPIPS_ALEX': 0,
|
||||
'CLIPIQA': 0,
|
||||
'MUSIQ': 0,
|
||||
}
|
||||
if args.niqe:
|
||||
metrics['NIQE'] = 0
|
||||
if args.dists:
|
||||
metrics['DISTS'] = 0
|
||||
if args.maniqa:
|
||||
metrics['MANIQA'] = 0
|
||||
if args.pi:
|
||||
metrics['PI'] = 0
|
||||
for ii, data in enumerate(dataloader):
|
||||
im_sr = data['image'].cuda() # N x h x w x 3, [0,1]
|
||||
im_gt = data['gt'].cuda() # N x h x w x 3, [0,1]
|
||||
current_bs = im_sr.shape[0]
|
||||
|
||||
if not (im_sr.shape == im_gt.shape):
|
||||
height = min(im_sr.shape[-2], im_gt.shape[-2])
|
||||
width = min(im_sr.shape[-1], im_gt.shape[-1])
|
||||
im_sr = im_sr[:, :, :height, :width]
|
||||
im_gt = im_gt[:, :, :height, :width]
|
||||
|
||||
current_psnr = psnr_metric(im_sr, im_gt).mean().item()
|
||||
current_ssim = ssim_metric(im_sr, im_gt).mean().item()
|
||||
current_lpips_vgg = loss_fn_vgg(
|
||||
(im_gt - 0.5) / 0.5,
|
||||
(im_sr - 0.5) / 0.5,
|
||||
).mean().item()
|
||||
current_lpips_alex = loss_fn_alex(
|
||||
(im_gt - 0.5) / 0.5,
|
||||
(im_sr - 0.5) / 0.5,
|
||||
).mean().item()
|
||||
if args.tocpu:
|
||||
current_clipiqa = clipiqa_metric(im_sr.cpu()).mean().item()
|
||||
current_musiq = musiq_metric(im_sr.cpu()).mean().item()
|
||||
else:
|
||||
current_clipiqa = clipiqa_metric(im_sr).mean().item()
|
||||
current_musiq = musiq_metric(im_sr).mean().item()
|
||||
if args.niqe:
|
||||
current_niqe = niqe_metric(im_sr).mean().item()
|
||||
if args.dists:
|
||||
current_dists = dists_metric(im_sr, im_gt).mean().item()
|
||||
if args.maniqa:
|
||||
current_maniqa = maniqa_metric(im_sr).mean().item()
|
||||
if args.pi:
|
||||
current_pi = pi_metric(im_sr).mean().item()
|
||||
|
||||
if (ii+1) % 30 == 0:
|
||||
log_str = ('Processing: {:03d}/{:03d}, PSNR={:5.2f}, LPIPS={:6.4f}/{:6.4f}, CLIPIQA={:6.4f}, MUSIQ={:6.4f}'.format(
|
||||
ii+1,
|
||||
math.ceil(len(dataset) /args.bs),
|
||||
current_psnr,
|
||||
current_lpips_vgg,
|
||||
current_lpips_alex,
|
||||
current_clipiqa,
|
||||
current_musiq,
|
||||
))
|
||||
logger.info(log_str)
|
||||
|
||||
metrics['PSNR'] += current_psnr * current_bs
|
||||
metrics['SSIM'] += current_ssim * current_bs
|
||||
metrics['LPIPS_VGG'] += current_lpips_vgg * current_bs
|
||||
metrics['LPIPS_ALEX'] += current_lpips_alex * current_bs
|
||||
metrics['CLIPIQA'] += current_clipiqa * current_bs
|
||||
metrics['MUSIQ'] += current_musiq * current_bs
|
||||
if args.niqe:
|
||||
metrics['NIQE'] += current_niqe * current_bs
|
||||
if args.dists:
|
||||
metrics['DISTS'] += current_dists * current_bs
|
||||
if args.maniqa:
|
||||
metrics['MANIQA'] += current_maniqa * current_bs
|
||||
if args.pi:
|
||||
metrics['PI'] += current_pi * current_bs
|
||||
|
||||
for key in metrics.keys():
|
||||
metrics[key] /= len(dataset)
|
||||
|
||||
if args.fid:
|
||||
metrics['FID'] = fid_metric(args.sr_dir, args.gt_dir)
|
||||
|
||||
logger.info(f"MEAN PSNR: {metrics['PSNR']:5.2f}")
|
||||
logger.info(f"MEAN SSIM: {metrics['SSIM']:6.4f}")
|
||||
logger.info(f"MEAN LPIPS(VGG): {metrics['LPIPS_VGG']:6.4f}")
|
||||
logger.info(f"MEAN LPIPS(ALEX): {metrics['LPIPS_ALEX']:6.4f}")
|
||||
logger.info(f"MEAN CLIPIQA: {metrics['CLIPIQA']:6.4f}")
|
||||
logger.info(f"MEAN MUSIQ: {metrics['MUSIQ']:6.4f}")
|
||||
if args.fid:
|
||||
logger.info(f"MEAN FID: {metrics['FID']:6.2f}")
|
||||
if args.niqe:
|
||||
logger.info(f"MEAN NIQE: {metrics['NIQE']:7.4f}")
|
||||
if args.dists:
|
||||
logger.info(f"MEAN DISTS: {metrics['DISTS']:6.4f}")
|
||||
if args.maniqa:
|
||||
logger.info(f"MEAN MANIQA: {metrics['MANIQA']:6.4f}")
|
||||
if args.pi:
|
||||
logger.info(f"MEAN PI: {metrics['PI']:7.4f}")
|
||||
|
||||
@@ -1,107 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding:utf-8 -*-
|
||||
# Power by Zongsheng Yue 2024-04-07 20:57:36
|
||||
|
||||
import os
|
||||
import torch
|
||||
import random
|
||||
import argparse
|
||||
import numpy as np
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
sys.path.append(str(Path(__file__).parents[1]))
|
||||
|
||||
from basicsr.data.realesrgan_dataset import RealESRGANDataset
|
||||
from utils import util_image
|
||||
from utils import util_common
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"-i",
|
||||
"--indir",
|
||||
type=str,
|
||||
default="/mnt/lustre/share/zhangwenwei/data/imagenet/val",
|
||||
help="Folder to save the checkpoints and training log",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-o",
|
||||
"--outdir",
|
||||
type=str,
|
||||
default="./ImageNet-Test",
|
||||
help="Folder to save the checkpoints and training log",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-r",
|
||||
"--resolution",
|
||||
type=int,
|
||||
default=1024,
|
||||
help="Image resolution of the ground truth",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num_imgs",
|
||||
type=int,
|
||||
default=-1,
|
||||
help="Number of images.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if Path(args.indir).is_dir():
|
||||
img_list = sorted([x for x in Path(args.indir).glob('*.[JjPp][PpNn]*[Gg]')])
|
||||
elif args.indir.endswith('txt'):
|
||||
img_list = util_common.readline_txt(args.indir)
|
||||
else:
|
||||
raise ValueError('Please input valid args.indir!')
|
||||
print(f'Number of images in the input folder: {len(img_list)}')
|
||||
|
||||
random.seed(10000)
|
||||
random.shuffle(img_list)
|
||||
|
||||
num_imgs = args.num_imgs
|
||||
if num_imgs > 0:
|
||||
assert num_imgs <= len(img_list)
|
||||
img_list = random.sample(img_list, k=num_imgs)
|
||||
|
||||
gt_dir = Path(args.outdir) / 'gt'
|
||||
if not gt_dir.exists():
|
||||
gt_dir.mkdir(parents=True)
|
||||
lq_dir = Path(args.outdir) / 'lq'
|
||||
if not lq_dir.exists():
|
||||
lq_dir.mkdir(parents=True)
|
||||
|
||||
# Loading configuration
|
||||
configs = OmegaConf.load('./configs/degradation_testing_realesrgan.yaml')
|
||||
opts, opts_degradation = configs.opts, configs.degradation
|
||||
opts['gt_size'] = args.resolution
|
||||
opts_degradation['gt_size'] = args.resolution
|
||||
|
||||
dataset = RealESRGANDataset(opts, mode='testing')
|
||||
dataset.image_paths = img_list
|
||||
dataset.text_paths = [None, ] * len(img_list)
|
||||
dataset.moment_paths = [None, ] * len(img_list)
|
||||
for ii in range(len(img_list)):
|
||||
data_dict1 = dataset.__getitem__(ii)
|
||||
if (ii + 1) % 100 == 0:
|
||||
print(f'Processing: {ii+1}/{len(img_list)}')
|
||||
prefix = 'realesrgan'
|
||||
data_dict2 = dataset.degrade_fun(
|
||||
opts_degradation,
|
||||
im_gt=data_dict1['gt'].unsqueeze(0),
|
||||
kernel1=data_dict1['kernel1'],
|
||||
kernel2=data_dict1['kernel2'],
|
||||
sinc_kernel=data_dict1['sinc_kernel'],
|
||||
)
|
||||
im_lq, im_gt = data_dict2['lq'], data_dict2['gt']
|
||||
im_lq, im_gt = util_image.tensor2img([im_lq, im_gt], rgb2bgr=True, min_max=(0,1) ) # uint8
|
||||
|
||||
im_name = Path(data_dict1['gt_path']).stem
|
||||
im_path_gt = gt_dir / f'{im_name}.png'
|
||||
util_image.imwrite(im_gt, im_path_gt, chn='bgr', dtype_in='uint8')
|
||||
|
||||
im_path_lq = lq_dir / f'{im_name}.png'
|
||||
util_image.imwrite(im_lq, im_path_lq, chn='bgr', dtype_in='uint8')
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
-1643
File diff suppressed because it is too large
Load Diff
@@ -1,98 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding:utf-8 -*-
|
||||
# Power by Zongsheng Yue 2021-11-24 20:29:36
|
||||
|
||||
import math
|
||||
import torch
|
||||
from pathlib import Path
|
||||
from copy import deepcopy
|
||||
from collections import OrderedDict
|
||||
import torch.nn.functional as F
|
||||
|
||||
def calculate_parameters(net):
|
||||
out = 0
|
||||
for param in net.parameters():
|
||||
out += param.numel()
|
||||
return out
|
||||
|
||||
def pad_input(x, mod):
|
||||
h, w = x.shape[-2:]
|
||||
bottom = int(math.ceil(h/mod)*mod -h)
|
||||
right = int(math.ceil(w/mod)*mod - w)
|
||||
x_pad = F.pad(x, pad=(0, right, 0, bottom), mode='reflect')
|
||||
return x_pad
|
||||
|
||||
def forward_chop(net, x, net_kwargs=None, scale=1, shave=10, min_size=160000):
|
||||
n_GPUs = 1
|
||||
b, c, h, w = x.size()
|
||||
h_half, w_half = h // 2, w // 2
|
||||
h_size, w_size = h_half + shave, w_half + shave
|
||||
lr_list = [
|
||||
x[:, :, 0:h_size, 0:w_size],
|
||||
x[:, :, 0:h_size, (w - w_size):w],
|
||||
x[:, :, (h - h_size):h, 0:w_size],
|
||||
x[:, :, (h - h_size):h, (w - w_size):w]]
|
||||
|
||||
if w_size * h_size < min_size:
|
||||
sr_list = []
|
||||
for i in range(0, 4, n_GPUs):
|
||||
lr_batch = torch.cat(lr_list[i:(i + n_GPUs)], dim=0)
|
||||
if net_kwargs is None:
|
||||
sr_batch = net(lr_batch)
|
||||
else:
|
||||
sr_batch = net(lr_batch, **net_kwargs)
|
||||
sr_list.extend(sr_batch.chunk(n_GPUs, dim=0))
|
||||
else:
|
||||
sr_list = [
|
||||
forward_chop(patch, shave=shave, min_size=min_size) \
|
||||
for patch in lr_list
|
||||
]
|
||||
|
||||
h, w = scale * h, scale * w
|
||||
h_half, w_half = scale * h_half, scale * w_half
|
||||
h_size, w_size = scale * h_size, scale * w_size
|
||||
shave *= scale
|
||||
|
||||
output = x.new(b, c, h, w)
|
||||
output[:, :, 0:h_half, 0:w_half] \
|
||||
= sr_list[0][:, :, 0:h_half, 0:w_half]
|
||||
output[:, :, 0:h_half, w_half:w] \
|
||||
= sr_list[1][:, :, 0:h_half, (w_size - w + w_half):w_size]
|
||||
output[:, :, h_half:h, 0:w_half] \
|
||||
= sr_list[2][:, :, (h_size - h + h_half):h_size, 0:w_half]
|
||||
output[:, :, h_half:h, w_half:w] \
|
||||
= sr_list[3][:, :, (h_size - h + h_half):h_size, (w_size - w + w_half):w_size]
|
||||
|
||||
return output
|
||||
|
||||
def measure_time(net, inputs, num_forward=100):
|
||||
'''
|
||||
Measuring the average runing time (seconds) for pytorch.
|
||||
out = net(*inputs)
|
||||
'''
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
|
||||
start.record()
|
||||
with torch.set_grad_enabled(False):
|
||||
for _ in range(num_forward):
|
||||
out = net(*inputs)
|
||||
end.record()
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
return start.elapsed_time(end) / 1000
|
||||
|
||||
def reload_model(model, ckpt):
|
||||
module_flag = list(ckpt.keys())[0].startswith('module.')
|
||||
compile_flag = '_orig_mod' in list(ckpt.keys())[0]
|
||||
|
||||
for source_key, source_value in model.state_dict().items():
|
||||
target_key = source_key
|
||||
if compile_flag and (not '_orig_mod.' in source_key):
|
||||
target_key = '_orig_mod.' + target_key
|
||||
if module_flag and (not source_key.startswith('module')):
|
||||
target_key = 'module.' + target_key
|
||||
|
||||
assert target_key in ckpt
|
||||
source_value.copy_(ckpt[target_key])
|
||||
@@ -1,12 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding:utf-8 -*-
|
||||
# Power by Zongsheng Yue 2024-08-15 16:25:07
|
||||
|
||||
def append_dims(x, target_dims:int):
|
||||
"""Appends dimensions to the end of a tensor until it has target_dims dimensions."""
|
||||
dims_to_append = target_dims - x.ndim
|
||||
if dims_to_append < 0:
|
||||
raise ValueError(
|
||||
f"input has {x.ndim} dims but target_dims is {target_dims}, which is less"
|
||||
)
|
||||
return x[(...,) + (None,) * dims_to_append]
|
||||
+1
-1
@@ -10,4 +10,4 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"InvSRSampler" : "InvSRSampler"
|
||||
}
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
from .inference_invsr import get_configs
|
||||
from .sampler_invsr import InvSamplerSR, BaseSampler
|
||||
from .noise_predictor import NoisePredictor
|
||||
from .time_aware_encoder import TimeAwareEncoder
|
||||
|
||||
__all__ = [
|
||||
"get_configs",
|
||||
"InvSamplerSR",
|
||||
"BaseSampler",
|
||||
"NoisePredictor",
|
||||
"TimeAwareEncoder"
|
||||
]
|
||||
@@ -0,0 +1,85 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding:utf-8 -*-
|
||||
# Power by Zongsheng Yue 2023-03-11 17:17:41
|
||||
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from omegaconf import OmegaConf
|
||||
from .sampler_invsr import InvSamplerSR, BaseSampler
|
||||
|
||||
from .utils import util_common
|
||||
from .utils.util_opts import str2bool
|
||||
from huggingface_hub import hf_hub_download
|
||||
from shutil import copy2
|
||||
|
||||
class Namespace:
|
||||
def __init__(self, **kwargs):
|
||||
for key, value in kwargs.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
def __repr__(self):
|
||||
items = [f"{key}={repr(value)}" for key, value in vars(self).items()]
|
||||
return f"Namespace({', '.join(items)})"
|
||||
|
||||
def get_configs(args, log=False):
|
||||
configs = OmegaConf.load(args.cfg_path)
|
||||
|
||||
if args.timesteps is not None:
|
||||
assert len(args.timesteps) == args.num_steps
|
||||
configs.timesteps = sorted(args.timesteps, reverse=True)
|
||||
else:
|
||||
if args.num_steps == 1:
|
||||
configs.timesteps = [200,]
|
||||
elif args.num_steps == 2:
|
||||
configs.timesteps = [200, 100]
|
||||
elif args.num_steps == 3:
|
||||
configs.timesteps = [200, 100, 50]
|
||||
elif args.num_steps == 4:
|
||||
configs.timesteps = [200, 150, 100, 50]
|
||||
elif args.num_steps == 5:
|
||||
configs.timesteps = [250, 200, 150, 100, 50]
|
||||
else:
|
||||
assert args.num_steps <= 250
|
||||
configs.timesteps = np.linspace(
|
||||
start=args.started_step, stop=0, num=args.num_steps, endpoint=False, dtype=np.int64()
|
||||
).tolist()
|
||||
if log:
|
||||
print(f'[InvSR] - Setting timesteps for inference: {configs.timesteps}')
|
||||
|
||||
# path to save Stable Diffusion
|
||||
sd_path = args.sd_path if args.sd_path else "./weights"
|
||||
util_common.mkdir(sd_path, delete=False, parents=True)
|
||||
configs.sd_pipe.params.cache_dir = sd_path
|
||||
|
||||
# path to save noise predictor
|
||||
started_ckpt_name = "noise_predictor_sd_turbo_v5.pth"
|
||||
|
||||
if getattr(args, "started_ckpt_dir", None) is not None:
|
||||
started_ckpt_dir = args.started_ckpt_dir
|
||||
else:
|
||||
started_ckpt_dir = "./weights"
|
||||
|
||||
if getattr(args, "started_ckpt_path", None) is not None:
|
||||
started_ckpt_path = args.started_ckpt_path
|
||||
else:
|
||||
started_ckpt_path = Path(started_ckpt_dir) / started_ckpt_name
|
||||
util_common.mkdir(started_ckpt_dir, delete=False, parents=True)
|
||||
|
||||
if not Path(started_ckpt_path).exists():
|
||||
temp_path = hf_hub_download(
|
||||
repo_id="OAOA/InvSR",
|
||||
filename="noise_predictor_sd_turbo_v5.pth",
|
||||
)
|
||||
copy2(temp_path, started_ckpt_path)
|
||||
configs.model_start.ckpt_path = str(started_ckpt_path)
|
||||
|
||||
configs.bs = args.bs
|
||||
configs.tiled_vae = args.tiled_vae
|
||||
configs.color_fix = args.color_fix
|
||||
configs.basesr.chopping.pch_size = args.chopping_size
|
||||
if args.bs > 1:
|
||||
configs.basesr.chopping.extra_bs = 1
|
||||
else:
|
||||
configs.basesr.chopping.extra_bs = args.chopping_bs
|
||||
|
||||
return configs
|
||||
Executable → Regular
Executable → Regular
@@ -0,0 +1,366 @@
|
||||
from typing import Dict, Optional, Tuple, Union
|
||||
import torch
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.loaders.single_file_model import FromOriginalModelMixin
|
||||
from diffusers.models.autoencoders.vae import (
|
||||
Decoder,
|
||||
DecoderOutput,
|
||||
DiagonalGaussianDistribution,
|
||||
Encoder,
|
||||
)
|
||||
from diffusers.models.attention_processor import (
|
||||
ADDED_KV_ATTENTION_PROCESSORS,
|
||||
CROSS_ATTENTION_PROCESSORS,
|
||||
AttentionProcessor,
|
||||
AttnAddedKVProcessor,
|
||||
AttnProcessor,
|
||||
)
|
||||
from diffusers.models.modeling_outputs import AutoencoderKLOutput
|
||||
from diffusers.utils.accelerate_utils import apply_forward_hook
|
||||
from .time_aware_encoder import TimeAwareEncoder
|
||||
|
||||
class NoisePredictor(ModelMixin, ConfigMixin, FromOriginalModelMixin):
|
||||
r"""
|
||||
A noise predicted model from the encoder of AutoencoderKL.
|
||||
|
||||
This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
|
||||
for all models (such as downloading or saving).
|
||||
|
||||
Parameters:
|
||||
in_channels (int, *optional*, defaults to 3): Number of channels in the input image.
|
||||
down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
|
||||
Tuple of downsample block types.
|
||||
up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`):
|
||||
Tuple of upsample block types.
|
||||
block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`):
|
||||
Tuple of block output channels.
|
||||
act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
|
||||
latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space.
|
||||
sample_size (`int`, *optional*, defaults to `32`): Sample input size.
|
||||
mid_block_add_attention (`bool`, *optional*, default to `True`):
|
||||
If enabled, the mid_block of the Encoder and Decoder will have attention blocks. If set to false, the
|
||||
mid_block will only have resnet blocks
|
||||
temb_channels (`int`, *optional*, default to 256): Number of channels for time embedding
|
||||
freq_shift (`int`, *optional*, defaults to 0): Frequency shift for Fourier time embedding.
|
||||
flip_sin_to_cos (`bool`, *optional*, defaults to `True`):
|
||||
Whether to flip sin to cos for Fourier time embedding.
|
||||
double_z (`bool`, *optional*, defaults to `True`):
|
||||
Whether to double the number of output channels for the last block.
|
||||
"""
|
||||
|
||||
_supports_gradient_checkpointing = True
|
||||
_no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D"]
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int = 3,
|
||||
down_block_types: Tuple[str] = ("DownEncoderBlock2D",),
|
||||
up_block_types: Tuple[str] = ("UpDecoderBlock2D",),
|
||||
block_out_channels: Tuple[int] = (64,),
|
||||
layers_per_block: int = 1,
|
||||
act_fn: str = "silu",
|
||||
latent_channels: int = 4,
|
||||
norm_num_groups: int = 32,
|
||||
sample_size: int = 32,
|
||||
mid_block_add_attention: bool = True,
|
||||
attention_head_dim: int = 1,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
temb_channels: int = 256,
|
||||
freq_shift: int = 0,
|
||||
flip_sin_to_cos: bool = True,
|
||||
double_z: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# pass init params to Encoder
|
||||
self.encoder = TimeAwareEncoder(
|
||||
in_channels=in_channels,
|
||||
out_channels=latent_channels,
|
||||
down_block_types=down_block_types,
|
||||
block_out_channels=block_out_channels,
|
||||
layers_per_block=layers_per_block,
|
||||
act_fn=act_fn,
|
||||
norm_num_groups=norm_num_groups,
|
||||
double_z=double_z,
|
||||
mid_block_add_attention=mid_block_add_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
temb_channels=temb_channels,
|
||||
freq_shift=freq_shift,
|
||||
flip_sin_to_cos=flip_sin_to_cos,
|
||||
attention_head_dim=attention_head_dim,
|
||||
)
|
||||
|
||||
self.use_slicing = False
|
||||
self.use_tiling = False
|
||||
self.double_z = double_z
|
||||
|
||||
# only relevant if vae tiling is enabled
|
||||
self.tile_sample_min_size = self.config.sample_size
|
||||
sample_size = (
|
||||
self.config.sample_size[0]
|
||||
if isinstance(self.config.sample_size, (list, tuple))
|
||||
else self.config.sample_size
|
||||
)
|
||||
self.tile_latent_min_size = int(
|
||||
sample_size / (2 ** (len(self.config.block_out_channels) - 1))
|
||||
)
|
||||
self.tile_overlap_factor = 0.25
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
if isinstance(module, (Encoder, Decoder)):
|
||||
module.gradient_checkpointing = value
|
||||
|
||||
def enable_tiling(self, use_tiling: bool = True):
|
||||
r"""
|
||||
Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
|
||||
compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
|
||||
processing larger images.
|
||||
"""
|
||||
self.use_tiling = use_tiling
|
||||
|
||||
def disable_tiling(self):
|
||||
r"""
|
||||
Disable tiled VAE decoding. If `enable_tiling` was previously enabled, this method will go back to computing
|
||||
decoding in one step.
|
||||
"""
|
||||
self.enable_tiling(False)
|
||||
|
||||
def enable_slicing(self):
|
||||
r"""
|
||||
Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
|
||||
compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
|
||||
"""
|
||||
self.use_slicing = True
|
||||
|
||||
def disable_slicing(self):
|
||||
r"""
|
||||
Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing
|
||||
decoding in one step.
|
||||
"""
|
||||
self.use_slicing = False
|
||||
|
||||
@property
|
||||
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
|
||||
def attn_processors(self) -> Dict[str, AttentionProcessor]:
|
||||
r"""
|
||||
Returns:
|
||||
`dict` of attention processors: A dictionary containing all attention processors used in the model with
|
||||
indexed by its weight name.
|
||||
"""
|
||||
# set recursively
|
||||
processors = {}
|
||||
|
||||
def fn_recursive_add_processors(
|
||||
name: str,
|
||||
module: torch.nn.Module,
|
||||
processors: Dict[str, AttentionProcessor],
|
||||
):
|
||||
if hasattr(module, "get_processor"):
|
||||
processors[f"{name}.processor"] = module.get_processor()
|
||||
|
||||
for sub_name, child in module.named_children():
|
||||
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
|
||||
|
||||
return processors
|
||||
|
||||
for name, module in self.named_children():
|
||||
fn_recursive_add_processors(name, module, processors)
|
||||
|
||||
return processors
|
||||
|
||||
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
|
||||
def set_attn_processor(
|
||||
self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]
|
||||
):
|
||||
r"""
|
||||
Sets the attention processor to use to compute attention.
|
||||
|
||||
Parameters:
|
||||
processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
|
||||
The instantiated processor class or a dictionary of processor classes that will be set as the processor
|
||||
for **all** `Attention` layers.
|
||||
|
||||
If `processor` is a dict, the key needs to define the path to the corresponding cross attention
|
||||
processor. This is strongly recommended when setting trainable attention processors.
|
||||
|
||||
"""
|
||||
count = len(self.attn_processors.keys())
|
||||
|
||||
if isinstance(processor, dict) and len(processor) != count:
|
||||
raise ValueError(
|
||||
f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
|
||||
f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
|
||||
)
|
||||
|
||||
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
|
||||
if hasattr(module, "set_processor"):
|
||||
if not isinstance(processor, dict):
|
||||
module.set_processor(processor)
|
||||
else:
|
||||
module.set_processor(processor.pop(f"{name}.processor"))
|
||||
|
||||
for sub_name, child in module.named_children():
|
||||
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
|
||||
|
||||
for name, module in self.named_children():
|
||||
fn_recursive_attn_processor(name, module, processor)
|
||||
|
||||
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
|
||||
def set_default_attn_processor(self):
|
||||
"""
|
||||
Disables custom attention processors and sets the default attention implementation.
|
||||
"""
|
||||
if all(
|
||||
proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS
|
||||
for proc in self.attn_processors.values()
|
||||
):
|
||||
processor = AttnAddedKVProcessor()
|
||||
elif all(
|
||||
proc.__class__ in CROSS_ATTENTION_PROCESSORS
|
||||
for proc in self.attn_processors.values()
|
||||
):
|
||||
processor = AttnProcessor()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
|
||||
)
|
||||
|
||||
self.set_attn_processor(processor)
|
||||
|
||||
@apply_forward_hook
|
||||
def encode(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
timestep: Union[int, torch.Tensor],
|
||||
return_dict: bool = True,
|
||||
) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
|
||||
"""
|
||||
Encode a batch of images into latents.
|
||||
|
||||
Args:
|
||||
x (`torch.Tensor`): Input batch of images.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
|
||||
|
||||
Returns:
|
||||
The latent representations of the encoded images. If `return_dict` is True, a
|
||||
[`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
|
||||
"""
|
||||
if self.use_tiling and (
|
||||
x.shape[-1] > self.tile_sample_min_size
|
||||
or x.shape[-2] > self.tile_sample_min_size
|
||||
):
|
||||
return self.tiled_encode(x, timestep, return_dict=return_dict)
|
||||
|
||||
if self.use_slicing and x.shape[0] > 1:
|
||||
encoded_slices = [self.encoder(x_slice, timestep) for x_slice in x.split(1)]
|
||||
h = torch.cat(encoded_slices)
|
||||
else:
|
||||
h = self.encoder(x, timestep)
|
||||
|
||||
if not self.double_z:
|
||||
return h
|
||||
|
||||
posterior = DiagonalGaussianDistribution(h)
|
||||
|
||||
if not return_dict:
|
||||
return (posterior,)
|
||||
|
||||
return AutoencoderKLOutput(latent_dist=posterior)
|
||||
|
||||
def tiled_encode(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
timestep: Union[int, torch.Tensor],
|
||||
return_dict: bool = True,
|
||||
) -> AutoencoderKLOutput:
|
||||
r"""Encode a batch of images using a tiled encoder.
|
||||
|
||||
When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
|
||||
steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
|
||||
different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
|
||||
tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
|
||||
output, but they should be much less noticeable.
|
||||
|
||||
Args:
|
||||
x (`torch.Tensor`): Input batch of images.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
|
||||
|
||||
Returns:
|
||||
[`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`:
|
||||
If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain
|
||||
`tuple` is returned.
|
||||
"""
|
||||
overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
|
||||
blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
|
||||
row_limit = self.tile_latent_min_size - blend_extent
|
||||
|
||||
# Split the image into 512x512 tiles and encode them separately.
|
||||
rows = []
|
||||
for i in range(0, x.shape[2], overlap_size):
|
||||
row = []
|
||||
for j in range(0, x.shape[3], overlap_size):
|
||||
tile = x[
|
||||
:,
|
||||
:,
|
||||
i : i + self.tile_sample_min_size,
|
||||
j : j + self.tile_sample_min_size,
|
||||
]
|
||||
tile = self.encoder(tile, timestep)
|
||||
if self.config.use_quant_conv:
|
||||
tile = self.quant_conv(tile)
|
||||
row.append(tile)
|
||||
rows.append(row)
|
||||
result_rows = []
|
||||
for i, row in enumerate(rows):
|
||||
result_row = []
|
||||
for j, tile in enumerate(row):
|
||||
# blend the above tile and the left tile
|
||||
# to the current tile and add the current tile to the result row
|
||||
if i > 0:
|
||||
tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
|
||||
if j > 0:
|
||||
tile = self.blend_h(row[j - 1], tile, blend_extent)
|
||||
result_row.append(tile[:, :, :row_limit, :row_limit])
|
||||
result_rows.append(torch.cat(result_row, dim=3))
|
||||
|
||||
moments = torch.cat(result_rows, dim=2)
|
||||
posterior = DiagonalGaussianDistribution(moments)
|
||||
|
||||
if not return_dict:
|
||||
return (posterior,)
|
||||
|
||||
return AutoencoderKLOutput(latent_dist=posterior)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
sample: torch.Tensor,
|
||||
timesteps: torch.Tensor,
|
||||
sample_posterior: bool = True,
|
||||
center_input_sample: bool = True,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
) -> Union[DecoderOutput, torch.Tensor]:
|
||||
r"""
|
||||
Args:
|
||||
sample (`torch.Tensor`): Input sample.
|
||||
sample_posterior (`bool`, *optional*, defaults to `False`):
|
||||
Whether to sample from the posterior.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
|
||||
"""
|
||||
if center_input_sample:
|
||||
sample = sample * 2 - 1.0
|
||||
|
||||
if not self.double_z:
|
||||
h = self.encode(sample, timesteps)
|
||||
return h
|
||||
else:
|
||||
posterior = self.encode(sample, timesteps).latent_dist
|
||||
|
||||
if sample_posterior:
|
||||
return posterior.sample()
|
||||
else:
|
||||
return posterior
|
||||
+14
-12
@@ -21,14 +21,14 @@ import torch
|
||||
from packaging import version
|
||||
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
|
||||
|
||||
from ...callbacks import MultiPipelineCallbacks, PipelineCallback
|
||||
from ...configuration_utils import FrozenDict
|
||||
from ...image_processor import PipelineImageInput, VaeImageProcessor
|
||||
from ...loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin
|
||||
from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel
|
||||
from ...models.lora import adjust_lora_scale_text_encoder
|
||||
from ...schedulers import KarrasDiffusionSchedulers
|
||||
from ...utils import (
|
||||
from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
|
||||
from diffusers.configuration_utils import FrozenDict
|
||||
from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
|
||||
from diffusers.loaders import FromSingleFileMixin, IPAdapterMixin, StableDiffusionLoraLoaderMixin, TextualInversionLoaderMixin
|
||||
from diffusers.models import AutoencoderKL, ImageProjection, UNet2DConditionModel
|
||||
from diffusers.models.lora import adjust_lora_scale_text_encoder
|
||||
from diffusers.schedulers import KarrasDiffusionSchedulers
|
||||
from diffusers.utils import (
|
||||
PIL_INTERPOLATION,
|
||||
USE_PEFT_BACKEND,
|
||||
deprecate,
|
||||
@@ -37,10 +37,12 @@ from ...utils import (
|
||||
scale_lora_layers,
|
||||
unscale_lora_layers,
|
||||
)
|
||||
from ...utils.torch_utils import randn_tensor
|
||||
from ..pipeline_utils import DiffusionPipeline, StableDiffusionMixin
|
||||
from . import StableDiffusionPipelineOutput
|
||||
from .safety_checker import StableDiffusionSafetyChecker
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
|
||||
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
|
||||
from diffusers.pipelines.stable_diffusion.safety_checker import (
|
||||
StableDiffusionSafetyChecker,
|
||||
)
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
@@ -4,25 +4,21 @@
|
||||
|
||||
import os, sys, math, random
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from loguru import logger
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.utils import util_net
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.utils import util_image
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.utils import util_common
|
||||
from custom_nodes.ComfyUI_InvSR.InvSR.utils import util_color_fix
|
||||
from comfy.utils import ProgressBar
|
||||
from .utils import util_net
|
||||
from .utils import util_image
|
||||
from .utils import util_common
|
||||
from .utils import util_color_fix
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
import torch.multiprocessing as mean_psnr
|
||||
|
||||
from .datapipe.datasets import create_dataset
|
||||
from diffusers import StableDiffusionInvEnhancePipeline, AutoencoderKL
|
||||
from .pipeline_stable_diffusion_inversion_sr import StableDiffusionInvEnhancePipeline
|
||||
from diffusers import AutoencoderKL
|
||||
|
||||
_positive= 'Cinematic, high-contrast, photo-realistic, 8k, ultra HD, ' +\
|
||||
'meticulous detailing, hyper sharpness, perfect without deformations'
|
||||
@@ -240,7 +236,3 @@ def get_torch_dtype(torch_dtype: str):
|
||||
return torch.float32
|
||||
else:
|
||||
raise ValueError(f'Unexpected torch dtype:{torch_dtype}')
|
||||
|
||||
if __name__ == '__main__':
|
||||
pass
|
||||
|
||||
@@ -0,0 +1,202 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Tuple, Union
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from diffusers.utils import is_torch_version
|
||||
from diffusers.models.unets.unet_2d_blocks import (
|
||||
UNetMidBlock2D,
|
||||
get_down_block,
|
||||
)
|
||||
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
|
||||
|
||||
|
||||
class TimeAwareEncoder(nn.Module):
|
||||
r"""
|
||||
The `TimeAwareEncoder` layer of a variational autoencoder that encodes its input into a latent representation.
|
||||
|
||||
Args:
|
||||
in_channels (`int`, *optional*, defaults to 3):
|
||||
The number of input channels.
|
||||
out_channels (`int`, *optional*, defaults to 3):
|
||||
The number of output channels.
|
||||
down_block_types (`Tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`):
|
||||
The types of down blocks to use. See `~diffusers.models.unet_2d_blocks.get_down_block` for available
|
||||
options.
|
||||
block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`):
|
||||
The number of output channels for each block.
|
||||
layers_per_block (`int`, *optional*, defaults to 2):
|
||||
The number of layers per block.
|
||||
norm_num_groups (`int`, *optional*, defaults to 32):
|
||||
The number of groups for normalization.
|
||||
act_fn (`str`, *optional*, defaults to `"silu"`):
|
||||
The activation function to use. See `~diffusers.models.activations.get_activation` for available options.
|
||||
double_z (`bool`, *optional*, defaults to `True`):
|
||||
Whether to double the number of output channels for the last block.
|
||||
resnet_time_scale_shift (`str`, defaults to `"default"`)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int = 3,
|
||||
out_channels: int = 3,
|
||||
down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",),
|
||||
block_out_channels: Tuple[int, ...] = (64,),
|
||||
layers_per_block: Union[int, Tuple[int, ...]] = 2,
|
||||
norm_num_groups: int = 32,
|
||||
act_fn: str = "silu",
|
||||
double_z: bool = True,
|
||||
mid_block_add_attention=True,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
temb_channels: int = 256,
|
||||
freq_shift: int = 0,
|
||||
flip_sin_to_cos: bool = True,
|
||||
attention_head_dim: int = 1,
|
||||
):
|
||||
super().__init__()
|
||||
if isinstance(layers_per_block, int):
|
||||
layers_per_block = (layers_per_block,) * len(down_block_types)
|
||||
self.layers_per_block = layers_per_block
|
||||
|
||||
timestep_input_dim = max(128, block_out_channels[0])
|
||||
self.time_proj = Timesteps(timestep_input_dim, flip_sin_to_cos, freq_shift)
|
||||
self.time_embedding = TimestepEmbedding(timestep_input_dim, temb_channels)
|
||||
|
||||
self.conv_in = nn.Conv2d(
|
||||
in_channels,
|
||||
block_out_channels[0],
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1,
|
||||
)
|
||||
|
||||
self.down_blocks = nn.ModuleList([])
|
||||
|
||||
# down
|
||||
output_channel = block_out_channels[0]
|
||||
for i, down_block_type in enumerate(down_block_types):
|
||||
input_channel = output_channel
|
||||
output_channel = block_out_channels[i]
|
||||
is_final_block = i == len(block_out_channels) - 1
|
||||
|
||||
down_block = get_down_block(
|
||||
down_block_type,
|
||||
num_layers=self.layers_per_block[i],
|
||||
in_channels=input_channel,
|
||||
out_channels=output_channel,
|
||||
add_downsample=not is_final_block,
|
||||
resnet_eps=1e-6,
|
||||
downsample_padding=0,
|
||||
resnet_act_fn=act_fn,
|
||||
resnet_groups=norm_num_groups,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
temb_channels=temb_channels,
|
||||
)
|
||||
self.down_blocks.append(down_block)
|
||||
|
||||
# mid
|
||||
self.mid_block = UNetMidBlock2D(
|
||||
in_channels=block_out_channels[-1],
|
||||
resnet_eps=1e-6,
|
||||
resnet_act_fn=act_fn,
|
||||
output_scale_factor=1,
|
||||
attention_head_dim=attention_head_dim,
|
||||
resnet_groups=norm_num_groups,
|
||||
add_attention=mid_block_add_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
temb_channels=temb_channels,
|
||||
)
|
||||
|
||||
# out
|
||||
self.conv_norm_out = nn.GroupNorm(
|
||||
num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6
|
||||
)
|
||||
self.conv_act = nn.SiLU()
|
||||
|
||||
conv_out_channels = 2 * out_channels if double_z else out_channels
|
||||
self.conv_out = nn.Conv2d(
|
||||
block_out_channels[-1], conv_out_channels, 3, padding=1
|
||||
)
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
sample: torch.Tensor,
|
||||
timesteps: Union[torch.Tensor, int],
|
||||
) -> torch.Tensor:
|
||||
r"""The forward method of the `Encoder` class."""
|
||||
|
||||
# time embedding
|
||||
if not torch.is_tensor(timesteps):
|
||||
timesteps = torch.tensor(
|
||||
[timesteps], dtype=torch.long, device=sample.device
|
||||
)
|
||||
elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0:
|
||||
timesteps = timesteps[None].to(sample.device)
|
||||
|
||||
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
||||
timesteps = timesteps * torch.ones(
|
||||
sample.shape[0], dtype=timesteps.dtype, device=timesteps.device
|
||||
)
|
||||
|
||||
t_emb = self.time_proj(timesteps)
|
||||
|
||||
# timesteps does not contain any weights and will always return f32 tensors
|
||||
# but time_embedding might actually be running in fp16. so we need to cast here.
|
||||
# there might be better ways to encapsulate this.
|
||||
t_emb = t_emb.to(dtype=list(self.time_embedding.parameters())[0].dtype)
|
||||
emb = self.time_embedding(t_emb)
|
||||
|
||||
sample = self.conv_in(sample)
|
||||
|
||||
if self.training and self.gradient_checkpointing:
|
||||
|
||||
def create_custom_forward(module):
|
||||
def custom_forward(*inputs):
|
||||
return module(*inputs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
# down
|
||||
if is_torch_version(">=", "1.11.0"):
|
||||
for down_block in self.down_blocks:
|
||||
sample = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(down_block),
|
||||
sample,
|
||||
emb,
|
||||
use_reentrant=False,
|
||||
)
|
||||
# middle
|
||||
sample = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(self.mid_block),
|
||||
sample,
|
||||
emb,
|
||||
use_reentrant=False,
|
||||
)
|
||||
else:
|
||||
for down_block in self.down_blocks:
|
||||
sample = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(down_block), sample, emb
|
||||
)
|
||||
# middle
|
||||
sample = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(self.mid_block), sample, emb
|
||||
)
|
||||
|
||||
else:
|
||||
# down
|
||||
for down_block in self.down_blocks:
|
||||
sample, _ = down_block(sample, emb)
|
||||
|
||||
# middle
|
||||
sample = self.mid_block(sample, emb)
|
||||
|
||||
# post-process
|
||||
sample = self.conv_norm_out(sample)
|
||||
sample = self.conv_act(sample)
|
||||
sample = self.conv_out(sample)
|
||||
|
||||
return sample
|
||||
@@ -1,136 +1,106 @@
|
||||
'''
|
||||
# --------------------------------------------------------------------------------
|
||||
# Color fixed script from Li Yi (https://github.com/pkuliyi2015/sd-webui-stablesr/blob/master/srmodule/colorfix.py)
|
||||
# --------------------------------------------------------------------------------
|
||||
'''
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torch import Tensor
|
||||
from torch.nn import functional as F
|
||||
|
||||
from torchvision.transforms import ToTensor, ToPILImage
|
||||
|
||||
from .util_image import rgb2ycbcrTorch, ycbcr2rgbTorch
|
||||
|
||||
def adain_color_fix(target: Image, source: Image):
|
||||
# Convert images to tensors
|
||||
to_tensor = ToTensor()
|
||||
target_tensor = to_tensor(target).unsqueeze(0)
|
||||
source_tensor = to_tensor(source).unsqueeze(0)
|
||||
|
||||
# Apply adaptive instance normalization
|
||||
result_tensor = adaptive_instance_normalization(target_tensor, source_tensor)
|
||||
|
||||
# Convert tensor back to image
|
||||
to_image = ToPILImage()
|
||||
result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0))
|
||||
|
||||
return result_image
|
||||
|
||||
def wavelet_color_fix(target: Image, source: Image):
|
||||
# Convert images to tensors
|
||||
to_tensor = ToTensor()
|
||||
target_tensor = to_tensor(target).unsqueeze(0)
|
||||
source_tensor = to_tensor(source).unsqueeze(0)
|
||||
|
||||
# Apply wavelet reconstruction
|
||||
result_tensor = wavelet_reconstruction(target_tensor, source_tensor)
|
||||
|
||||
# Convert tensor back to image
|
||||
to_image = ToPILImage()
|
||||
result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0))
|
||||
|
||||
return result_image
|
||||
|
||||
def calc_mean_std(feat: Tensor, eps=1e-5):
|
||||
"""Calculate mean and std for adaptive_instance_normalization.
|
||||
Args:
|
||||
feat (Tensor): 4D tensor.
|
||||
eps (float): A small value added to the variance to avoid
|
||||
divide-by-zero. Default: 1e-5.
|
||||
"""
|
||||
size = feat.size()
|
||||
assert len(size) == 4, 'The input feature should be 4D tensor.'
|
||||
b, c = size[:2]
|
||||
feat_var = feat.reshape(b, c, -1).var(dim=2) + eps
|
||||
feat_std = feat_var.sqrt().reshape(b, c, 1, 1)
|
||||
feat_mean = feat.reshape(b, c, -1).mean(dim=2).reshape(b, c, 1, 1)
|
||||
return feat_mean, feat_std
|
||||
|
||||
def adaptive_instance_normalization(content_feat:Tensor, style_feat:Tensor):
|
||||
"""Adaptive instance normalization.
|
||||
Adjust the reference features to have the similar color and illuminations
|
||||
as those in the degradate features.
|
||||
Args:
|
||||
content_feat (Tensor): The reference feature.
|
||||
style_feat (Tensor): The degradate features.
|
||||
"""
|
||||
size = content_feat.size()
|
||||
style_mean, style_std = calc_mean_std(style_feat)
|
||||
content_mean, content_std = calc_mean_std(content_feat)
|
||||
normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
|
||||
return normalized_feat * style_std.expand(size) + style_mean.expand(size)
|
||||
|
||||
def wavelet_blur(image: Tensor, radius: int):
|
||||
"""
|
||||
Apply wavelet blur to the input tensor.
|
||||
"""
|
||||
# input shape: (1, 3, H, W)
|
||||
# convolution kernel
|
||||
kernel_vals = [
|
||||
[0.0625, 0.125, 0.0625],
|
||||
[0.125, 0.25, 0.125],
|
||||
[0.0625, 0.125, 0.0625],
|
||||
]
|
||||
kernel = torch.tensor(kernel_vals, dtype=image.dtype, device=image.device)
|
||||
# add channel dimensions to the kernel to make it a 4D tensor
|
||||
kernel = kernel[None, None]
|
||||
# repeat the kernel across all input channels
|
||||
kernel = kernel.repeat(3, 1, 1, 1)
|
||||
image = F.pad(image, (radius, radius, radius, radius), mode='replicate')
|
||||
# apply convolution
|
||||
output = F.conv2d(image, kernel, groups=3, dilation=radius)
|
||||
return output
|
||||
|
||||
def wavelet_decomposition(image: Tensor, levels=5):
|
||||
"""
|
||||
Apply wavelet decomposition to the input tensor.
|
||||
This function only returns the low frequency & the high frequency.
|
||||
"""
|
||||
high_freq = torch.zeros_like(image)
|
||||
for i in range(levels):
|
||||
radius = 2 ** i
|
||||
low_freq = wavelet_blur(image, radius)
|
||||
high_freq += (image - low_freq)
|
||||
image = low_freq
|
||||
|
||||
return high_freq, low_freq
|
||||
|
||||
def wavelet_reconstruction(content_feat:Tensor, style_feat:Tensor):
|
||||
"""
|
||||
Apply wavelet decomposition, so that the content will have the same color as the style.
|
||||
"""
|
||||
# calculate the wavelet decomposition of the content feature
|
||||
content_high_freq, content_low_freq = wavelet_decomposition(content_feat)
|
||||
del content_low_freq
|
||||
# calculate the wavelet decomposition of the style feature
|
||||
style_high_freq, style_low_freq = wavelet_decomposition(style_feat)
|
||||
del style_high_freq
|
||||
# reconstruct the content feature with the style's high frequency
|
||||
return content_high_freq + style_low_freq
|
||||
|
||||
def ycbcr_color_replace(content_feat:Tensor, style_feat:Tensor):
|
||||
"""
|
||||
Apply ycbcr decomposition, so that the content will have the same color as the style.
|
||||
"""
|
||||
content_y = rgb2ycbcrTorch(content_feat, only_y=True)
|
||||
style_ycbcr = rgb2ycbcrTorch(style_feat, only_y=False)
|
||||
|
||||
target_ycbcr = torch.cat([content_y, style_ycbcr[:, 1:,]], dim=1)
|
||||
|
||||
target_rgb = ycbcr2rgbTorch(target_ycbcr)
|
||||
|
||||
return target_rgb
|
||||
|
||||
|
||||
'''
|
||||
# --------------------------------------------------------------------------------
|
||||
# Color fixed script from Li Yi (https://github.com/pkuliyi2015/sd-webui-stablesr/blob/master/srmodule/colorfix.py)
|
||||
# --------------------------------------------------------------------------------
|
||||
'''
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.nn import functional as F
|
||||
|
||||
from torchvision.transforms import ToTensor, ToPILImage
|
||||
|
||||
from .util_image import rgb2ycbcrTorch, ycbcr2rgbTorch
|
||||
|
||||
|
||||
def calc_mean_std(feat: Tensor, eps=1e-5):
|
||||
"""Calculate mean and std for adaptive_instance_normalization.
|
||||
Args:
|
||||
feat (Tensor): 4D tensor.
|
||||
eps (float): A small value added to the variance to avoid
|
||||
divide-by-zero. Default: 1e-5.
|
||||
"""
|
||||
size = feat.size()
|
||||
assert len(size) == 4, 'The input feature should be 4D tensor.'
|
||||
b, c = size[:2]
|
||||
feat_var = feat.reshape(b, c, -1).var(dim=2) + eps
|
||||
feat_std = feat_var.sqrt().reshape(b, c, 1, 1)
|
||||
feat_mean = feat.reshape(b, c, -1).mean(dim=2).reshape(b, c, 1, 1)
|
||||
return feat_mean, feat_std
|
||||
|
||||
def adaptive_instance_normalization(content_feat:Tensor, style_feat:Tensor):
|
||||
"""Adaptive instance normalization.
|
||||
Adjust the reference features to have the similar color and illuminations
|
||||
as those in the degradate features.
|
||||
Args:
|
||||
content_feat (Tensor): The reference feature.
|
||||
style_feat (Tensor): The degradate features.
|
||||
"""
|
||||
size = content_feat.size()
|
||||
style_mean, style_std = calc_mean_std(style_feat)
|
||||
content_mean, content_std = calc_mean_std(content_feat)
|
||||
normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
|
||||
return normalized_feat * style_std.expand(size) + style_mean.expand(size)
|
||||
|
||||
def wavelet_blur(image: Tensor, radius: int):
|
||||
"""
|
||||
Apply wavelet blur to the input tensor.
|
||||
"""
|
||||
# input shape: (1, 3, H, W)
|
||||
# convolution kernel
|
||||
kernel_vals = [
|
||||
[0.0625, 0.125, 0.0625],
|
||||
[0.125, 0.25, 0.125],
|
||||
[0.0625, 0.125, 0.0625],
|
||||
]
|
||||
kernel = torch.tensor(kernel_vals, dtype=image.dtype, device=image.device)
|
||||
# add channel dimensions to the kernel to make it a 4D tensor
|
||||
kernel = kernel[None, None]
|
||||
# repeat the kernel across all input channels
|
||||
kernel = kernel.repeat(3, 1, 1, 1)
|
||||
image = F.pad(image, (radius, radius, radius, radius), mode='replicate')
|
||||
# apply convolution
|
||||
output = F.conv2d(image, kernel, groups=3, dilation=radius)
|
||||
return output
|
||||
|
||||
def wavelet_decomposition(image: Tensor, levels=5):
|
||||
"""
|
||||
Apply wavelet decomposition to the input tensor.
|
||||
This function only returns the low frequency & the high frequency.
|
||||
"""
|
||||
high_freq = torch.zeros_like(image)
|
||||
for i in range(levels):
|
||||
radius = 2 ** i
|
||||
low_freq = wavelet_blur(image, radius)
|
||||
high_freq += (image - low_freq)
|
||||
image = low_freq
|
||||
|
||||
return high_freq, low_freq
|
||||
|
||||
def wavelet_reconstruction(content_feat:Tensor, style_feat:Tensor):
|
||||
"""
|
||||
Apply wavelet decomposition, so that the content will have the same color as the style.
|
||||
"""
|
||||
# calculate the wavelet decomposition of the content feature
|
||||
content_high_freq, content_low_freq = wavelet_decomposition(content_feat)
|
||||
del content_low_freq
|
||||
# calculate the wavelet decomposition of the style feature
|
||||
style_high_freq, style_low_freq = wavelet_decomposition(style_feat)
|
||||
del style_high_freq
|
||||
# reconstruct the content feature with the style's high frequency
|
||||
return content_high_freq + style_low_freq
|
||||
|
||||
def ycbcr_color_replace(content_feat:Tensor, style_feat:Tensor):
|
||||
"""
|
||||
Apply ycbcr decomposition, so that the content will have the same color as the style.
|
||||
"""
|
||||
content_y = rgb2ycbcrTorch(content_feat, only_y=True)
|
||||
style_ycbcr = rgb2ycbcrTorch(style_feat, only_y=False)
|
||||
|
||||
target_ycbcr = torch.cat([content_y, style_ycbcr[:, 1:,]], dim=1)
|
||||
|
||||
target_rgb = ycbcr2rgbTorch(target_ycbcr)
|
||||
|
||||
return target_rgb
|
||||
|
||||
|
||||
@@ -7,7 +7,6 @@ import random
|
||||
import requests
|
||||
import importlib
|
||||
from pathlib import Path
|
||||
from PIL import Image
|
||||
|
||||
def mkdir(dir_path, delete=False, parents=True):
|
||||
import shutil
|
||||
@@ -20,11 +19,15 @@ def mkdir(dir_path, delete=False, parents=True):
|
||||
dir_path.mkdir(parents=parents)
|
||||
|
||||
def get_obj_from_str(string, reload=False):
|
||||
current_package = __package__.rsplit(".", 1)[0]
|
||||
is_relative_import = string.startswith(".")
|
||||
package = current_package if is_relative_import else None
|
||||
|
||||
module, cls = string.rsplit(".", 1)
|
||||
if reload:
|
||||
module_imp = importlib.import_module(module)
|
||||
module_imp = importlib.import_module(module, package=package)
|
||||
importlib.reload(module_imp)
|
||||
return getattr(importlib.import_module(module, package=None), cls)
|
||||
return getattr(importlib.import_module(module, package=package), cls)
|
||||
|
||||
def instantiate_from_config(config):
|
||||
if not "target" in config:
|
||||
@@ -125,31 +128,3 @@ def write_path_to_txt(
|
||||
with open(txt_path, mode='w') as ff:
|
||||
for line in path_list:
|
||||
ff.write(line+'\n')
|
||||
|
||||
def download_image_from_url(url, dir="./"):
|
||||
# Download a file from a given URI, including minimal checks
|
||||
|
||||
# Download
|
||||
f = str(Path(dir) / os.path.basename(url)) # filename
|
||||
try:
|
||||
with open(f, "wb") as file:
|
||||
file.write(requests.get(url, timeout=10).content)
|
||||
except:
|
||||
print(f'Skip the url: {f}!')
|
||||
|
||||
# Rename (remove wildcard characters)
|
||||
src = f # original name
|
||||
for c in ["%20", "%", "*", "~", "(", ")"]:
|
||||
f = f.replace(c, "_")
|
||||
f = f[: f.index("?")] if "?" in f else f # new name
|
||||
if src != f:
|
||||
os.rename(src, f) # rename
|
||||
|
||||
# Add suffix (if missing)
|
||||
if Path(f).suffix == "":
|
||||
src = f # original name
|
||||
try:
|
||||
f += f".{Image.open(f).format.lower()}"
|
||||
os.rename(src, f) # rename
|
||||
except:
|
||||
Path(f).unlink()
|
||||
@@ -8,10 +8,7 @@ import math
|
||||
import torch
|
||||
import random
|
||||
import numpy as np
|
||||
from scipy import fft
|
||||
from pathlib import Path
|
||||
from einops import rearrange
|
||||
from skimage import img_as_ubyte, img_as_float32
|
||||
|
||||
# --------------------------Metrics----------------------------
|
||||
def ssim(img1, img2):
|
||||
@@ -273,44 +270,6 @@ def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)):
|
||||
result = result[0]
|
||||
return result
|
||||
|
||||
def img2tensor(imgs, bgr2rgb=False, out_type=torch.float32):
|
||||
"""Convert image numpy arrays into torch tensor.
|
||||
Args:
|
||||
imgs (Array or list[array]): Accept shapes:
|
||||
3) list of numpy arrays
|
||||
1) 3D numpy array of shape (H x W x 3/1);
|
||||
2) 2D Tensor of shape (H x W).
|
||||
Tensor channel should be in RGB order.
|
||||
|
||||
Returns:
|
||||
(array or list): 4D ndarray of shape (1 x C x H x W)
|
||||
"""
|
||||
|
||||
def _img2tensor(img):
|
||||
if img.ndim == 2:
|
||||
tensor = torch.from_numpy(img[None, None,]).type(out_type)
|
||||
elif img.ndim == 3:
|
||||
if bgr2rgb:
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
tensor = torch.from_numpy(rearrange(img, 'h w c -> c h w')).type(out_type).unsqueeze(0)
|
||||
else:
|
||||
raise TypeError(f'2D or 3D numpy array expected, got{img.ndim}D array')
|
||||
return tensor
|
||||
|
||||
if not (isinstance(imgs, np.ndarray) or (isinstance(imgs, list) and all(isinstance(t, np.ndarray) for t in imgs))):
|
||||
raise TypeError(f'Numpy array or list of numpy array expected, got {type(imgs)}')
|
||||
|
||||
flag_numpy = isinstance(imgs, np.ndarray)
|
||||
if flag_numpy:
|
||||
imgs = [imgs,]
|
||||
result = []
|
||||
for _img in imgs:
|
||||
result.append(_img2tensor(_img))
|
||||
|
||||
if len(result) == 1 and flag_numpy:
|
||||
result = result[0]
|
||||
return result
|
||||
|
||||
# ------------------------Image resize-----------------------------
|
||||
def imresize_np(img, scale, antialiasing=True):
|
||||
# Now the scale should be the same for H and W
|
||||
@@ -492,53 +451,6 @@ def imread(path, chn='rgb', dtype='float32', force_gray2rgb=True, force_rgba2rgb
|
||||
|
||||
return im
|
||||
|
||||
def imwrite(im_in, path, chn='rgb', dtype_in='float32', qf=None):
|
||||
'''
|
||||
Save image.
|
||||
Input:
|
||||
im: h x w x c, numpy tensor
|
||||
path: the saving path
|
||||
chn: the channel order of the im,
|
||||
'''
|
||||
im = im_in.copy()
|
||||
if isinstance(path, str):
|
||||
path = Path(path)
|
||||
if dtype_in != 'uint8':
|
||||
im = img_as_ubyte(im)
|
||||
|
||||
if chn.lower() == 'rgb' and im.ndim == 3:
|
||||
im = rgb2bgr(im)
|
||||
|
||||
if qf is not None and path.suffix.lower() in ['.jpg', '.jpeg']:
|
||||
flag = cv2.imwrite(str(path), im, [int(cv2.IMWRITE_JPEG_QUALITY), int(qf)])
|
||||
else:
|
||||
flag = cv2.imwrite(str(path), im)
|
||||
|
||||
return flag
|
||||
|
||||
def jpeg_compress(im, qf, chn_in='rgb'):
|
||||
'''
|
||||
Input:
|
||||
im: h x w x 3 array
|
||||
qf: compress factor, (0, 100]
|
||||
chn_in: 'rgb' or 'bgr'
|
||||
Return:
|
||||
Compressed Image with channel order: chn_in
|
||||
'''
|
||||
# transform to BGR channle and uint8 data type
|
||||
im_bgr = rgb2bgr(im) if chn_in.lower() == 'rgb' else im
|
||||
if im.dtype != np.dtype('uint8'): im_bgr = img_as_ubyte(im_bgr)
|
||||
|
||||
# JPEG compress
|
||||
flag, encimg = cv2.imencode('.jpg', im_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), qf])
|
||||
assert flag
|
||||
im_jpg_bgr = cv2.imdecode(encimg, 1) # uint8, BGR
|
||||
|
||||
# transform back to original channel and the original data type
|
||||
im_out = bgr2rgb(im_jpg_bgr) if chn_in.lower() == 'rgb' else im_jpg_bgr
|
||||
if im.dtype != np.dtype('uint8'): im_out = img_as_float32(im_out).astype(im.dtype)
|
||||
return im_out
|
||||
|
||||
# ------------------------Augmentation-----------------------------
|
||||
def data_aug_np(image, mode):
|
||||
'''
|
||||
@@ -680,73 +592,6 @@ def imgrad(im, pading_mode='mirror'):
|
||||
|
||||
return {'gradx': gradx, 'grady': grady, 'grad':grad}
|
||||
|
||||
def imgrad_fft(im):
|
||||
'''
|
||||
Calculate image gradient.
|
||||
Input:
|
||||
im: h x w x c numpy array
|
||||
'''
|
||||
wx = np.rot90(np.array([[0, 0, 0],
|
||||
[-1, 1, 0],
|
||||
[0, 0, 0]], dtype=np.float32), k=2)
|
||||
gradx = convfft(im, wx)
|
||||
wy = np.rot90(np.array([[0, -1, 0],
|
||||
[0, 1, 0],
|
||||
[0, 0, 0]], dtype=np.float32), k=2)
|
||||
grady = convfft(im, wy)
|
||||
grad = np.concatenate((gradx, grady), axis=2)
|
||||
|
||||
return {'gradx': gradx, 'grady': grady, 'grad':grad}
|
||||
|
||||
def convfft(im, weight):
|
||||
'''
|
||||
Convolution with FFT
|
||||
Input:
|
||||
im: h1 x w1 x c numpy array
|
||||
weight: h2 x w2 numpy array
|
||||
Output:
|
||||
out: h1 x w1 x c numpy array
|
||||
'''
|
||||
axes = (0,1)
|
||||
otf = psf2otf(weight, im.shape[:2])
|
||||
if im.ndim == 3:
|
||||
otf = np.tile(otf[:, :, None], (1,1,im.shape[2]))
|
||||
out = fft.ifft2(fft.fft2(im, axes=axes) * otf, axes=axes).real
|
||||
return out
|
||||
|
||||
def psf2otf(psf, shape):
|
||||
"""
|
||||
MATLAB psf2otf function.
|
||||
Borrowed from https://github.com/aboucaud/pypher/blob/master/pypher/pypher.py.
|
||||
Input:
|
||||
psf : h x w numpy array
|
||||
shape : list or tuple, output shape of the OTF array
|
||||
Output:
|
||||
otf : OTF array with the desirable shape
|
||||
"""
|
||||
if np.all(psf == 0):
|
||||
return np.zeros_like(psf)
|
||||
|
||||
inshape = psf.shape
|
||||
# Pad the PSF to outsize
|
||||
psf = zero_pad(psf, shape, position='corner')
|
||||
|
||||
# Circularly shift OTF so that the 'center' of the PSF is [0,0] element of the array
|
||||
for axis, axis_size in enumerate(inshape):
|
||||
psf = np.roll(psf, -int(axis_size / 2), axis=axis)
|
||||
|
||||
# Compute the OTF
|
||||
otf = fft.fft2(psf)
|
||||
|
||||
# Estimate the rough number of operations involved in the FFT
|
||||
# and discard the PSF imaginary part if within roundoff error
|
||||
# roundoff error = machine epsilon = sys.float_info.epsilon
|
||||
# or np.finfo().eps
|
||||
n_ops = np.sum(psf.size * np.log2(psf.shape))
|
||||
otf = np.real_if_close(otf, tol=n_ops)
|
||||
|
||||
return otf
|
||||
|
||||
def convtorch(im, weight, mode='reflect'):
|
||||
'''
|
||||
Image convolution with pytorch
|
||||
@@ -1136,14 +981,3 @@ class SpatialAug:
|
||||
else:
|
||||
out = data_aug_np(im, flag)
|
||||
return out
|
||||
|
||||
if __name__ == '__main__':
|
||||
im = np.random.randn(64, 64, 3).astype(np.float32)
|
||||
|
||||
grad1 = imgrad(im)['grad']
|
||||
grad2 = imgrad_fft(im)['grad']
|
||||
|
||||
error = np.abs(grad1 -grad2).max()
|
||||
mean_error = np.abs(grad1 -grad2).mean()
|
||||
print('The largest error is {:.2e}'.format(error))
|
||||
print('The mean error is {:.2e}'.format(mean_error))
|
||||
@@ -0,0 +1,17 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding:utf-8 -*-
|
||||
# Power by Zongsheng Yue 2021-11-24 20:29:36
|
||||
|
||||
def reload_model(model, ckpt):
|
||||
module_flag = list(ckpt.keys())[0].startswith('module.')
|
||||
compile_flag = '_orig_mod' in list(ckpt.keys())[0]
|
||||
|
||||
for source_key, source_value in model.state_dict().items():
|
||||
target_key = source_key
|
||||
if compile_flag and (not '_orig_mod.' in source_key):
|
||||
target_key = '_orig_mod.' + target_key
|
||||
if module_flag and (not source_key.startswith('module')):
|
||||
target_key = 'module.' + target_key
|
||||
|
||||
assert target_key in ckpt
|
||||
source_value.copy_(ckpt[target_key])
|
||||
@@ -36,7 +36,7 @@ sd_pipe:
|
||||
torch_dtype: torch.float16
|
||||
|
||||
model_start:
|
||||
target: diffusers.models.autoencoders.NoisePredictor
|
||||
target: .noise_predictor.NoisePredictor
|
||||
ckpt_path: ~ # For initializing
|
||||
params:
|
||||
in_channels: 3
|
||||
@@ -65,7 +65,7 @@ model_start:
|
||||
double_z: True
|
||||
|
||||
model_middle:
|
||||
target: diffusers.models.autoencoders.NoisePredictor
|
||||
target: .noise_predictor.NoisePredictor
|
||||
params:
|
||||
in_channels: 3
|
||||
down_block_types:
|
||||
@@ -31,7 +31,7 @@ llpips:
|
||||
verbose: True
|
||||
|
||||
model:
|
||||
target: diffusers.models.autoencoders.NoisePredictor
|
||||
target: .noise_predictor.NoisePredictor
|
||||
ckpt_start_path: ~ # only used for training the intermidiate model
|
||||
ckpt_path: ~ # For initializing
|
||||
compile: False
|
||||
@@ -1,7 +1,8 @@
|
||||
from .InvSR.inference_invsr import get_configs
|
||||
from .InvSR.sampler_invsr import InvSamplerSR, BaseSampler
|
||||
from .comfyui_invsr_trimmed import get_configs, InvSamplerSR, BaseSampler
|
||||
import torch
|
||||
from comfy.utils import ProgressBar
|
||||
from folder_paths import get_full_path, get_folder_paths
|
||||
import os
|
||||
|
||||
def split_tensor_into_batches(tensor, batch_size):
|
||||
"""
|
||||
@@ -73,8 +74,24 @@ class LoadInvSRModels:
|
||||
case "bf16":
|
||||
dtype = "torch.bfloat16"
|
||||
|
||||
|
||||
args = Namespace(bs=1, chopping_bs=8, timesteps=None, num_steps=1, cfg_path='custom_nodes/ComfyUI_InvSR/InvSR/configs/sample-sd-turbo.yaml', sd_path='models/diffusers', started_ckpt_path='custom_nodes/ComfyUI_InvSR/weights/noise_predictor_sd_turbo_v5.pth', tiled_vae=tiled_vae, color_fix='', chopping_size=128)
|
||||
cfg_path = os.path.join(
|
||||
os.path.dirname(__file__), "configs", "sample-sd-turbo.yaml"
|
||||
)
|
||||
sd_path = get_folder_paths("diffusers")[0]
|
||||
ckpt_path = get_full_path("invsr", "noise_predictor_sd_turbo_v5.pth")
|
||||
|
||||
args = Namespace(
|
||||
bs=1,
|
||||
chopping_bs=8,
|
||||
timesteps=None,
|
||||
num_steps=1,
|
||||
cfg_path=cfg_path,
|
||||
sd_path=sd_path,
|
||||
started_ckpt_path=ckpt_path,
|
||||
tiled_vae=tiled_vae,
|
||||
color_fix="",
|
||||
chopping_size=128,
|
||||
)
|
||||
configs = get_configs(args)
|
||||
configs["sd_pipe"]["params"]["torch_dtype"] = dtype
|
||||
base_sampler = BaseSampler(configs)
|
||||
@@ -107,7 +124,24 @@ class InvSRSampler:
|
||||
if color_fix == "none":
|
||||
color_fix = ""
|
||||
|
||||
args = Namespace(bs=batch_size, chopping_bs=chopping_batch_size, timesteps=None, num_steps=num_steps, cfg_path='custom_nodes/ComfyUI_InvSR/InvSR/configs/sample-sd-turbo.yaml', sd_path='models/diffusers', started_ckpt_path='custom_nodes/ComfyUI_InvSR/weights/noise_predictor_sd_turbo_v5.pth', tiled_vae=base_sampler.configs.tiled_vae, color_fix=color_fix, chopping_size=chopping_size)
|
||||
cfg_path = os.path.join(
|
||||
os.path.dirname(__file__), "configs", "sample-sd-turbo.yaml"
|
||||
)
|
||||
sd_path = get_folder_paths("diffusers")[0]
|
||||
ckpt_path = get_full_path("invsr", "noise_predictor_sd_turbo_v5.pth")
|
||||
|
||||
args = Namespace(
|
||||
bs=batch_size,
|
||||
chopping_bs=chopping_batch_size,
|
||||
timesteps=None,
|
||||
num_steps=num_steps,
|
||||
cfg_path=cfg_path,
|
||||
sd_path=sd_path,
|
||||
started_ckpt_path=ckpt_path,
|
||||
tiled_vae=base_sampler.configs.tiled_vae,
|
||||
color_fix=color_fix,
|
||||
chopping_size=chopping_size,
|
||||
)
|
||||
configs = get_configs(args)
|
||||
base_sampler.configs = get_configs(args, log=True)
|
||||
base_sampler.setup_seed(seed)
|
||||
@@ -124,4 +158,4 @@ class InvSRSampler:
|
||||
results.append(torch.from_numpy(result))
|
||||
pbar.update(1)
|
||||
|
||||
return (torch.cat(results, dim=0),)
|
||||
return (torch.cat(results, dim=0),)
|
||||
|
||||
+5
-7
@@ -1,8 +1,6 @@
|
||||
pyiqa==0.1.12
|
||||
opencv-python
|
||||
albumentations==1.4.18
|
||||
bitsandbytes
|
||||
protobuf
|
||||
python-box
|
||||
opencv-contrib-python-headless
|
||||
omegaconf
|
||||
loguru
|
||||
diffusers>=0.28.0
|
||||
numpy<2
|
||||
huggingface-hub
|
||||
transformers
|
||||
|
||||
@@ -1,286 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
Simple check list from AllenNLP repo: https://github.com/allenai/allennlp/blob/main/setup.py
|
||||
|
||||
To create the package for PyPI.
|
||||
|
||||
1. Run `make pre-release` (or `make pre-patch` for a patch release) then run `make fix-copies` to fix the index of the
|
||||
documentation.
|
||||
|
||||
If releasing on a special branch, copy the updated README.md on the main branch for the commit you will make
|
||||
for the post-release and run `make fix-copies` on the main branch as well.
|
||||
|
||||
2. Unpin specific versions from setup.py that use a git install.
|
||||
|
||||
3. Checkout the release branch (v<RELEASE>-release, for example v4.19-release), and commit these changes with the
|
||||
message: "Release: <RELEASE>" and push.
|
||||
|
||||
4. Manually trigger the "Nightly and release tests on main/release branch" workflow from the release branch. Wait for
|
||||
the tests to complete. We can safely ignore the known test failures.
|
||||
|
||||
5. Wait for the tests on main to be completed and be green (otherwise revert and fix bugs).
|
||||
|
||||
6. Add a tag in git to mark the release: "git tag v<RELEASE> -m 'Adds tag v<RELEASE> for PyPI'"
|
||||
Push the tag to git: git push --tags origin v<RELEASE>-release
|
||||
|
||||
7. Build both the sources and the wheel. Do not change anything in setup.py between
|
||||
creating the wheel and the source distribution (obviously).
|
||||
|
||||
For the wheel, run: "python setup.py bdist_wheel" in the top level directory
|
||||
(This will build a wheel for the Python version you use to build it).
|
||||
|
||||
For the sources, run: "python setup.py sdist"
|
||||
You should now have a /dist directory with both .whl and .tar.gz source versions.
|
||||
|
||||
Long story cut short, you need to run both before you can upload the distribution to the
|
||||
test PyPI and the actual PyPI servers:
|
||||
|
||||
python setup.py bdist_wheel && python setup.py sdist
|
||||
|
||||
8. Check that everything looks correct by uploading the package to the PyPI test server:
|
||||
|
||||
twine upload dist/* -r pypitest
|
||||
(pypi suggests using twine as other methods upload files via plaintext.)
|
||||
You may have to specify the repository url, use the following command then:
|
||||
twine upload dist/* -r pypitest --repository-url=https://test.pypi.org/legacy/
|
||||
|
||||
Check that you can install it in a virtualenv by running:
|
||||
pip install -i https://testpypi.python.org/pypi diffusers
|
||||
|
||||
If you are testing from a Colab Notebook, for instance, then do:
|
||||
pip install diffusers && pip uninstall diffusers
|
||||
pip install -i https://testpypi.python.org/pypi diffusers
|
||||
|
||||
Check you can run the following commands:
|
||||
python -c "from diffusers import __version__; print(__version__)"
|
||||
python -c "from diffusers import DiffusionPipeline; pipe = DiffusionPipeline.from_pretrained('fusing/unet-ldm-dummy-update'); pipe()"
|
||||
python -c "from diffusers import DiffusionPipeline; pipe = DiffusionPipeline.from_pretrained('hf-internal-testing/tiny-stable-diffusion-pipe', safety_checker=None); pipe('ah suh du')"
|
||||
python -c "from diffusers import *"
|
||||
|
||||
9. Upload the final version to the actual PyPI:
|
||||
twine upload dist/* -r pypi
|
||||
|
||||
10. Prepare the release notes and publish them on GitHub once everything is looking hunky-dory. You can use the following
|
||||
Space to fetch all the commits applicable for the release: https://huggingface.co/spaces/lysandre/github-release. Repo should
|
||||
be `huggingface/diffusers`. `tag` should be the previous release tag (v0.26.1, for example), and `branch` should be
|
||||
the latest release branch (v0.27.0-release, for example). It denotes all commits that have happened on branch
|
||||
v0.27.0-release after the tag v0.26.1 was created.
|
||||
|
||||
11. Run `make post-release` (or, for a patch release, `make post-patch`). If you were on a branch for the release,
|
||||
you need to go back to main before executing this.
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
|
||||
from setuptools import Command, find_packages, setup
|
||||
|
||||
|
||||
# IMPORTANT:
|
||||
# 1. all dependencies should be listed here with their version requirements if any
|
||||
# 2. once modified, run: `make deps_table_update` to update src/diffusers/dependency_versions_table.py
|
||||
_deps = [
|
||||
"Pillow", # keep the PIL.Image.Resampling deprecation away
|
||||
"accelerate>=0.31.0",
|
||||
"compel==0.1.8",
|
||||
"datasets",
|
||||
"filelock",
|
||||
"flax>=0.4.1",
|
||||
"hf-doc-builder>=0.3.0",
|
||||
"huggingface-hub>=0.23.2",
|
||||
"requests-mock==1.10.0",
|
||||
"importlib_metadata",
|
||||
"invisible-watermark>=0.2.0",
|
||||
"isort>=5.5.4",
|
||||
"jax>=0.4.1",
|
||||
"jaxlib>=0.4.1",
|
||||
"Jinja2",
|
||||
"k-diffusion>=0.0.12",
|
||||
"torchsde",
|
||||
"note_seq",
|
||||
"librosa",
|
||||
"numpy",
|
||||
"parameterized",
|
||||
"peft>=0.6.0",
|
||||
"protobuf>=3.20.3,<4",
|
||||
"pytest",
|
||||
"pytest-timeout",
|
||||
"pytest-xdist",
|
||||
"python>=3.8.0",
|
||||
"ruff==0.1.5",
|
||||
"safetensors>=0.3.1",
|
||||
"sentencepiece>=0.1.91,!=0.1.92",
|
||||
"GitPython<3.1.19",
|
||||
"scipy",
|
||||
"onnx",
|
||||
"regex!=2019.12.17",
|
||||
"requests",
|
||||
"tensorboard",
|
||||
"torch>=1.4",
|
||||
"torchvision",
|
||||
"transformers>=4.41.2",
|
||||
"urllib3<=2.0.0",
|
||||
"black",
|
||||
]
|
||||
|
||||
# this is a lookup table with items like:
|
||||
#
|
||||
# tokenizers: "huggingface-hub==0.8.0"
|
||||
# packaging: "packaging"
|
||||
#
|
||||
# some of the values are versioned whereas others aren't.
|
||||
deps = {b: a for a, b in (re.findall(r"^(([^!=<>~]+)(?:[!=<>~].*)?$)", x)[0] for x in _deps)}
|
||||
|
||||
# since we save this data in src/diffusers/dependency_versions_table.py it can be easily accessed from
|
||||
# anywhere. If you need to quickly access the data from this table in a shell, you can do so easily with:
|
||||
#
|
||||
# python -c 'import sys; from diffusers.dependency_versions_table import deps; \
|
||||
# print(" ".join([deps[x] for x in sys.argv[1:]]))' tokenizers datasets
|
||||
#
|
||||
# Just pass the desired package names to that script as it's shown with 2 packages above.
|
||||
#
|
||||
# If diffusers is not yet installed and the work is done from the cloned repo remember to add `PYTHONPATH=src` to the script above
|
||||
#
|
||||
# You can then feed this for example to `pip`:
|
||||
#
|
||||
# pip install -U $(python -c 'import sys; from diffusers.dependency_versions_table import deps; \
|
||||
# print(" ".join([deps[x] for x in sys.argv[1:]]))' tokenizers datasets)
|
||||
#
|
||||
|
||||
|
||||
def deps_list(*pkgs):
|
||||
return [deps[pkg] for pkg in pkgs]
|
||||
|
||||
|
||||
class DepsTableUpdateCommand(Command):
|
||||
"""
|
||||
A custom command that updates the dependency table.
|
||||
usage: python setup.py deps_table_update
|
||||
"""
|
||||
|
||||
description = "build runtime dependency table"
|
||||
user_options = [
|
||||
# format: (long option, short option, description).
|
||||
(
|
||||
"dep-table-update",
|
||||
None,
|
||||
"updates src/diffusers/dependency_versions_table.py",
|
||||
),
|
||||
]
|
||||
|
||||
def initialize_options(self):
|
||||
pass
|
||||
|
||||
def finalize_options(self):
|
||||
pass
|
||||
|
||||
def run(self):
|
||||
entries = "\n".join([f' "{k}": "{v}",' for k, v in deps.items()])
|
||||
content = [
|
||||
"# THIS FILE HAS BEEN AUTOGENERATED. To update:",
|
||||
"# 1. modify the `_deps` dict in setup.py",
|
||||
"# 2. run `make deps_table_update`",
|
||||
"deps = {",
|
||||
entries,
|
||||
"}",
|
||||
"",
|
||||
]
|
||||
target = "src/diffusers/dependency_versions_table.py"
|
||||
print(f"updating {target}")
|
||||
with open(target, "w", encoding="utf-8", newline="\n") as f:
|
||||
f.write("\n".join(content))
|
||||
|
||||
|
||||
extras = {}
|
||||
extras["quality"] = deps_list("urllib3", "isort", "ruff", "hf-doc-builder")
|
||||
extras["docs"] = deps_list("hf-doc-builder")
|
||||
extras["training"] = deps_list("accelerate", "datasets", "protobuf", "tensorboard", "Jinja2", "peft")
|
||||
extras["test"] = deps_list(
|
||||
"compel",
|
||||
"GitPython",
|
||||
"datasets",
|
||||
"Jinja2",
|
||||
"invisible-watermark",
|
||||
"k-diffusion",
|
||||
"librosa",
|
||||
"parameterized",
|
||||
"pytest",
|
||||
"pytest-timeout",
|
||||
"pytest-xdist",
|
||||
"requests-mock",
|
||||
"safetensors",
|
||||
"sentencepiece",
|
||||
"scipy",
|
||||
"torchvision",
|
||||
"transformers",
|
||||
)
|
||||
extras["torch"] = deps_list("torch", "accelerate")
|
||||
|
||||
if os.name == "nt": # windows
|
||||
extras["flax"] = [] # jax is not supported on windows
|
||||
else:
|
||||
extras["flax"] = deps_list("jax", "jaxlib", "flax")
|
||||
|
||||
extras["dev"] = (
|
||||
extras["quality"] + extras["test"] + extras["training"] + extras["docs"] + extras["torch"] + extras["flax"]
|
||||
)
|
||||
|
||||
install_requires = [
|
||||
deps["importlib_metadata"],
|
||||
deps["filelock"],
|
||||
deps["huggingface-hub"],
|
||||
deps["numpy"],
|
||||
deps["regex"],
|
||||
deps["requests"],
|
||||
deps["safetensors"],
|
||||
deps["Pillow"],
|
||||
]
|
||||
|
||||
version_range_max = max(sys.version_info[1], 10) + 1
|
||||
|
||||
setup(
|
||||
name="diffusers",
|
||||
version="0.30.0.dev0", # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots)
|
||||
description="State-of-the-art diffusion in PyTorch and JAX.",
|
||||
long_description=open("README.md", "r", encoding="utf-8").read(),
|
||||
long_description_content_type="text/markdown",
|
||||
keywords="deep learning diffusion jax pytorch stable diffusion audioldm",
|
||||
license="Apache 2.0 License",
|
||||
author="The Hugging Face team (past and future) with the help of all our contributors (https://github.com/huggingface/diffusers/graphs/contributors)",
|
||||
author_email="diffusers@huggingface.co",
|
||||
url="https://github.com/huggingface/diffusers",
|
||||
package_dir={"": "src"},
|
||||
packages=find_packages("src"),
|
||||
package_data={"diffusers": ["py.typed"]},
|
||||
include_package_data=True,
|
||||
python_requires=">=3.8.0",
|
||||
install_requires=list(install_requires),
|
||||
extras_require=extras,
|
||||
entry_points={"console_scripts": ["diffusers-cli=diffusers.commands.diffusers_cli:main"]},
|
||||
classifiers=[
|
||||
"Development Status :: 5 - Production/Stable",
|
||||
"Intended Audience :: Developers",
|
||||
"Intended Audience :: Education",
|
||||
"Intended Audience :: Science/Research",
|
||||
"License :: OSI Approved :: Apache Software License",
|
||||
"Operating System :: OS Independent",
|
||||
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||
"Programming Language :: Python :: 3",
|
||||
]
|
||||
+ [f"Programming Language :: Python :: 3.{i}" for i in range(8, version_range_max)],
|
||||
cmdclass={"deps_table_update": DepsTableUpdateCommand},
|
||||
)
|
||||
@@ -1,142 +0,0 @@
|
||||
Metadata-Version: 2.1
|
||||
Name: diffusers
|
||||
Version: 0.30.0.dev0
|
||||
Summary: State-of-the-art diffusion in PyTorch and JAX.
|
||||
Home-page: https://github.com/huggingface/diffusers
|
||||
Author: The Hugging Face team (past and future) with the help of all our contributors (https://github.com/huggingface/diffusers/graphs/contributors)
|
||||
Author-email: diffusers@huggingface.co
|
||||
License: Apache 2.0 License
|
||||
Keywords: deep learning diffusion jax pytorch stable diffusion audioldm
|
||||
Platform: UNKNOWN
|
||||
Classifier: Development Status :: 5 - Production/Stable
|
||||
Classifier: Intended Audience :: Developers
|
||||
Classifier: Intended Audience :: Education
|
||||
Classifier: Intended Audience :: Science/Research
|
||||
Classifier: License :: OSI Approved :: Apache Software License
|
||||
Classifier: Operating System :: OS Independent
|
||||
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
||||
Classifier: Programming Language :: Python :: 3
|
||||
Classifier: Programming Language :: Python :: 3.8
|
||||
Classifier: Programming Language :: Python :: 3.9
|
||||
Classifier: Programming Language :: Python :: 3.10
|
||||
Requires-Python: >=3.8.0
|
||||
Description-Content-Type: text/markdown
|
||||
Provides-Extra: dev
|
||||
Provides-Extra: docs
|
||||
Provides-Extra: flax
|
||||
Provides-Extra: quality
|
||||
Provides-Extra: test
|
||||
Provides-Extra: torch
|
||||
Provides-Extra: training
|
||||
License-File: LICENSE
|
||||
|
||||
# Arbitrary-steps Image Super-resolution via Diffusion Inversion
|
||||
|
||||
[Zongsheng Yue](https://zsyoaoa.github.io/), [Kang Liao](https://kangliao929.github.io/), [Chen Change Loy](https://www.mmlab-ntu.com/person/ccloy/)
|
||||
|
||||
[](https://arxiv.org/abs/2412.09013) [](https://replicate.com/zsyoaoa/invsr) [](https://huggingface.co/spaces/OAOA/InvSR) <a href="https://colab.research.google.com/drive/1hjgCFnAU4oUUhh9VRfTwsFN1AiIjdcSR?usp=sharing"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google colab logo"></a> 
|
||||
|
||||
<!--[](https://replicate.com/cjwbw/resshift)-->
|
||||
|
||||
|
||||
:star: If you've found InvSR useful for your research or projects, please show your support by starring this repo. Thanks! :hugs:
|
||||
|
||||
---
|
||||
>This study presents a new image super-resolution (SR) technique based on diffusion inversion, aiming at harnessing the rich image priors encapsulated in large pre-trained diffusion models to improve SR performance. We design a \textit{Partial noise Prediction} strategy to construct an intermediate state of the diffusion model, which serves as the starting sampling point. Central to our approach is a deep noise predictor to estimate the optimal noise maps for the forward diffusion process. Once trained, this noise predictor can be used to initialize the sampling process partially along the diffusion trajectory, generating the desirable high-resolution result. Compared to existing approaches, our method offers a flexible and efficient sampling mechanism that supports an arbitrary number of sampling steps, ranging from one to five. Even with a single sampling step, our method demonstrates superior or comparable performance to recent state-of-the-art approaches.
|
||||
><img src="./assets/framework.png" align="middle" width="800">
|
||||
---
|
||||
## Update
|
||||
- **2025.01.08**: Update gradio demo for batch processing.
|
||||
- **2024.12.14**: Add [](https://replicate.com/zsyoaoa/invsr).
|
||||
- **2024.12.13**: Add [](https://huggingface.co/spaces/OAOA/InvSR) and <a href="https://colab.research.google.com/drive/1hjgCFnAU4oUUhh9VRfTwsFN1AiIjdcSR?usp=sharing"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google colab logo"></a>.
|
||||
- **2024.12.11**: Create this repo.
|
||||
|
||||
## Requirements
|
||||
* Python 3.10, Pytorch 2.4.0, [xformers](https://github.com/facebookresearch/xformers) 0.0.27.post2
|
||||
* More detail (See [environment.yaml](environment.yaml))
|
||||
* A suitable [conda](https://conda.io/) environment named `invsr` can be created and activated with:
|
||||
|
||||
```
|
||||
conda create -n invsr python=3.10
|
||||
conda activate invsr
|
||||
pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu121
|
||||
pip install -U xformers==0.0.27.post2 --index-url https://download.pytorch.org/whl/cu121
|
||||
pip install -e ".[torch]"
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
## Applications
|
||||
### :point_right: Real-world Image Super-resolution
|
||||
[<img src="assets/real-7.png" height="235"/>](https://imgsli.com/MzI2MTU5) [<img src="assets/real-1.png" height="235"/>](https://imgsli.com/MzI2MTUx) [<img src="assets/real-2.png" height="235"/>](https://imgsli.com/MzI2MTUy)
|
||||
[<img src="assets/real-4.png" height="361"/>](https://imgsli.com/MzI2MTU0) [<img src="assets/real-6.png" height="361"/>](https://imgsli.com/MzI2MTU3) [<img src="assets/real-5.png" height="361"/>](https://imgsli.com/MzI2MTU1)
|
||||
|
||||
<!--### :point_right: General Image Enhancement-->
|
||||
<!--[<img src="assets/enhance-1.png" height="246.5"/>](https://imgsli.com/MzI2MTYw) [<img src="assets/enhance-2.png" height="246.5"/>](https://imgsli.com/MzI2MTYy) -->
|
||||
<!--[<img src="assets/enhance-3.png" height="207"/>](https://imgsli.com/MzI2MjAx) [<img src="assets/enhance-4.png" height="207"/>](https://imgsli.com/MzI2NTk1) [<img src="assets/enhance-5.png" height="207"/>](https://imgsli.com/MzI2MjA0)-->
|
||||
|
||||
### :point_right: AIGC Image Enhancement
|
||||
[<img src="assets/sdxl-1.png" height="272"/>](https://imgsli.com/MzI2MjQy) [<img src="assets/sdxl-2.png" height="272"/>](https://imgsli.com/MzI2MjQ1) [<img src="assets/sdxl-3.png" height="272"/>](https://imgsli.com/MzI2MjQ3)
|
||||
[<img src="assets/flux-1.png" height="272"/>](https://imgsli.com/MzI2MjQ5) [<img src="assets/flux-2.png" height="272"/>](https://imgsli.com/MzI2MjUw) [<img src="assets/flux-3.png" height="272"/>](https://imgsli.com/MzI2MjUx)
|
||||
|
||||
|
||||
## Inference
|
||||
### :rocket: Fast testing
|
||||
```
|
||||
python inference_invsr.py -i [image folder/image path] -o [result folder] --num_steps 1
|
||||
```
|
||||
1. **To deal with large images, e.g., 1k---->4k, we recommend adding the option** ``--chopping_size 256``.
|
||||
2. Other options:
|
||||
+ Specify the pre-downloaded [SD Turbo](https://huggingface.co/stabilityai/sd-turbo) Model: ``--sd_path``.
|
||||
+ Specify the pre-downloaded noise predictor: ``--started_ckpt_path``.
|
||||
+ The number of sampling steps: ``--num_steps``.
|
||||
+ If your GPU memory is limited, please add the option ``--chopping_bs 1``.
|
||||
|
||||
### :railway_car: Online Demo
|
||||
You can try our method through an online demo:
|
||||
```
|
||||
python app.py
|
||||
```
|
||||
|
||||
### :whale: Now also available in Docker
|
||||
```bash
|
||||
docker compose up -d # Go to http://127.0.0.1:7860/
|
||||
```
|
||||
|
||||
### :airplane: Reproducing our paper results
|
||||
+ Synthetic dataset of ImageNet-Test: [Google Drive](https://drive.google.com/file/d/1PRGrujx3OFilgJ7I6nW7ETIR00wlAl2m/view?usp=sharing).
|
||||
|
||||
+ Real data for image super-resolution: [RealSRV3](https://github.com/csjcai/RealSR) | [RealSet80](testdata/RealSet80)
|
||||
|
||||
+ To reproduce the quantitative results on Imagenet-Test and RealSRV3, please add the color fixing options by ``--color_fix wavelet``.
|
||||
|
||||
## Training
|
||||
### :turtle: Preparing stage
|
||||
1. Download the finetuned LPIPS model from this [link](https://huggingface.co/OAOA/InvSR/resolve/main/vgg16_sdturbo_lpips.pth?download=true) and put it in the folder of "weights".
|
||||
2. Prepare the [config](configs/sd-turbo-sr-ldis.yaml) file:
|
||||
+ SD-Turbo path: configs.sd_pipe.params.cache_dir.
|
||||
+ Training data path: data.train.params.data_source.
|
||||
+ Validation data path: data.val.params.dir_path (low-quality image) and data.val.params.extra_dir_path (high-quality image).
|
||||
+ Batchsize: configs.train.batch and configs.train.microbatch (total batchsize = microbatch * #GPUS * num_grad_accumulation)
|
||||
|
||||
### :dolphin: Begin training
|
||||
```
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --standalone --nproc_per_node=4 --nnodes=1 main.py --save_dir [Logging Folder]
|
||||
```
|
||||
|
||||
### :whale: Resume from interruption
|
||||
```
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --standalone --nproc_per_node=4 --nnodes=1 main.py --save_dir [Logging Folder] --resume save_dir/ckpts/model_xx.pth
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
This project is licensed under [NTU S-Lab License 1.0](LICENSE). Redistribution and use should follow this license.
|
||||
|
||||
## Acknowledgement
|
||||
|
||||
This project is based on [BasicSR](https://github.com/XPixelGroup/BasicSR) and [diffusers](https://github.com/huggingface/diffusers). Thanks for their awesome works.
|
||||
|
||||
### Contact
|
||||
If you have any questions, please feel free to contact me via `zsyzam@gmail.com`.
|
||||
|
||||
|
||||
@@ -1,486 +0,0 @@
|
||||
LICENSE
|
||||
README.md
|
||||
setup.py
|
||||
src/diffusers/__init__.py
|
||||
src/diffusers/callbacks.py
|
||||
src/diffusers/configuration_utils.py
|
||||
src/diffusers/dependency_versions_check.py
|
||||
src/diffusers/dependency_versions_table.py
|
||||
src/diffusers/image_processor.py
|
||||
src/diffusers/optimization.py
|
||||
src/diffusers/py.typed
|
||||
src/diffusers/training_utils.py
|
||||
src/diffusers/video_processor.py
|
||||
src/diffusers.egg-info/PKG-INFO
|
||||
src/diffusers.egg-info/SOURCES.txt
|
||||
src/diffusers.egg-info/dependency_links.txt
|
||||
src/diffusers.egg-info/entry_points.txt
|
||||
src/diffusers.egg-info/requires.txt
|
||||
src/diffusers.egg-info/top_level.txt
|
||||
src/diffusers/commands/__init__.py
|
||||
src/diffusers/commands/diffusers_cli.py
|
||||
src/diffusers/commands/env.py
|
||||
src/diffusers/commands/fp16_safetensors.py
|
||||
src/diffusers/experimental/__init__.py
|
||||
src/diffusers/experimental/rl/__init__.py
|
||||
src/diffusers/experimental/rl/value_guided_sampling.py
|
||||
src/diffusers/loaders/__init__.py
|
||||
src/diffusers/loaders/ip_adapter.py
|
||||
src/diffusers/loaders/lora_base.py
|
||||
src/diffusers/loaders/lora_conversion_utils.py
|
||||
src/diffusers/loaders/lora_pipeline.py
|
||||
src/diffusers/loaders/peft.py
|
||||
src/diffusers/loaders/single_file.py
|
||||
src/diffusers/loaders/single_file_model.py
|
||||
src/diffusers/loaders/single_file_utils.py
|
||||
src/diffusers/loaders/textual_inversion.py
|
||||
src/diffusers/loaders/unet.py
|
||||
src/diffusers/loaders/unet_loader_utils.py
|
||||
src/diffusers/loaders/utils.py
|
||||
src/diffusers/models/__init__.py
|
||||
src/diffusers/models/activations.py
|
||||
src/diffusers/models/adapter.py
|
||||
src/diffusers/models/attention.py
|
||||
src/diffusers/models/attention_flax.py
|
||||
src/diffusers/models/attention_processor.py
|
||||
src/diffusers/models/controlnet.py
|
||||
src/diffusers/models/controlnet_flax.py
|
||||
src/diffusers/models/controlnet_hunyuan.py
|
||||
src/diffusers/models/controlnet_sd3.py
|
||||
src/diffusers/models/controlnet_sparsectrl.py
|
||||
src/diffusers/models/controlnet_xs.py
|
||||
src/diffusers/models/downsampling.py
|
||||
src/diffusers/models/embeddings.py
|
||||
src/diffusers/models/embeddings_flax.py
|
||||
src/diffusers/models/lora.py
|
||||
src/diffusers/models/model_loading_utils.py
|
||||
src/diffusers/models/modeling_flax_pytorch_utils.py
|
||||
src/diffusers/models/modeling_flax_utils.py
|
||||
src/diffusers/models/modeling_outputs.py
|
||||
src/diffusers/models/modeling_pytorch_flax_utils.py
|
||||
src/diffusers/models/modeling_utils.py
|
||||
src/diffusers/models/normalization.py
|
||||
src/diffusers/models/resnet.py
|
||||
src/diffusers/models/resnet_flax.py
|
||||
src/diffusers/models/upsampling.py
|
||||
src/diffusers/models/vae_flax.py
|
||||
src/diffusers/models/vq_model.py
|
||||
src/diffusers/models/autoencoders/__init__.py
|
||||
src/diffusers/models/autoencoders/autoencoder_asym_kl.py
|
||||
src/diffusers/models/autoencoders/autoencoder_kl.py
|
||||
src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py
|
||||
src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py
|
||||
src/diffusers/models/autoencoders/autoencoder_oobleck.py
|
||||
src/diffusers/models/autoencoders/autoencoder_tiny.py
|
||||
src/diffusers/models/autoencoders/consistency_decoder_vae.py
|
||||
src/diffusers/models/autoencoders/vae.py
|
||||
src/diffusers/models/autoencoders/vq_model.py
|
||||
src/diffusers/models/transformers/__init__.py
|
||||
src/diffusers/models/transformers/auraflow_transformer_2d.py
|
||||
src/diffusers/models/transformers/cogvideox_transformer_3d.py
|
||||
src/diffusers/models/transformers/dit_transformer_2d.py
|
||||
src/diffusers/models/transformers/dual_transformer_2d.py
|
||||
src/diffusers/models/transformers/hunyuan_transformer_2d.py
|
||||
src/diffusers/models/transformers/latte_transformer_3d.py
|
||||
src/diffusers/models/transformers/lumina_nextdit2d.py
|
||||
src/diffusers/models/transformers/pixart_transformer_2d.py
|
||||
src/diffusers/models/transformers/prior_transformer.py
|
||||
src/diffusers/models/transformers/stable_audio_transformer.py
|
||||
src/diffusers/models/transformers/t5_film_transformer.py
|
||||
src/diffusers/models/transformers/transformer_2d.py
|
||||
src/diffusers/models/transformers/transformer_flux.py
|
||||
src/diffusers/models/transformers/transformer_sd3.py
|
||||
src/diffusers/models/transformers/transformer_temporal.py
|
||||
src/diffusers/models/unets/__init__.py
|
||||
src/diffusers/models/unets/unet_1d.py
|
||||
src/diffusers/models/unets/unet_1d_blocks.py
|
||||
src/diffusers/models/unets/unet_2d.py
|
||||
src/diffusers/models/unets/unet_2d_blocks.py
|
||||
src/diffusers/models/unets/unet_2d_blocks_flax.py
|
||||
src/diffusers/models/unets/unet_2d_condition.py
|
||||
src/diffusers/models/unets/unet_2d_condition_discriminator.py
|
||||
src/diffusers/models/unets/unet_2d_condition_flax.py
|
||||
src/diffusers/models/unets/unet_3d_blocks.py
|
||||
src/diffusers/models/unets/unet_3d_condition.py
|
||||
src/diffusers/models/unets/unet_i2vgen_xl.py
|
||||
src/diffusers/models/unets/unet_kandinsky3.py
|
||||
src/diffusers/models/unets/unet_motion_model.py
|
||||
src/diffusers/models/unets/unet_spatio_temporal_condition.py
|
||||
src/diffusers/models/unets/unet_stable_cascade.py
|
||||
src/diffusers/models/unets/uvit_2d.py
|
||||
src/diffusers/pipelines/__init__.py
|
||||
src/diffusers/pipelines/auto_pipeline.py
|
||||
src/diffusers/pipelines/free_init_utils.py
|
||||
src/diffusers/pipelines/free_noise_utils.py
|
||||
src/diffusers/pipelines/onnx_utils.py
|
||||
src/diffusers/pipelines/pipeline_flax_utils.py
|
||||
src/diffusers/pipelines/pipeline_loading_utils.py
|
||||
src/diffusers/pipelines/pipeline_utils.py
|
||||
src/diffusers/pipelines/amused/__init__.py
|
||||
src/diffusers/pipelines/amused/pipeline_amused.py
|
||||
src/diffusers/pipelines/amused/pipeline_amused_img2img.py
|
||||
src/diffusers/pipelines/amused/pipeline_amused_inpaint.py
|
||||
src/diffusers/pipelines/animatediff/__init__.py
|
||||
src/diffusers/pipelines/animatediff/pipeline_animatediff.py
|
||||
src/diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py
|
||||
src/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py
|
||||
src/diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py
|
||||
src/diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py
|
||||
src/diffusers/pipelines/animatediff/pipeline_output.py
|
||||
src/diffusers/pipelines/audioldm/__init__.py
|
||||
src/diffusers/pipelines/audioldm/pipeline_audioldm.py
|
||||
src/diffusers/pipelines/audioldm2/__init__.py
|
||||
src/diffusers/pipelines/audioldm2/modeling_audioldm2.py
|
||||
src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py
|
||||
src/diffusers/pipelines/aura_flow/__init__.py
|
||||
src/diffusers/pipelines/aura_flow/pipeline_aura_flow.py
|
||||
src/diffusers/pipelines/blip_diffusion/__init__.py
|
||||
src/diffusers/pipelines/blip_diffusion/blip_image_processing.py
|
||||
src/diffusers/pipelines/blip_diffusion/modeling_blip2.py
|
||||
src/diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py
|
||||
src/diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py
|
||||
src/diffusers/pipelines/cogvideo/__init__.py
|
||||
src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py
|
||||
src/diffusers/pipelines/consistency_models/__init__.py
|
||||
src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py
|
||||
src/diffusers/pipelines/controlnet/__init__.py
|
||||
src/diffusers/pipelines/controlnet/multicontrolnet.py
|
||||
src/diffusers/pipelines/controlnet/pipeline_controlnet.py
|
||||
src/diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py
|
||||
src/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py
|
||||
src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py
|
||||
src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py
|
||||
src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py
|
||||
src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py
|
||||
src/diffusers/pipelines/controlnet/pipeline_flax_controlnet.py
|
||||
src/diffusers/pipelines/controlnet_hunyuandit/__init__.py
|
||||
src/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py
|
||||
src/diffusers/pipelines/controlnet_sd3/__init__.py
|
||||
src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py
|
||||
src/diffusers/pipelines/controlnet_xs/__init__.py
|
||||
src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py
|
||||
src/diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py
|
||||
src/diffusers/pipelines/dance_diffusion/__init__.py
|
||||
src/diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py
|
||||
src/diffusers/pipelines/ddim/__init__.py
|
||||
src/diffusers/pipelines/ddim/pipeline_ddim.py
|
||||
src/diffusers/pipelines/ddpm/__init__.py
|
||||
src/diffusers/pipelines/ddpm/pipeline_ddpm.py
|
||||
src/diffusers/pipelines/deepfloyd_if/__init__.py
|
||||
src/diffusers/pipelines/deepfloyd_if/pipeline_if.py
|
||||
src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py
|
||||
src/diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py
|
||||
src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py
|
||||
src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py
|
||||
src/diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py
|
||||
src/diffusers/pipelines/deepfloyd_if/pipeline_output.py
|
||||
src/diffusers/pipelines/deepfloyd_if/safety_checker.py
|
||||
src/diffusers/pipelines/deepfloyd_if/timesteps.py
|
||||
src/diffusers/pipelines/deepfloyd_if/watermark.py
|
||||
src/diffusers/pipelines/deprecated/__init__.py
|
||||
src/diffusers/pipelines/deprecated/alt_diffusion/__init__.py
|
||||
src/diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py
|
||||
src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py
|
||||
src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py
|
||||
src/diffusers/pipelines/deprecated/alt_diffusion/pipeline_output.py
|
||||
src/diffusers/pipelines/deprecated/audio_diffusion/__init__.py
|
||||
src/diffusers/pipelines/deprecated/audio_diffusion/mel.py
|
||||
src/diffusers/pipelines/deprecated/audio_diffusion/pipeline_audio_diffusion.py
|
||||
src/diffusers/pipelines/deprecated/latent_diffusion_uncond/__init__.py
|
||||
src/diffusers/pipelines/deprecated/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py
|
||||
src/diffusers/pipelines/deprecated/pndm/__init__.py
|
||||
src/diffusers/pipelines/deprecated/pndm/pipeline_pndm.py
|
||||
src/diffusers/pipelines/deprecated/repaint/__init__.py
|
||||
src/diffusers/pipelines/deprecated/repaint/pipeline_repaint.py
|
||||
src/diffusers/pipelines/deprecated/score_sde_ve/__init__.py
|
||||
src/diffusers/pipelines/deprecated/score_sde_ve/pipeline_score_sde_ve.py
|
||||
src/diffusers/pipelines/deprecated/spectrogram_diffusion/__init__.py
|
||||
src/diffusers/pipelines/deprecated/spectrogram_diffusion/continuous_encoder.py
|
||||
src/diffusers/pipelines/deprecated/spectrogram_diffusion/midi_utils.py
|
||||
src/diffusers/pipelines/deprecated/spectrogram_diffusion/notes_encoder.py
|
||||
src/diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py
|
||||
src/diffusers/pipelines/deprecated/stable_diffusion_variants/__init__.py
|
||||
src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py
|
||||
src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_onnx_stable_diffusion_inpaint_legacy.py
|
||||
src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py
|
||||
src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py
|
||||
src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py
|
||||
src/diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py
|
||||
src/diffusers/pipelines/deprecated/stochastic_karras_ve/__init__.py
|
||||
src/diffusers/pipelines/deprecated/stochastic_karras_ve/pipeline_stochastic_karras_ve.py
|
||||
src/diffusers/pipelines/deprecated/versatile_diffusion/__init__.py
|
||||
src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py
|
||||
src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py
|
||||
src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py
|
||||
src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py
|
||||
src/diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py
|
||||
src/diffusers/pipelines/deprecated/vq_diffusion/__init__.py
|
||||
src/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py
|
||||
src/diffusers/pipelines/dit/__init__.py
|
||||
src/diffusers/pipelines/dit/pipeline_dit.py
|
||||
src/diffusers/pipelines/flux/__init__.py
|
||||
src/diffusers/pipelines/flux/pipeline_flux.py
|
||||
src/diffusers/pipelines/flux/pipeline_output.py
|
||||
src/diffusers/pipelines/hunyuandit/__init__.py
|
||||
src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py
|
||||
src/diffusers/pipelines/i2vgen_xl/__init__.py
|
||||
src/diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py
|
||||
src/diffusers/pipelines/kandinsky/__init__.py
|
||||
src/diffusers/pipelines/kandinsky/pipeline_kandinsky.py
|
||||
src/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py
|
||||
src/diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py
|
||||
src/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py
|
||||
src/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py
|
||||
src/diffusers/pipelines/kandinsky/text_encoder.py
|
||||
src/diffusers/pipelines/kandinsky2_2/__init__.py
|
||||
src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py
|
||||
src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py
|
||||
src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py
|
||||
src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py
|
||||
src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py
|
||||
src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py
|
||||
src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py
|
||||
src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py
|
||||
src/diffusers/pipelines/kandinsky3/__init__.py
|
||||
src/diffusers/pipelines/kandinsky3/convert_kandinsky3_unet.py
|
||||
src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py
|
||||
src/diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py
|
||||
src/diffusers/pipelines/kolors/__init__.py
|
||||
src/diffusers/pipelines/kolors/pipeline_kolors.py
|
||||
src/diffusers/pipelines/kolors/pipeline_kolors_img2img.py
|
||||
src/diffusers/pipelines/kolors/pipeline_output.py
|
||||
src/diffusers/pipelines/kolors/text_encoder.py
|
||||
src/diffusers/pipelines/kolors/tokenizer.py
|
||||
src/diffusers/pipelines/latent_consistency_models/__init__.py
|
||||
src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py
|
||||
src/diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py
|
||||
src/diffusers/pipelines/latent_diffusion/__init__.py
|
||||
src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py
|
||||
src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py
|
||||
src/diffusers/pipelines/latte/__init__.py
|
||||
src/diffusers/pipelines/latte/pipeline_latte.py
|
||||
src/diffusers/pipelines/ledits_pp/__init__.py
|
||||
src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py
|
||||
src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py
|
||||
src/diffusers/pipelines/ledits_pp/pipeline_output.py
|
||||
src/diffusers/pipelines/lumina/__init__.py
|
||||
src/diffusers/pipelines/lumina/pipeline_lumina.py
|
||||
src/diffusers/pipelines/marigold/__init__.py
|
||||
src/diffusers/pipelines/marigold/marigold_image_processing.py
|
||||
src/diffusers/pipelines/marigold/pipeline_marigold_depth.py
|
||||
src/diffusers/pipelines/marigold/pipeline_marigold_normals.py
|
||||
src/diffusers/pipelines/musicldm/__init__.py
|
||||
src/diffusers/pipelines/musicldm/pipeline_musicldm.py
|
||||
src/diffusers/pipelines/pag/__init__.py
|
||||
src/diffusers/pipelines/pag/pag_utils.py
|
||||
src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py
|
||||
src/diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py
|
||||
src/diffusers/pipelines/pag/pipeline_pag_hunyuandit.py
|
||||
src/diffusers/pipelines/pag/pipeline_pag_kolors.py
|
||||
src/diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py
|
||||
src/diffusers/pipelines/pag/pipeline_pag_sd.py
|
||||
src/diffusers/pipelines/pag/pipeline_pag_sd_3.py
|
||||
src/diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py
|
||||
src/diffusers/pipelines/pag/pipeline_pag_sd_xl.py
|
||||
src/diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py
|
||||
src/diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py
|
||||
src/diffusers/pipelines/paint_by_example/__init__.py
|
||||
src/diffusers/pipelines/paint_by_example/image_encoder.py
|
||||
src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py
|
||||
src/diffusers/pipelines/pia/__init__.py
|
||||
src/diffusers/pipelines/pia/pipeline_pia.py
|
||||
src/diffusers/pipelines/pixart_alpha/__init__.py
|
||||
src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py
|
||||
src/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py
|
||||
src/diffusers/pipelines/semantic_stable_diffusion/__init__.py
|
||||
src/diffusers/pipelines/semantic_stable_diffusion/pipeline_output.py
|
||||
src/diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py
|
||||
src/diffusers/pipelines/shap_e/__init__.py
|
||||
src/diffusers/pipelines/shap_e/camera.py
|
||||
src/diffusers/pipelines/shap_e/pipeline_shap_e.py
|
||||
src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py
|
||||
src/diffusers/pipelines/shap_e/renderer.py
|
||||
src/diffusers/pipelines/stable_audio/__init__.py
|
||||
src/diffusers/pipelines/stable_audio/modeling_stable_audio.py
|
||||
src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py
|
||||
src/diffusers/pipelines/stable_cascade/__init__.py
|
||||
src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py
|
||||
src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py
|
||||
src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py
|
||||
src/diffusers/pipelines/stable_diffusion/__init__.py
|
||||
src/diffusers/pipelines/stable_diffusion/clip_image_project_model.py
|
||||
src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_img2img.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion_inpaint.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_output.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inversion_sr.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py
|
||||
src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py
|
||||
src/diffusers/pipelines/stable_diffusion/safety_checker.py
|
||||
src/diffusers/pipelines/stable_diffusion/safety_checker_flax.py
|
||||
src/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py
|
||||
src/diffusers/pipelines/stable_diffusion_3/__init__.py
|
||||
src/diffusers/pipelines/stable_diffusion_3/pipeline_output.py
|
||||
src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py
|
||||
src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py
|
||||
src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py
|
||||
src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inversion_sr.py
|
||||
src/diffusers/pipelines/stable_diffusion_attend_and_excite/__init__.py
|
||||
src/diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py
|
||||
src/diffusers/pipelines/stable_diffusion_diffedit/__init__.py
|
||||
src/diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py
|
||||
src/diffusers/pipelines/stable_diffusion_gligen/__init__.py
|
||||
src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py
|
||||
src/diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py
|
||||
src/diffusers/pipelines/stable_diffusion_k_diffusion/__init__.py
|
||||
src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py
|
||||
src/diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py
|
||||
src/diffusers/pipelines/stable_diffusion_ldm3d/__init__.py
|
||||
src/diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py
|
||||
src/diffusers/pipelines/stable_diffusion_panorama/__init__.py
|
||||
src/diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py
|
||||
src/diffusers/pipelines/stable_diffusion_safe/__init__.py
|
||||
src/diffusers/pipelines/stable_diffusion_safe/pipeline_output.py
|
||||
src/diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py
|
||||
src/diffusers/pipelines/stable_diffusion_safe/safety_checker.py
|
||||
src/diffusers/pipelines/stable_diffusion_sag/__init__.py
|
||||
src/diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py
|
||||
src/diffusers/pipelines/stable_diffusion_xl/__init__.py
|
||||
src/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py
|
||||
src/diffusers/pipelines/stable_diffusion_xl/pipeline_output.py
|
||||
src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py
|
||||
src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py
|
||||
src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py
|
||||
src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py
|
||||
src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inversion_sr.py
|
||||
src/diffusers/pipelines/stable_diffusion_xl/watermark.py
|
||||
src/diffusers/pipelines/stable_video_diffusion/__init__.py
|
||||
src/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py
|
||||
src/diffusers/pipelines/t2i_adapter/__init__.py
|
||||
src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py
|
||||
src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py
|
||||
src/diffusers/pipelines/text_to_video_synthesis/__init__.py
|
||||
src/diffusers/pipelines/text_to_video_synthesis/pipeline_output.py
|
||||
src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py
|
||||
src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py
|
||||
src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py
|
||||
src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py
|
||||
src/diffusers/pipelines/unclip/__init__.py
|
||||
src/diffusers/pipelines/unclip/pipeline_unclip.py
|
||||
src/diffusers/pipelines/unclip/pipeline_unclip_image_variation.py
|
||||
src/diffusers/pipelines/unclip/text_proj.py
|
||||
src/diffusers/pipelines/unidiffuser/__init__.py
|
||||
src/diffusers/pipelines/unidiffuser/modeling_text_decoder.py
|
||||
src/diffusers/pipelines/unidiffuser/modeling_uvit.py
|
||||
src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py
|
||||
src/diffusers/pipelines/wuerstchen/__init__.py
|
||||
src/diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py
|
||||
src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py
|
||||
src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_diffnext.py
|
||||
src/diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py
|
||||
src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py
|
||||
src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py
|
||||
src/diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py
|
||||
src/diffusers/quantizers/__init__.py
|
||||
src/diffusers/quantizers/auto.py
|
||||
src/diffusers/quantizers/base.py
|
||||
src/diffusers/quantizers/quantization_config.py
|
||||
src/diffusers/quantizers/bitsandbytes/__init__.py
|
||||
src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py
|
||||
src/diffusers/quantizers/bitsandbytes/utils.py
|
||||
src/diffusers/schedulers/__init__.py
|
||||
src/diffusers/schedulers/scheduling_amused.py
|
||||
src/diffusers/schedulers/scheduling_consistency_decoder.py
|
||||
src/diffusers/schedulers/scheduling_consistency_models.py
|
||||
src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py
|
||||
src/diffusers/schedulers/scheduling_ddim.py
|
||||
src/diffusers/schedulers/scheduling_ddim_cogvideox.py
|
||||
src/diffusers/schedulers/scheduling_ddim_flax.py
|
||||
src/diffusers/schedulers/scheduling_ddim_inverse.py
|
||||
src/diffusers/schedulers/scheduling_ddim_parallel.py
|
||||
src/diffusers/schedulers/scheduling_ddpm.py
|
||||
src/diffusers/schedulers/scheduling_ddpm_flax.py
|
||||
src/diffusers/schedulers/scheduling_ddpm_parallel.py
|
||||
src/diffusers/schedulers/scheduling_ddpm_wuerstchen.py
|
||||
src/diffusers/schedulers/scheduling_deis_multistep.py
|
||||
src/diffusers/schedulers/scheduling_dpm_cogvideox.py
|
||||
src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
|
||||
src/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py
|
||||
src/diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py
|
||||
src/diffusers/schedulers/scheduling_dpmsolver_sde.py
|
||||
src/diffusers/schedulers/scheduling_dpmsolver_singlestep.py
|
||||
src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py
|
||||
src/diffusers/schedulers/scheduling_edm_euler.py
|
||||
src/diffusers/schedulers/scheduling_euler_ancestral_discrete.py
|
||||
src/diffusers/schedulers/scheduling_euler_discrete.py
|
||||
src/diffusers/schedulers/scheduling_euler_discrete_flax.py
|
||||
src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py
|
||||
src/diffusers/schedulers/scheduling_flow_match_heun_discrete.py
|
||||
src/diffusers/schedulers/scheduling_heun_discrete.py
|
||||
src/diffusers/schedulers/scheduling_ipndm.py
|
||||
src/diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py
|
||||
src/diffusers/schedulers/scheduling_k_dpm_2_discrete.py
|
||||
src/diffusers/schedulers/scheduling_karras_ve_flax.py
|
||||
src/diffusers/schedulers/scheduling_lcm.py
|
||||
src/diffusers/schedulers/scheduling_lms_discrete.py
|
||||
src/diffusers/schedulers/scheduling_lms_discrete_flax.py
|
||||
src/diffusers/schedulers/scheduling_pndm.py
|
||||
src/diffusers/schedulers/scheduling_pndm_flax.py
|
||||
src/diffusers/schedulers/scheduling_repaint.py
|
||||
src/diffusers/schedulers/scheduling_sasolver.py
|
||||
src/diffusers/schedulers/scheduling_sde_ve.py
|
||||
src/diffusers/schedulers/scheduling_sde_ve_flax.py
|
||||
src/diffusers/schedulers/scheduling_tcd.py
|
||||
src/diffusers/schedulers/scheduling_unclip.py
|
||||
src/diffusers/schedulers/scheduling_unipc_multistep.py
|
||||
src/diffusers/schedulers/scheduling_utils.py
|
||||
src/diffusers/schedulers/scheduling_utils_flax.py
|
||||
src/diffusers/schedulers/scheduling_vq_diffusion.py
|
||||
src/diffusers/schedulers/deprecated/__init__.py
|
||||
src/diffusers/schedulers/deprecated/scheduling_karras_ve.py
|
||||
src/diffusers/schedulers/deprecated/scheduling_sde_vp.py
|
||||
src/diffusers/utils/__init__.py
|
||||
src/diffusers/utils/accelerate_utils.py
|
||||
src/diffusers/utils/constants.py
|
||||
src/diffusers/utils/deprecation_utils.py
|
||||
src/diffusers/utils/doc_utils.py
|
||||
src/diffusers/utils/dummy_flax_and_transformers_objects.py
|
||||
src/diffusers/utils/dummy_flax_objects.py
|
||||
src/diffusers/utils/dummy_note_seq_objects.py
|
||||
src/diffusers/utils/dummy_onnx_objects.py
|
||||
src/diffusers/utils/dummy_pt_objects.py
|
||||
src/diffusers/utils/dummy_torch_and_librosa_objects.py
|
||||
src/diffusers/utils/dummy_torch_and_scipy_objects.py
|
||||
src/diffusers/utils/dummy_torch_and_torchsde_objects.py
|
||||
src/diffusers/utils/dummy_torch_and_transformers_and_k_diffusion_objects.py
|
||||
src/diffusers/utils/dummy_torch_and_transformers_and_onnx_objects.py
|
||||
src/diffusers/utils/dummy_torch_and_transformers_and_sentencepiece_objects.py
|
||||
src/diffusers/utils/dummy_torch_and_transformers_objects.py
|
||||
src/diffusers/utils/dummy_transformers_and_torch_and_note_seq_objects.py
|
||||
src/diffusers/utils/dynamic_modules_utils.py
|
||||
src/diffusers/utils/export_utils.py
|
||||
src/diffusers/utils/hub_utils.py
|
||||
src/diffusers/utils/import_utils.py
|
||||
src/diffusers/utils/loading_utils.py
|
||||
src/diffusers/utils/logging.py
|
||||
src/diffusers/utils/model_card_template.md
|
||||
src/diffusers/utils/outputs.py
|
||||
src/diffusers/utils/peft_utils.py
|
||||
src/diffusers/utils/pil_utils.py
|
||||
src/diffusers/utils/state_dict_utils.py
|
||||
src/diffusers/utils/testing_utils.py
|
||||
src/diffusers/utils/torch_utils.py
|
||||
src/diffusers/utils/versions.py
|
||||
@@ -1 +0,0 @@
|
||||
|
||||
@@ -1,3 +0,0 @@
|
||||
[console_scripts]
|
||||
diffusers-cli = diffusers.commands.diffusers_cli:main
|
||||
|
||||
@@ -1,88 +0,0 @@
|
||||
Pillow
|
||||
filelock
|
||||
huggingface-hub>=0.23.2
|
||||
importlib_metadata
|
||||
numpy
|
||||
regex!=2019.12.17
|
||||
requests
|
||||
safetensors>=0.3.1
|
||||
|
||||
[dev]
|
||||
GitPython<3.1.19
|
||||
Jinja2
|
||||
Jinja2
|
||||
accelerate>=0.31.0
|
||||
accelerate>=0.31.0
|
||||
compel==0.1.8
|
||||
datasets
|
||||
datasets
|
||||
flax>=0.4.1
|
||||
hf-doc-builder>=0.3.0
|
||||
hf-doc-builder>=0.3.0
|
||||
invisible-watermark>=0.2.0
|
||||
isort>=5.5.4
|
||||
jax>=0.4.1
|
||||
jaxlib>=0.4.1
|
||||
k-diffusion>=0.0.12
|
||||
librosa
|
||||
parameterized
|
||||
peft>=0.6.0
|
||||
protobuf<4,>=3.20.3
|
||||
pytest
|
||||
pytest-timeout
|
||||
pytest-xdist
|
||||
requests-mock==1.10.0
|
||||
ruff==0.1.5
|
||||
safetensors>=0.3.1
|
||||
scipy
|
||||
sentencepiece!=0.1.92,>=0.1.91
|
||||
tensorboard
|
||||
torch>=1.4
|
||||
torchvision
|
||||
transformers>=4.41.2
|
||||
urllib3<=2.0.0
|
||||
|
||||
[docs]
|
||||
hf-doc-builder>=0.3.0
|
||||
|
||||
[flax]
|
||||
flax>=0.4.1
|
||||
jax>=0.4.1
|
||||
jaxlib>=0.4.1
|
||||
|
||||
[quality]
|
||||
hf-doc-builder>=0.3.0
|
||||
isort>=5.5.4
|
||||
ruff==0.1.5
|
||||
urllib3<=2.0.0
|
||||
|
||||
[test]
|
||||
GitPython<3.1.19
|
||||
Jinja2
|
||||
compel==0.1.8
|
||||
datasets
|
||||
invisible-watermark>=0.2.0
|
||||
k-diffusion>=0.0.12
|
||||
librosa
|
||||
parameterized
|
||||
pytest
|
||||
pytest-timeout
|
||||
pytest-xdist
|
||||
requests-mock==1.10.0
|
||||
safetensors>=0.3.1
|
||||
scipy
|
||||
sentencepiece!=0.1.92,>=0.1.91
|
||||
torchvision
|
||||
transformers>=4.41.2
|
||||
|
||||
[torch]
|
||||
accelerate>=0.31.0
|
||||
torch>=1.4
|
||||
|
||||
[training]
|
||||
Jinja2
|
||||
accelerate>=0.31.0
|
||||
datasets
|
||||
peft>=0.6.0
|
||||
protobuf<4,>=3.20.3
|
||||
tensorboard
|
||||
@@ -1 +0,0 @@
|
||||
diffusers
|
||||
@@ -1,959 +0,0 @@
|
||||
__version__ = "0.32.0.dev0"
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from .utils import (
|
||||
DIFFUSERS_SLOW_IMPORT,
|
||||
OptionalDependencyNotAvailable,
|
||||
_LazyModule,
|
||||
is_flax_available,
|
||||
is_k_diffusion_available,
|
||||
is_librosa_available,
|
||||
is_note_seq_available,
|
||||
is_onnx_available,
|
||||
is_scipy_available,
|
||||
is_sentencepiece_available,
|
||||
is_torch_available,
|
||||
is_torchsde_available,
|
||||
is_transformers_available,
|
||||
)
|
||||
|
||||
|
||||
# Lazy Import based on
|
||||
# https://github.com/huggingface/transformers/blob/main/src/transformers/__init__.py
|
||||
|
||||
# When adding a new object to this init, please add it to `_import_structure`. The `_import_structure` is a dictionary submodule to list of object names,
|
||||
# and is used to defer the actual importing for when the objects are requested.
|
||||
# This way `import diffusers` provides the names in the namespace without actually importing anything (and especially none of the backends).
|
||||
|
||||
_import_structure = {
|
||||
"configuration_utils": ["ConfigMixin"],
|
||||
"loaders": ["FromOriginalModelMixin"],
|
||||
"models": [],
|
||||
"pipelines": [],
|
||||
"quantizers.quantization_config": ["BitsAndBytesConfig"],
|
||||
"schedulers": [],
|
||||
"utils": [
|
||||
"OptionalDependencyNotAvailable",
|
||||
"is_flax_available",
|
||||
"is_inflect_available",
|
||||
"is_invisible_watermark_available",
|
||||
"is_k_diffusion_available",
|
||||
"is_k_diffusion_version",
|
||||
"is_librosa_available",
|
||||
"is_note_seq_available",
|
||||
"is_onnx_available",
|
||||
"is_scipy_available",
|
||||
"is_torch_available",
|
||||
"is_torchsde_available",
|
||||
"is_transformers_available",
|
||||
"is_transformers_version",
|
||||
"is_unidecode_available",
|
||||
"logging",
|
||||
],
|
||||
}
|
||||
|
||||
try:
|
||||
if not is_onnx_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils import dummy_onnx_objects # noqa F403
|
||||
|
||||
_import_structure["utils.dummy_onnx_objects"] = [
|
||||
name for name in dir(dummy_onnx_objects) if not name.startswith("_")
|
||||
]
|
||||
|
||||
else:
|
||||
_import_structure["pipelines"].extend(["OnnxRuntimeModel"])
|
||||
|
||||
try:
|
||||
if not is_torch_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils import dummy_pt_objects # noqa F403
|
||||
|
||||
_import_structure["utils.dummy_pt_objects"] = [name for name in dir(dummy_pt_objects) if not name.startswith("_")]
|
||||
|
||||
else:
|
||||
_import_structure["models"].extend(
|
||||
[
|
||||
"AsymmetricAutoencoderKL",
|
||||
"AuraFlowTransformer2DModel",
|
||||
"AutoencoderKL",
|
||||
"AutoencoderKLCogVideoX",
|
||||
"AutoencoderKLTemporalDecoder",
|
||||
"AutoencoderOobleck",
|
||||
"AutoencoderTiny",
|
||||
"CogVideoXTransformer3DModel",
|
||||
"CogView3PlusTransformer2DModel",
|
||||
"ConsistencyDecoderVAE",
|
||||
"ControlNetModel",
|
||||
"ControlNetXSAdapter",
|
||||
"DiTTransformer2DModel",
|
||||
"FluxControlNetModel",
|
||||
"FluxMultiControlNetModel",
|
||||
"FluxTransformer2DModel",
|
||||
"HunyuanDiT2DControlNetModel",
|
||||
"HunyuanDiT2DModel",
|
||||
"HunyuanDiT2DMultiControlNetModel",
|
||||
"I2VGenXLUNet",
|
||||
"Kandinsky3UNet",
|
||||
"LatteTransformer3DModel",
|
||||
"LuminaNextDiT2DModel",
|
||||
"ModelMixin",
|
||||
"MotionAdapter",
|
||||
"MultiAdapter",
|
||||
"PixArtTransformer2DModel",
|
||||
"PriorTransformer",
|
||||
"SD3ControlNetModel",
|
||||
"SD3MultiControlNetModel",
|
||||
"SD3Transformer2DModel",
|
||||
"SparseControlNetModel",
|
||||
"StableAudioDiTModel",
|
||||
"StableCascadeUNet",
|
||||
"T2IAdapter",
|
||||
"T5FilmDecoder",
|
||||
"Transformer2DModel",
|
||||
"UNet1DModel",
|
||||
"UNet2DConditionModel",
|
||||
"UNet2DModel",
|
||||
"UNet3DConditionModel",
|
||||
"UNetControlNetXSModel",
|
||||
"UNetMotionModel",
|
||||
"UNetSpatioTemporalConditionModel",
|
||||
"UVit2DModel",
|
||||
"VQModel",
|
||||
]
|
||||
)
|
||||
_import_structure["optimization"] = [
|
||||
"get_constant_schedule",
|
||||
"get_constant_schedule_with_warmup",
|
||||
"get_cosine_schedule_with_warmup",
|
||||
"get_cosine_with_hard_restarts_schedule_with_warmup",
|
||||
"get_linear_schedule_with_warmup",
|
||||
"get_polynomial_decay_schedule_with_warmup",
|
||||
"get_scheduler",
|
||||
]
|
||||
_import_structure["pipelines"].extend(
|
||||
[
|
||||
"AudioPipelineOutput",
|
||||
"AutoPipelineForImage2Image",
|
||||
"AutoPipelineForInpainting",
|
||||
"AutoPipelineForText2Image",
|
||||
"ConsistencyModelPipeline",
|
||||
"DanceDiffusionPipeline",
|
||||
"DDIMPipeline",
|
||||
"DDPMPipeline",
|
||||
"DiffusionPipeline",
|
||||
"DiTPipeline",
|
||||
"ImagePipelineOutput",
|
||||
"KarrasVePipeline",
|
||||
"LDMPipeline",
|
||||
"LDMSuperResolutionPipeline",
|
||||
"PNDMPipeline",
|
||||
"RePaintPipeline",
|
||||
"ScoreSdeVePipeline",
|
||||
"StableDiffusionMixin",
|
||||
]
|
||||
)
|
||||
_import_structure["quantizers"] = ["DiffusersQuantizer"]
|
||||
_import_structure["schedulers"].extend(
|
||||
[
|
||||
"AmusedScheduler",
|
||||
"CMStochasticIterativeScheduler",
|
||||
"CogVideoXDDIMScheduler",
|
||||
"CogVideoXDPMScheduler",
|
||||
"DDIMInverseScheduler",
|
||||
"DDIMParallelScheduler",
|
||||
"DDIMScheduler",
|
||||
"DDPMParallelScheduler",
|
||||
"DDPMScheduler",
|
||||
"DDPMWuerstchenScheduler",
|
||||
"DEISMultistepScheduler",
|
||||
"DPMSolverMultistepInverseScheduler",
|
||||
"DPMSolverMultistepScheduler",
|
||||
"DPMSolverSinglestepScheduler",
|
||||
"EDMDPMSolverMultistepScheduler",
|
||||
"EDMEulerScheduler",
|
||||
"EulerAncestralDiscreteScheduler",
|
||||
"EulerDiscreteScheduler",
|
||||
"FlowMatchEulerDiscreteScheduler",
|
||||
"FlowMatchHeunDiscreteScheduler",
|
||||
"HeunDiscreteScheduler",
|
||||
"IPNDMScheduler",
|
||||
"KarrasVeScheduler",
|
||||
"KDPM2AncestralDiscreteScheduler",
|
||||
"KDPM2DiscreteScheduler",
|
||||
"LCMScheduler",
|
||||
"PNDMScheduler",
|
||||
"RePaintScheduler",
|
||||
"SASolverScheduler",
|
||||
"SchedulerMixin",
|
||||
"ScoreSdeVeScheduler",
|
||||
"TCDScheduler",
|
||||
"UnCLIPScheduler",
|
||||
"UniPCMultistepScheduler",
|
||||
"VQDiffusionScheduler",
|
||||
]
|
||||
)
|
||||
_import_structure["training_utils"] = ["EMAModel"]
|
||||
|
||||
try:
|
||||
if not (is_torch_available() and is_scipy_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils import dummy_torch_and_scipy_objects # noqa F403
|
||||
|
||||
_import_structure["utils.dummy_torch_and_scipy_objects"] = [
|
||||
name for name in dir(dummy_torch_and_scipy_objects) if not name.startswith("_")
|
||||
]
|
||||
|
||||
else:
|
||||
_import_structure["schedulers"].extend(["LMSDiscreteScheduler"])
|
||||
|
||||
try:
|
||||
if not (is_torch_available() and is_torchsde_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils import dummy_torch_and_torchsde_objects # noqa F403
|
||||
|
||||
_import_structure["utils.dummy_torch_and_torchsde_objects"] = [
|
||||
name for name in dir(dummy_torch_and_torchsde_objects) if not name.startswith("_")
|
||||
]
|
||||
|
||||
else:
|
||||
_import_structure["schedulers"].extend(["CosineDPMSolverMultistepScheduler", "DPMSolverSDEScheduler"])
|
||||
|
||||
try:
|
||||
if not (is_torch_available() and is_transformers_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils import dummy_torch_and_transformers_objects # noqa F403
|
||||
|
||||
_import_structure["utils.dummy_torch_and_transformers_objects"] = [
|
||||
name for name in dir(dummy_torch_and_transformers_objects) if not name.startswith("_")
|
||||
]
|
||||
|
||||
else:
|
||||
_import_structure["pipelines"].extend(
|
||||
[
|
||||
"AltDiffusionImg2ImgPipeline",
|
||||
"AltDiffusionPipeline",
|
||||
"AmusedImg2ImgPipeline",
|
||||
"AmusedInpaintPipeline",
|
||||
"AmusedPipeline",
|
||||
"AnimateDiffControlNetPipeline",
|
||||
"AnimateDiffPAGPipeline",
|
||||
"AnimateDiffPipeline",
|
||||
"AnimateDiffSDXLPipeline",
|
||||
"AnimateDiffSparseControlNetPipeline",
|
||||
"AnimateDiffVideoToVideoControlNetPipeline",
|
||||
"AnimateDiffVideoToVideoPipeline",
|
||||
"AudioLDM2Pipeline",
|
||||
"AudioLDM2ProjectionModel",
|
||||
"AudioLDM2UNet2DConditionModel",
|
||||
"AudioLDMPipeline",
|
||||
"AuraFlowPipeline",
|
||||
"BlipDiffusionControlNetPipeline",
|
||||
"BlipDiffusionPipeline",
|
||||
"CLIPImageProjection",
|
||||
"CogVideoXFunControlPipeline",
|
||||
"CogVideoXImageToVideoPipeline",
|
||||
"CogVideoXPipeline",
|
||||
"CogVideoXVideoToVideoPipeline",
|
||||
"CogView3PlusPipeline",
|
||||
"CycleDiffusionPipeline",
|
||||
"FluxControlNetImg2ImgPipeline",
|
||||
"FluxControlNetInpaintPipeline",
|
||||
"FluxControlNetPipeline",
|
||||
"FluxImg2ImgPipeline",
|
||||
"FluxInpaintPipeline",
|
||||
"FluxPipeline",
|
||||
"HunyuanDiTControlNetPipeline",
|
||||
"HunyuanDiTPAGPipeline",
|
||||
"HunyuanDiTPipeline",
|
||||
"I2VGenXLPipeline",
|
||||
"IFImg2ImgPipeline",
|
||||
"IFImg2ImgSuperResolutionPipeline",
|
||||
"IFInpaintingPipeline",
|
||||
"IFInpaintingSuperResolutionPipeline",
|
||||
"IFPipeline",
|
||||
"IFSuperResolutionPipeline",
|
||||
"ImageTextPipelineOutput",
|
||||
"Kandinsky3Img2ImgPipeline",
|
||||
"Kandinsky3Pipeline",
|
||||
"KandinskyCombinedPipeline",
|
||||
"KandinskyImg2ImgCombinedPipeline",
|
||||
"KandinskyImg2ImgPipeline",
|
||||
"KandinskyInpaintCombinedPipeline",
|
||||
"KandinskyInpaintPipeline",
|
||||
"KandinskyPipeline",
|
||||
"KandinskyPriorPipeline",
|
||||
"KandinskyV22CombinedPipeline",
|
||||
"KandinskyV22ControlnetImg2ImgPipeline",
|
||||
"KandinskyV22ControlnetPipeline",
|
||||
"KandinskyV22Img2ImgCombinedPipeline",
|
||||
"KandinskyV22Img2ImgPipeline",
|
||||
"KandinskyV22InpaintCombinedPipeline",
|
||||
"KandinskyV22InpaintPipeline",
|
||||
"KandinskyV22Pipeline",
|
||||
"KandinskyV22PriorEmb2EmbPipeline",
|
||||
"KandinskyV22PriorPipeline",
|
||||
"LatentConsistencyModelImg2ImgPipeline",
|
||||
"LatentConsistencyModelPipeline",
|
||||
"LattePipeline",
|
||||
"LDMTextToImagePipeline",
|
||||
"LEditsPPPipelineStableDiffusion",
|
||||
"LEditsPPPipelineStableDiffusionXL",
|
||||
"LuminaText2ImgPipeline",
|
||||
"MarigoldDepthPipeline",
|
||||
"MarigoldNormalsPipeline",
|
||||
"MusicLDMPipeline",
|
||||
"PaintByExamplePipeline",
|
||||
"PIAPipeline",
|
||||
"PixArtAlphaPipeline",
|
||||
"PixArtSigmaPAGPipeline",
|
||||
"PixArtSigmaPipeline",
|
||||
"SemanticStableDiffusionPipeline",
|
||||
"ShapEImg2ImgPipeline",
|
||||
"ShapEPipeline",
|
||||
"StableAudioPipeline",
|
||||
"StableAudioProjectionModel",
|
||||
"StableCascadeCombinedPipeline",
|
||||
"StableCascadeDecoderPipeline",
|
||||
"StableCascadePriorPipeline",
|
||||
"StableDiffusion3ControlNetInpaintingPipeline",
|
||||
"StableDiffusion3ControlNetPipeline",
|
||||
"StableDiffusion3Img2ImgPipeline",
|
||||
"StableDiffusion3InpaintPipeline",
|
||||
"StableDiffusion3PAGPipeline",
|
||||
"StableDiffusion3Pipeline",
|
||||
"StableDiffusion3InvEnhancePipeline",
|
||||
"StableDiffusionAdapterPipeline",
|
||||
"StableDiffusionAttendAndExcitePipeline",
|
||||
"StableDiffusionControlNetImg2ImgPipeline",
|
||||
"StableDiffusionControlNetInpaintPipeline",
|
||||
"StableDiffusionControlNetPAGInpaintPipeline",
|
||||
"StableDiffusionControlNetPAGPipeline",
|
||||
"StableDiffusionControlNetPipeline",
|
||||
"StableDiffusionControlNetXSPipeline",
|
||||
"StableDiffusionDepth2ImgPipeline",
|
||||
"StableDiffusionDiffEditPipeline",
|
||||
"StableDiffusionGLIGENPipeline",
|
||||
"StableDiffusionGLIGENTextImagePipeline",
|
||||
"StableDiffusionImageVariationPipeline",
|
||||
"StableDiffusionImg2ImgPipeline",
|
||||
"StableDiffusionInpaintPipeline",
|
||||
"StableDiffusionInpaintPipelineLegacy",
|
||||
"StableDiffusionInstructPix2PixPipeline",
|
||||
"StableDiffusionLatentUpscalePipeline",
|
||||
"StableDiffusionLDM3DPipeline",
|
||||
"StableDiffusionModelEditingPipeline",
|
||||
"StableDiffusionPAGImg2ImgPipeline",
|
||||
"StableDiffusionPAGPipeline",
|
||||
"StableDiffusionPanoramaPipeline",
|
||||
"StableDiffusionParadigmsPipeline",
|
||||
"StableDiffusionPipeline",
|
||||
"StableDiffusionInvEnhancePipeline",
|
||||
"StableDiffusionPipelineSafe",
|
||||
"StableDiffusionPix2PixZeroPipeline",
|
||||
"StableDiffusionSAGPipeline",
|
||||
"StableDiffusionUpscalePipeline",
|
||||
"StableDiffusionXLAdapterPipeline",
|
||||
"StableDiffusionXLControlNetImg2ImgPipeline",
|
||||
"StableDiffusionXLControlNetInpaintPipeline",
|
||||
"StableDiffusionXLControlNetPAGImg2ImgPipeline",
|
||||
"StableDiffusionXLControlNetPAGPipeline",
|
||||
"StableDiffusionXLControlNetPipeline",
|
||||
"StableDiffusionXLControlNetXSPipeline",
|
||||
"StableDiffusionXLImg2ImgPipeline",
|
||||
"StableDiffusionXLInpaintPipeline",
|
||||
"StableDiffusionXLInstructPix2PixPipeline",
|
||||
"StableDiffusionXLPAGImg2ImgPipeline",
|
||||
"StableDiffusionXLPAGInpaintPipeline",
|
||||
"StableDiffusionXLPAGPipeline",
|
||||
"StableDiffusionXLPipeline",
|
||||
"StableDiffusionXLInvEnhancePipeline",
|
||||
"StableUnCLIPImg2ImgPipeline",
|
||||
"StableUnCLIPPipeline",
|
||||
"StableVideoDiffusionPipeline",
|
||||
"TextToVideoSDPipeline",
|
||||
"TextToVideoZeroPipeline",
|
||||
"TextToVideoZeroSDXLPipeline",
|
||||
"UnCLIPImageVariationPipeline",
|
||||
"UnCLIPPipeline",
|
||||
"UniDiffuserModel",
|
||||
"UniDiffuserPipeline",
|
||||
"UniDiffuserTextDecoder",
|
||||
"VersatileDiffusionDualGuidedPipeline",
|
||||
"VersatileDiffusionImageVariationPipeline",
|
||||
"VersatileDiffusionPipeline",
|
||||
"VersatileDiffusionTextToImagePipeline",
|
||||
"VideoToVideoSDPipeline",
|
||||
"VQDiffusionPipeline",
|
||||
"WuerstchenCombinedPipeline",
|
||||
"WuerstchenDecoderPipeline",
|
||||
"WuerstchenPriorPipeline",
|
||||
]
|
||||
)
|
||||
|
||||
try:
|
||||
if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils import dummy_torch_and_transformers_and_k_diffusion_objects # noqa F403
|
||||
|
||||
_import_structure["utils.dummy_torch_and_transformers_and_k_diffusion_objects"] = [
|
||||
name for name in dir(dummy_torch_and_transformers_and_k_diffusion_objects) if not name.startswith("_")
|
||||
]
|
||||
|
||||
else:
|
||||
_import_structure["pipelines"].extend(["StableDiffusionKDiffusionPipeline", "StableDiffusionXLKDiffusionPipeline"])
|
||||
|
||||
try:
|
||||
if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils import dummy_torch_and_transformers_and_sentencepiece_objects # noqa F403
|
||||
|
||||
_import_structure["utils.dummy_torch_and_transformers_and_sentencepiece_objects"] = [
|
||||
name for name in dir(dummy_torch_and_transformers_and_sentencepiece_objects) if not name.startswith("_")
|
||||
]
|
||||
|
||||
else:
|
||||
_import_structure["pipelines"].extend(["KolorsImg2ImgPipeline", "KolorsPAGPipeline", "KolorsPipeline"])
|
||||
|
||||
try:
|
||||
if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils import dummy_torch_and_transformers_and_onnx_objects # noqa F403
|
||||
|
||||
_import_structure["utils.dummy_torch_and_transformers_and_onnx_objects"] = [
|
||||
name for name in dir(dummy_torch_and_transformers_and_onnx_objects) if not name.startswith("_")
|
||||
]
|
||||
|
||||
else:
|
||||
_import_structure["pipelines"].extend(
|
||||
[
|
||||
"OnnxStableDiffusionImg2ImgPipeline",
|
||||
"OnnxStableDiffusionInpaintPipeline",
|
||||
"OnnxStableDiffusionInpaintPipelineLegacy",
|
||||
"OnnxStableDiffusionPipeline",
|
||||
"OnnxStableDiffusionUpscalePipeline",
|
||||
"StableDiffusionOnnxPipeline",
|
||||
]
|
||||
)
|
||||
|
||||
try:
|
||||
if not (is_torch_available() and is_librosa_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils import dummy_torch_and_librosa_objects # noqa F403
|
||||
|
||||
_import_structure["utils.dummy_torch_and_librosa_objects"] = [
|
||||
name for name in dir(dummy_torch_and_librosa_objects) if not name.startswith("_")
|
||||
]
|
||||
|
||||
else:
|
||||
_import_structure["pipelines"].extend(["AudioDiffusionPipeline", "Mel"])
|
||||
|
||||
try:
|
||||
if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils import dummy_transformers_and_torch_and_note_seq_objects # noqa F403
|
||||
|
||||
_import_structure["utils.dummy_transformers_and_torch_and_note_seq_objects"] = [
|
||||
name for name in dir(dummy_transformers_and_torch_and_note_seq_objects) if not name.startswith("_")
|
||||
]
|
||||
|
||||
|
||||
else:
|
||||
_import_structure["pipelines"].extend(["SpectrogramDiffusionPipeline"])
|
||||
|
||||
try:
|
||||
if not is_flax_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils import dummy_flax_objects # noqa F403
|
||||
|
||||
_import_structure["utils.dummy_flax_objects"] = [
|
||||
name for name in dir(dummy_flax_objects) if not name.startswith("_")
|
||||
]
|
||||
|
||||
|
||||
else:
|
||||
_import_structure["models.controlnet_flax"] = ["FlaxControlNetModel"]
|
||||
_import_structure["models.modeling_flax_utils"] = ["FlaxModelMixin"]
|
||||
_import_structure["models.unets.unet_2d_condition_flax"] = ["FlaxUNet2DConditionModel"]
|
||||
_import_structure["models.vae_flax"] = ["FlaxAutoencoderKL"]
|
||||
_import_structure["pipelines"].extend(["FlaxDiffusionPipeline"])
|
||||
_import_structure["schedulers"].extend(
|
||||
[
|
||||
"FlaxDDIMScheduler",
|
||||
"FlaxDDPMScheduler",
|
||||
"FlaxDPMSolverMultistepScheduler",
|
||||
"FlaxEulerDiscreteScheduler",
|
||||
"FlaxKarrasVeScheduler",
|
||||
"FlaxLMSDiscreteScheduler",
|
||||
"FlaxPNDMScheduler",
|
||||
"FlaxSchedulerMixin",
|
||||
"FlaxScoreSdeVeScheduler",
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
try:
|
||||
if not (is_flax_available() and is_transformers_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils import dummy_flax_and_transformers_objects # noqa F403
|
||||
|
||||
_import_structure["utils.dummy_flax_and_transformers_objects"] = [
|
||||
name for name in dir(dummy_flax_and_transformers_objects) if not name.startswith("_")
|
||||
]
|
||||
|
||||
|
||||
else:
|
||||
_import_structure["pipelines"].extend(
|
||||
[
|
||||
"FlaxStableDiffusionControlNetPipeline",
|
||||
"FlaxStableDiffusionImg2ImgPipeline",
|
||||
"FlaxStableDiffusionInpaintPipeline",
|
||||
"FlaxStableDiffusionPipeline",
|
||||
"FlaxStableDiffusionXLPipeline",
|
||||
]
|
||||
)
|
||||
|
||||
try:
|
||||
if not (is_note_seq_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils import dummy_note_seq_objects # noqa F403
|
||||
|
||||
_import_structure["utils.dummy_note_seq_objects"] = [
|
||||
name for name in dir(dummy_note_seq_objects) if not name.startswith("_")
|
||||
]
|
||||
|
||||
|
||||
else:
|
||||
_import_structure["pipelines"].extend(["MidiProcessor"])
|
||||
|
||||
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
|
||||
from .configuration_utils import ConfigMixin
|
||||
from .quantizers.quantization_config import BitsAndBytesConfig
|
||||
|
||||
try:
|
||||
if not is_onnx_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils.dummy_onnx_objects import * # noqa F403
|
||||
else:
|
||||
from .pipelines import OnnxRuntimeModel
|
||||
|
||||
try:
|
||||
if not is_torch_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils.dummy_pt_objects import * # noqa F403
|
||||
else:
|
||||
from .models import (
|
||||
AsymmetricAutoencoderKL,
|
||||
AuraFlowTransformer2DModel,
|
||||
AutoencoderKL,
|
||||
AutoencoderKLCogVideoX,
|
||||
AutoencoderKLTemporalDecoder,
|
||||
AutoencoderOobleck,
|
||||
AutoencoderTiny,
|
||||
CogVideoXTransformer3DModel,
|
||||
CogView3PlusTransformer2DModel,
|
||||
ConsistencyDecoderVAE,
|
||||
ControlNetModel,
|
||||
ControlNetXSAdapter,
|
||||
DiTTransformer2DModel,
|
||||
FluxControlNetModel,
|
||||
FluxMultiControlNetModel,
|
||||
FluxTransformer2DModel,
|
||||
HunyuanDiT2DControlNetModel,
|
||||
HunyuanDiT2DModel,
|
||||
HunyuanDiT2DMultiControlNetModel,
|
||||
I2VGenXLUNet,
|
||||
Kandinsky3UNet,
|
||||
LatteTransformer3DModel,
|
||||
LuminaNextDiT2DModel,
|
||||
ModelMixin,
|
||||
MotionAdapter,
|
||||
MultiAdapter,
|
||||
PixArtTransformer2DModel,
|
||||
PriorTransformer,
|
||||
SD3ControlNetModel,
|
||||
SD3MultiControlNetModel,
|
||||
SD3Transformer2DModel,
|
||||
SparseControlNetModel,
|
||||
StableAudioDiTModel,
|
||||
T2IAdapter,
|
||||
T5FilmDecoder,
|
||||
Transformer2DModel,
|
||||
UNet1DModel,
|
||||
UNet2DConditionModel,
|
||||
UNet2DModel,
|
||||
UNet3DConditionModel,
|
||||
UNetControlNetXSModel,
|
||||
UNetMotionModel,
|
||||
UNetSpatioTemporalConditionModel,
|
||||
UVit2DModel,
|
||||
VQModel,
|
||||
)
|
||||
from .optimization import (
|
||||
get_constant_schedule,
|
||||
get_constant_schedule_with_warmup,
|
||||
get_cosine_schedule_with_warmup,
|
||||
get_cosine_with_hard_restarts_schedule_with_warmup,
|
||||
get_linear_schedule_with_warmup,
|
||||
get_polynomial_decay_schedule_with_warmup,
|
||||
get_scheduler,
|
||||
)
|
||||
from .pipelines import (
|
||||
AudioPipelineOutput,
|
||||
AutoPipelineForImage2Image,
|
||||
AutoPipelineForInpainting,
|
||||
AutoPipelineForText2Image,
|
||||
BlipDiffusionControlNetPipeline,
|
||||
BlipDiffusionPipeline,
|
||||
CLIPImageProjection,
|
||||
ConsistencyModelPipeline,
|
||||
DanceDiffusionPipeline,
|
||||
DDIMPipeline,
|
||||
DDPMPipeline,
|
||||
DiffusionPipeline,
|
||||
DiTPipeline,
|
||||
ImagePipelineOutput,
|
||||
KarrasVePipeline,
|
||||
LDMPipeline,
|
||||
LDMSuperResolutionPipeline,
|
||||
PNDMPipeline,
|
||||
RePaintPipeline,
|
||||
ScoreSdeVePipeline,
|
||||
StableDiffusionMixin,
|
||||
)
|
||||
from .quantizers import DiffusersQuantizer
|
||||
from .schedulers import (
|
||||
AmusedScheduler,
|
||||
CMStochasticIterativeScheduler,
|
||||
CogVideoXDDIMScheduler,
|
||||
CogVideoXDPMScheduler,
|
||||
DDIMInverseScheduler,
|
||||
DDIMParallelScheduler,
|
||||
DDIMScheduler,
|
||||
DDPMParallelScheduler,
|
||||
DDPMScheduler,
|
||||
DDPMWuerstchenScheduler,
|
||||
DEISMultistepScheduler,
|
||||
DPMSolverMultistepInverseScheduler,
|
||||
DPMSolverMultistepScheduler,
|
||||
DPMSolverSinglestepScheduler,
|
||||
EDMDPMSolverMultistepScheduler,
|
||||
EDMEulerScheduler,
|
||||
EulerAncestralDiscreteScheduler,
|
||||
EulerDiscreteScheduler,
|
||||
FlowMatchEulerDiscreteScheduler,
|
||||
FlowMatchHeunDiscreteScheduler,
|
||||
HeunDiscreteScheduler,
|
||||
IPNDMScheduler,
|
||||
KarrasVeScheduler,
|
||||
KDPM2AncestralDiscreteScheduler,
|
||||
KDPM2DiscreteScheduler,
|
||||
LCMScheduler,
|
||||
PNDMScheduler,
|
||||
RePaintScheduler,
|
||||
SASolverScheduler,
|
||||
SchedulerMixin,
|
||||
ScoreSdeVeScheduler,
|
||||
TCDScheduler,
|
||||
UnCLIPScheduler,
|
||||
UniPCMultistepScheduler,
|
||||
VQDiffusionScheduler,
|
||||
)
|
||||
from .training_utils import EMAModel
|
||||
|
||||
try:
|
||||
if not (is_torch_available() and is_scipy_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils.dummy_torch_and_scipy_objects import * # noqa F403
|
||||
else:
|
||||
from .schedulers import LMSDiscreteScheduler
|
||||
|
||||
try:
|
||||
if not (is_torch_available() and is_torchsde_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils.dummy_torch_and_torchsde_objects import * # noqa F403
|
||||
else:
|
||||
from .schedulers import CosineDPMSolverMultistepScheduler, DPMSolverSDEScheduler
|
||||
|
||||
try:
|
||||
if not (is_torch_available() and is_transformers_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils.dummy_torch_and_transformers_objects import * # noqa F403
|
||||
else:
|
||||
from .pipelines import (
|
||||
AltDiffusionImg2ImgPipeline,
|
||||
AltDiffusionPipeline,
|
||||
AmusedImg2ImgPipeline,
|
||||
AmusedInpaintPipeline,
|
||||
AmusedPipeline,
|
||||
AnimateDiffControlNetPipeline,
|
||||
AnimateDiffPAGPipeline,
|
||||
AnimateDiffPipeline,
|
||||
AnimateDiffSDXLPipeline,
|
||||
AnimateDiffSparseControlNetPipeline,
|
||||
AnimateDiffVideoToVideoControlNetPipeline,
|
||||
AnimateDiffVideoToVideoPipeline,
|
||||
AudioLDM2Pipeline,
|
||||
AudioLDM2ProjectionModel,
|
||||
AudioLDM2UNet2DConditionModel,
|
||||
AudioLDMPipeline,
|
||||
AuraFlowPipeline,
|
||||
CLIPImageProjection,
|
||||
CogVideoXFunControlPipeline,
|
||||
CogVideoXImageToVideoPipeline,
|
||||
CogVideoXPipeline,
|
||||
CogVideoXVideoToVideoPipeline,
|
||||
CogView3PlusPipeline,
|
||||
CycleDiffusionPipeline,
|
||||
FluxControlNetImg2ImgPipeline,
|
||||
FluxControlNetInpaintPipeline,
|
||||
FluxControlNetPipeline,
|
||||
FluxImg2ImgPipeline,
|
||||
FluxInpaintPipeline,
|
||||
FluxPipeline,
|
||||
HunyuanDiTControlNetPipeline,
|
||||
HunyuanDiTPAGPipeline,
|
||||
HunyuanDiTPipeline,
|
||||
I2VGenXLPipeline,
|
||||
IFImg2ImgPipeline,
|
||||
IFImg2ImgSuperResolutionPipeline,
|
||||
IFInpaintingPipeline,
|
||||
IFInpaintingSuperResolutionPipeline,
|
||||
IFPipeline,
|
||||
IFSuperResolutionPipeline,
|
||||
ImageTextPipelineOutput,
|
||||
Kandinsky3Img2ImgPipeline,
|
||||
Kandinsky3Pipeline,
|
||||
KandinskyCombinedPipeline,
|
||||
KandinskyImg2ImgCombinedPipeline,
|
||||
KandinskyImg2ImgPipeline,
|
||||
KandinskyInpaintCombinedPipeline,
|
||||
KandinskyInpaintPipeline,
|
||||
KandinskyPipeline,
|
||||
KandinskyPriorPipeline,
|
||||
KandinskyV22CombinedPipeline,
|
||||
KandinskyV22ControlnetImg2ImgPipeline,
|
||||
KandinskyV22ControlnetPipeline,
|
||||
KandinskyV22Img2ImgCombinedPipeline,
|
||||
KandinskyV22Img2ImgPipeline,
|
||||
KandinskyV22InpaintCombinedPipeline,
|
||||
KandinskyV22InpaintPipeline,
|
||||
KandinskyV22Pipeline,
|
||||
KandinskyV22PriorEmb2EmbPipeline,
|
||||
KandinskyV22PriorPipeline,
|
||||
LatentConsistencyModelImg2ImgPipeline,
|
||||
LatentConsistencyModelPipeline,
|
||||
LattePipeline,
|
||||
LDMTextToImagePipeline,
|
||||
LEditsPPPipelineStableDiffusion,
|
||||
LEditsPPPipelineStableDiffusionXL,
|
||||
LuminaText2ImgPipeline,
|
||||
MarigoldDepthPipeline,
|
||||
MarigoldNormalsPipeline,
|
||||
MusicLDMPipeline,
|
||||
PaintByExamplePipeline,
|
||||
PIAPipeline,
|
||||
PixArtAlphaPipeline,
|
||||
PixArtSigmaPAGPipeline,
|
||||
PixArtSigmaPipeline,
|
||||
SemanticStableDiffusionPipeline,
|
||||
ShapEImg2ImgPipeline,
|
||||
ShapEPipeline,
|
||||
StableAudioPipeline,
|
||||
StableAudioProjectionModel,
|
||||
StableCascadeCombinedPipeline,
|
||||
StableCascadeDecoderPipeline,
|
||||
StableCascadePriorPipeline,
|
||||
StableDiffusion3ControlNetPipeline,
|
||||
StableDiffusion3Img2ImgPipeline,
|
||||
StableDiffusion3InpaintPipeline,
|
||||
StableDiffusion3PAGPipeline,
|
||||
StableDiffusion3Pipeline,
|
||||
StableDiffusion3InvEnhancePipeline,
|
||||
StableDiffusionAdapterPipeline,
|
||||
StableDiffusionAttendAndExcitePipeline,
|
||||
StableDiffusionControlNetImg2ImgPipeline,
|
||||
StableDiffusionControlNetInpaintPipeline,
|
||||
StableDiffusionControlNetPAGInpaintPipeline,
|
||||
StableDiffusionControlNetPAGPipeline,
|
||||
StableDiffusionControlNetPipeline,
|
||||
StableDiffusionControlNetXSPipeline,
|
||||
StableDiffusionDepth2ImgPipeline,
|
||||
StableDiffusionDiffEditPipeline,
|
||||
StableDiffusionGLIGENPipeline,
|
||||
StableDiffusionGLIGENTextImagePipeline,
|
||||
StableDiffusionImageVariationPipeline,
|
||||
StableDiffusionImg2ImgPipeline,
|
||||
StableDiffusionInpaintPipeline,
|
||||
StableDiffusionInpaintPipelineLegacy,
|
||||
StableDiffusionInstructPix2PixPipeline,
|
||||
StableDiffusionLatentUpscalePipeline,
|
||||
StableDiffusionLDM3DPipeline,
|
||||
StableDiffusionModelEditingPipeline,
|
||||
StableDiffusionPAGImg2ImgPipeline,
|
||||
StableDiffusionPAGPipeline,
|
||||
StableDiffusionPanoramaPipeline,
|
||||
StableDiffusionParadigmsPipeline,
|
||||
StableDiffusionPipeline,
|
||||
StableDiffusionInvEnhancePipeline,
|
||||
StableDiffusionPipelineSafe,
|
||||
StableDiffusionPix2PixZeroPipeline,
|
||||
StableDiffusionSAGPipeline,
|
||||
StableDiffusionUpscalePipeline,
|
||||
StableDiffusionXLAdapterPipeline,
|
||||
StableDiffusionXLControlNetImg2ImgPipeline,
|
||||
StableDiffusionXLControlNetInpaintPipeline,
|
||||
StableDiffusionXLControlNetPAGImg2ImgPipeline,
|
||||
StableDiffusionXLControlNetPAGPipeline,
|
||||
StableDiffusionXLControlNetPipeline,
|
||||
StableDiffusionXLControlNetXSPipeline,
|
||||
StableDiffusionXLImg2ImgPipeline,
|
||||
StableDiffusionXLInpaintPipeline,
|
||||
StableDiffusionXLInstructPix2PixPipeline,
|
||||
StableDiffusionXLPAGImg2ImgPipeline,
|
||||
StableDiffusionXLPAGInpaintPipeline,
|
||||
StableDiffusionXLPAGPipeline,
|
||||
StableDiffusionXLPipeline,
|
||||
StableDiffusionXLInvEnhancePipeline,
|
||||
StableUnCLIPImg2ImgPipeline,
|
||||
StableUnCLIPPipeline,
|
||||
StableVideoDiffusionPipeline,
|
||||
TextToVideoSDPipeline,
|
||||
TextToVideoZeroPipeline,
|
||||
TextToVideoZeroSDXLPipeline,
|
||||
UnCLIPImageVariationPipeline,
|
||||
UnCLIPPipeline,
|
||||
UniDiffuserModel,
|
||||
UniDiffuserPipeline,
|
||||
UniDiffuserTextDecoder,
|
||||
VersatileDiffusionDualGuidedPipeline,
|
||||
VersatileDiffusionImageVariationPipeline,
|
||||
VersatileDiffusionPipeline,
|
||||
VersatileDiffusionTextToImagePipeline,
|
||||
VideoToVideoSDPipeline,
|
||||
VQDiffusionPipeline,
|
||||
WuerstchenCombinedPipeline,
|
||||
WuerstchenDecoderPipeline,
|
||||
WuerstchenPriorPipeline,
|
||||
)
|
||||
|
||||
try:
|
||||
if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import * # noqa F403
|
||||
else:
|
||||
from .pipelines import StableDiffusionKDiffusionPipeline, StableDiffusionXLKDiffusionPipeline
|
||||
|
||||
try:
|
||||
if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils.dummy_torch_and_transformers_and_sentencepiece_objects import * # noqa F403
|
||||
else:
|
||||
from .pipelines import KolorsImg2ImgPipeline, KolorsPAGPipeline, KolorsPipeline
|
||||
try:
|
||||
if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils.dummy_torch_and_transformers_and_onnx_objects import * # noqa F403
|
||||
else:
|
||||
from .pipelines import (
|
||||
OnnxStableDiffusionImg2ImgPipeline,
|
||||
OnnxStableDiffusionInpaintPipeline,
|
||||
OnnxStableDiffusionInpaintPipelineLegacy,
|
||||
OnnxStableDiffusionPipeline,
|
||||
OnnxStableDiffusionUpscalePipeline,
|
||||
StableDiffusionOnnxPipeline,
|
||||
)
|
||||
|
||||
try:
|
||||
if not (is_torch_available() and is_librosa_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils.dummy_torch_and_librosa_objects import * # noqa F403
|
||||
else:
|
||||
from .pipelines import AudioDiffusionPipeline, Mel
|
||||
|
||||
try:
|
||||
if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils.dummy_transformers_and_torch_and_note_seq_objects import * # noqa F403
|
||||
else:
|
||||
from .pipelines import SpectrogramDiffusionPipeline
|
||||
|
||||
try:
|
||||
if not is_flax_available():
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils.dummy_flax_objects import * # noqa F403
|
||||
else:
|
||||
from .models.controlnet_flax import FlaxControlNetModel
|
||||
from .models.modeling_flax_utils import FlaxModelMixin
|
||||
from .models.unets.unet_2d_condition_flax import FlaxUNet2DConditionModel
|
||||
from .models.vae_flax import FlaxAutoencoderKL
|
||||
from .pipelines import FlaxDiffusionPipeline
|
||||
from .schedulers import (
|
||||
FlaxDDIMScheduler,
|
||||
FlaxDDPMScheduler,
|
||||
FlaxDPMSolverMultistepScheduler,
|
||||
FlaxEulerDiscreteScheduler,
|
||||
FlaxKarrasVeScheduler,
|
||||
FlaxLMSDiscreteScheduler,
|
||||
FlaxPNDMScheduler,
|
||||
FlaxSchedulerMixin,
|
||||
FlaxScoreSdeVeScheduler,
|
||||
)
|
||||
|
||||
try:
|
||||
if not (is_flax_available() and is_transformers_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils.dummy_flax_and_transformers_objects import * # noqa F403
|
||||
else:
|
||||
from .pipelines import (
|
||||
FlaxStableDiffusionControlNetPipeline,
|
||||
FlaxStableDiffusionImg2ImgPipeline,
|
||||
FlaxStableDiffusionInpaintPipeline,
|
||||
FlaxStableDiffusionPipeline,
|
||||
FlaxStableDiffusionXLPipeline,
|
||||
)
|
||||
|
||||
try:
|
||||
if not (is_note_seq_available()):
|
||||
raise OptionalDependencyNotAvailable()
|
||||
except OptionalDependencyNotAvailable:
|
||||
from .utils.dummy_note_seq_objects import * # noqa F403
|
||||
else:
|
||||
from .pipelines import MidiProcessor
|
||||
|
||||
else:
|
||||
import sys
|
||||
|
||||
sys.modules[__name__] = _LazyModule(
|
||||
__name__,
|
||||
globals()["__file__"],
|
||||
_import_structure,
|
||||
module_spec=__spec__,
|
||||
extra_objects={"__version__": __version__},
|
||||
)
|
||||
@@ -1,156 +0,0 @@
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from .configuration_utils import ConfigMixin, register_to_config
|
||||
from .utils import CONFIG_NAME
|
||||
|
||||
|
||||
class PipelineCallback(ConfigMixin):
|
||||
"""
|
||||
Base class for all the official callbacks used in a pipeline. This class provides a structure for implementing
|
||||
custom callbacks and ensures that all callbacks have a consistent interface.
|
||||
|
||||
Please implement the following:
|
||||
`tensor_inputs`: This should return a list of tensor inputs specific to your callback. You will only be able to
|
||||
include
|
||||
variables listed in the `._callback_tensor_inputs` attribute of your pipeline class.
|
||||
`callback_fn`: This method defines the core functionality of your callback.
|
||||
"""
|
||||
|
||||
config_name = CONFIG_NAME
|
||||
|
||||
@register_to_config
|
||||
def __init__(self, cutoff_step_ratio=1.0, cutoff_step_index=None):
|
||||
super().__init__()
|
||||
|
||||
if (cutoff_step_ratio is None and cutoff_step_index is None) or (
|
||||
cutoff_step_ratio is not None and cutoff_step_index is not None
|
||||
):
|
||||
raise ValueError("Either cutoff_step_ratio or cutoff_step_index should be provided, not both or none.")
|
||||
|
||||
if cutoff_step_ratio is not None and (
|
||||
not isinstance(cutoff_step_ratio, float) or not (0.0 <= cutoff_step_ratio <= 1.0)
|
||||
):
|
||||
raise ValueError("cutoff_step_ratio must be a float between 0.0 and 1.0.")
|
||||
|
||||
@property
|
||||
def tensor_inputs(self) -> List[str]:
|
||||
raise NotImplementedError(f"You need to set the attribute `tensor_inputs` for {self.__class__}")
|
||||
|
||||
def callback_fn(self, pipeline, step_index, timesteps, callback_kwargs) -> Dict[str, Any]:
|
||||
raise NotImplementedError(f"You need to implement the method `callback_fn` for {self.__class__}")
|
||||
|
||||
def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]:
|
||||
return self.callback_fn(pipeline, step_index, timestep, callback_kwargs)
|
||||
|
||||
|
||||
class MultiPipelineCallbacks:
|
||||
"""
|
||||
This class is designed to handle multiple pipeline callbacks. It accepts a list of PipelineCallback objects and
|
||||
provides a unified interface for calling all of them.
|
||||
"""
|
||||
|
||||
def __init__(self, callbacks: List[PipelineCallback]):
|
||||
self.callbacks = callbacks
|
||||
|
||||
@property
|
||||
def tensor_inputs(self) -> List[str]:
|
||||
return [input for callback in self.callbacks for input in callback.tensor_inputs]
|
||||
|
||||
def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]:
|
||||
"""
|
||||
Calls all the callbacks in order with the given arguments and returns the final callback_kwargs.
|
||||
"""
|
||||
for callback in self.callbacks:
|
||||
callback_kwargs = callback(pipeline, step_index, timestep, callback_kwargs)
|
||||
|
||||
return callback_kwargs
|
||||
|
||||
|
||||
class SDCFGCutoffCallback(PipelineCallback):
|
||||
"""
|
||||
Callback function for Stable Diffusion Pipelines. After certain number of steps (set by `cutoff_step_ratio` or
|
||||
`cutoff_step_index`), this callback will disable the CFG.
|
||||
|
||||
Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step.
|
||||
"""
|
||||
|
||||
tensor_inputs = ["prompt_embeds"]
|
||||
|
||||
def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]:
|
||||
cutoff_step_ratio = self.config.cutoff_step_ratio
|
||||
cutoff_step_index = self.config.cutoff_step_index
|
||||
|
||||
# Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
|
||||
cutoff_step = (
|
||||
cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
|
||||
)
|
||||
|
||||
if step_index == cutoff_step:
|
||||
prompt_embeds = callback_kwargs[self.tensor_inputs[0]]
|
||||
prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens.
|
||||
|
||||
pipeline._guidance_scale = 0.0
|
||||
|
||||
callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
|
||||
return callback_kwargs
|
||||
|
||||
|
||||
class SDXLCFGCutoffCallback(PipelineCallback):
|
||||
"""
|
||||
Callback function for Stable Diffusion XL Pipelines. After certain number of steps (set by `cutoff_step_ratio` or
|
||||
`cutoff_step_index`), this callback will disable the CFG.
|
||||
|
||||
Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step.
|
||||
"""
|
||||
|
||||
tensor_inputs = ["prompt_embeds", "add_text_embeds", "add_time_ids"]
|
||||
|
||||
def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]:
|
||||
cutoff_step_ratio = self.config.cutoff_step_ratio
|
||||
cutoff_step_index = self.config.cutoff_step_index
|
||||
|
||||
# Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
|
||||
cutoff_step = (
|
||||
cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
|
||||
)
|
||||
|
||||
if step_index == cutoff_step:
|
||||
prompt_embeds = callback_kwargs[self.tensor_inputs[0]]
|
||||
prompt_embeds = prompt_embeds[-1:] # "-1" denotes the embeddings for conditional text tokens.
|
||||
|
||||
add_text_embeds = callback_kwargs[self.tensor_inputs[1]]
|
||||
add_text_embeds = add_text_embeds[-1:] # "-1" denotes the embeddings for conditional pooled text tokens
|
||||
|
||||
add_time_ids = callback_kwargs[self.tensor_inputs[2]]
|
||||
add_time_ids = add_time_ids[-1:] # "-1" denotes the embeddings for conditional added time vector
|
||||
|
||||
pipeline._guidance_scale = 0.0
|
||||
|
||||
callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
|
||||
callback_kwargs[self.tensor_inputs[1]] = add_text_embeds
|
||||
callback_kwargs[self.tensor_inputs[2]] = add_time_ids
|
||||
return callback_kwargs
|
||||
|
||||
|
||||
class IPAdapterScaleCutoffCallback(PipelineCallback):
|
||||
"""
|
||||
Callback function for any pipeline that inherits `IPAdapterMixin`. After certain number of steps (set by
|
||||
`cutoff_step_ratio` or `cutoff_step_index`), this callback will set the IP Adapter scale to `0.0`.
|
||||
|
||||
Note: This callback mutates the IP Adapter attention processors by setting the scale to 0.0 after the cutoff step.
|
||||
"""
|
||||
|
||||
tensor_inputs = []
|
||||
|
||||
def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> Dict[str, Any]:
|
||||
cutoff_step_ratio = self.config.cutoff_step_ratio
|
||||
cutoff_step_index = self.config.cutoff_step_index
|
||||
|
||||
# Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
|
||||
cutoff_step = (
|
||||
cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
|
||||
)
|
||||
|
||||
if step_index == cutoff_step:
|
||||
pipeline.set_ip_adapter_scale(0.0)
|
||||
return callback_kwargs
|
||||
@@ -1,27 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from argparse import ArgumentParser
|
||||
|
||||
|
||||
class BaseDiffusersCLICommand(ABC):
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def register_subcommand(parser: ArgumentParser):
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def run(self):
|
||||
raise NotImplementedError()
|
||||
@@ -1,43 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
# Copyright 2024 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from argparse import ArgumentParser
|
||||
|
||||
from .env import EnvironmentCommand
|
||||
from .fp16_safetensors import FP16SafetensorsCommand
|
||||
|
||||
|
||||
def main():
|
||||
parser = ArgumentParser("Diffusers CLI tool", usage="diffusers-cli <command> [<args>]")
|
||||
commands_parser = parser.add_subparsers(help="diffusers-cli command helpers")
|
||||
|
||||
# Register commands
|
||||
EnvironmentCommand.register_subcommand(commands_parser)
|
||||
FP16SafetensorsCommand.register_subcommand(commands_parser)
|
||||
|
||||
# Let's go
|
||||
args = parser.parse_args()
|
||||
|
||||
if not hasattr(args, "func"):
|
||||
parser.print_help()
|
||||
exit(1)
|
||||
|
||||
# Run
|
||||
service = args.func(args)
|
||||
service.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,180 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import platform
|
||||
import subprocess
|
||||
from argparse import ArgumentParser
|
||||
|
||||
import huggingface_hub
|
||||
|
||||
from .. import __version__ as version
|
||||
from ..utils import (
|
||||
is_accelerate_available,
|
||||
is_bitsandbytes_available,
|
||||
is_flax_available,
|
||||
is_google_colab,
|
||||
is_peft_available,
|
||||
is_safetensors_available,
|
||||
is_torch_available,
|
||||
is_transformers_available,
|
||||
is_xformers_available,
|
||||
)
|
||||
from . import BaseDiffusersCLICommand
|
||||
|
||||
|
||||
def info_command_factory(_):
|
||||
return EnvironmentCommand()
|
||||
|
||||
|
||||
class EnvironmentCommand(BaseDiffusersCLICommand):
|
||||
@staticmethod
|
||||
def register_subcommand(parser: ArgumentParser) -> None:
|
||||
download_parser = parser.add_parser("env")
|
||||
download_parser.set_defaults(func=info_command_factory)
|
||||
|
||||
def run(self) -> dict:
|
||||
hub_version = huggingface_hub.__version__
|
||||
|
||||
safetensors_version = "not installed"
|
||||
if is_safetensors_available():
|
||||
import safetensors
|
||||
|
||||
safetensors_version = safetensors.__version__
|
||||
|
||||
pt_version = "not installed"
|
||||
pt_cuda_available = "NA"
|
||||
if is_torch_available():
|
||||
import torch
|
||||
|
||||
pt_version = torch.__version__
|
||||
pt_cuda_available = torch.cuda.is_available()
|
||||
|
||||
flax_version = "not installed"
|
||||
jax_version = "not installed"
|
||||
jaxlib_version = "not installed"
|
||||
jax_backend = "NA"
|
||||
if is_flax_available():
|
||||
import flax
|
||||
import jax
|
||||
import jaxlib
|
||||
|
||||
flax_version = flax.__version__
|
||||
jax_version = jax.__version__
|
||||
jaxlib_version = jaxlib.__version__
|
||||
jax_backend = jax.lib.xla_bridge.get_backend().platform
|
||||
|
||||
transformers_version = "not installed"
|
||||
if is_transformers_available():
|
||||
import transformers
|
||||
|
||||
transformers_version = transformers.__version__
|
||||
|
||||
accelerate_version = "not installed"
|
||||
if is_accelerate_available():
|
||||
import accelerate
|
||||
|
||||
accelerate_version = accelerate.__version__
|
||||
|
||||
peft_version = "not installed"
|
||||
if is_peft_available():
|
||||
import peft
|
||||
|
||||
peft_version = peft.__version__
|
||||
|
||||
bitsandbytes_version = "not installed"
|
||||
if is_bitsandbytes_available():
|
||||
import bitsandbytes
|
||||
|
||||
bitsandbytes_version = bitsandbytes.__version__
|
||||
|
||||
xformers_version = "not installed"
|
||||
if is_xformers_available():
|
||||
import xformers
|
||||
|
||||
xformers_version = xformers.__version__
|
||||
|
||||
platform_info = platform.platform()
|
||||
|
||||
is_google_colab_str = "Yes" if is_google_colab() else "No"
|
||||
|
||||
accelerator = "NA"
|
||||
if platform.system() in {"Linux", "Windows"}:
|
||||
try:
|
||||
sp = subprocess.Popen(
|
||||
["nvidia-smi", "--query-gpu=gpu_name,memory.total", "--format=csv,noheader"],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
)
|
||||
out_str, _ = sp.communicate()
|
||||
out_str = out_str.decode("utf-8")
|
||||
|
||||
if len(out_str) > 0:
|
||||
accelerator = out_str.strip()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
elif platform.system() == "Darwin": # Mac OS
|
||||
try:
|
||||
sp = subprocess.Popen(
|
||||
["system_profiler", "SPDisplaysDataType"],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
)
|
||||
out_str, _ = sp.communicate()
|
||||
out_str = out_str.decode("utf-8")
|
||||
|
||||
start = out_str.find("Chipset Model:")
|
||||
if start != -1:
|
||||
start += len("Chipset Model:")
|
||||
end = out_str.find("\n", start)
|
||||
accelerator = out_str[start:end].strip()
|
||||
|
||||
start = out_str.find("VRAM (Total):")
|
||||
if start != -1:
|
||||
start += len("VRAM (Total):")
|
||||
end = out_str.find("\n", start)
|
||||
accelerator += " VRAM: " + out_str[start:end].strip()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
else:
|
||||
print("It seems you are running an unusual OS. Could you fill in the accelerator manually?")
|
||||
|
||||
info = {
|
||||
"🤗 Diffusers version": version,
|
||||
"Platform": platform_info,
|
||||
"Running on Google Colab?": is_google_colab_str,
|
||||
"Python version": platform.python_version(),
|
||||
"PyTorch version (GPU?)": f"{pt_version} ({pt_cuda_available})",
|
||||
"Flax version (CPU?/GPU?/TPU?)": f"{flax_version} ({jax_backend})",
|
||||
"Jax version": jax_version,
|
||||
"JaxLib version": jaxlib_version,
|
||||
"Huggingface_hub version": hub_version,
|
||||
"Transformers version": transformers_version,
|
||||
"Accelerate version": accelerate_version,
|
||||
"PEFT version": peft_version,
|
||||
"Bitsandbytes version": bitsandbytes_version,
|
||||
"Safetensors version": safetensors_version,
|
||||
"xFormers version": xformers_version,
|
||||
"Accelerator": accelerator,
|
||||
"Using GPU in script?": "<fill in>",
|
||||
"Using distributed or parallel set-up in script?": "<fill in>",
|
||||
}
|
||||
|
||||
print("\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n")
|
||||
print(self.format_dict(info))
|
||||
|
||||
return info
|
||||
|
||||
@staticmethod
|
||||
def format_dict(d: dict) -> str:
|
||||
return "\n".join([f"- {prop}: {val}" for prop, val in d.items()]) + "\n"
|
||||
@@ -1,132 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
Usage example:
|
||||
diffusers-cli fp16_safetensors --ckpt_id=openai/shap-e --fp16 --use_safetensors
|
||||
"""
|
||||
|
||||
import glob
|
||||
import json
|
||||
import warnings
|
||||
from argparse import ArgumentParser, Namespace
|
||||
from importlib import import_module
|
||||
|
||||
import huggingface_hub
|
||||
import torch
|
||||
from huggingface_hub import hf_hub_download
|
||||
from packaging import version
|
||||
|
||||
from ..utils import logging
|
||||
from . import BaseDiffusersCLICommand
|
||||
|
||||
|
||||
def conversion_command_factory(args: Namespace):
|
||||
if args.use_auth_token:
|
||||
warnings.warn(
|
||||
"The `--use_auth_token` flag is deprecated and will be removed in a future version. Authentication is now"
|
||||
" handled automatically if user is logged in."
|
||||
)
|
||||
return FP16SafetensorsCommand(args.ckpt_id, args.fp16, args.use_safetensors)
|
||||
|
||||
|
||||
class FP16SafetensorsCommand(BaseDiffusersCLICommand):
|
||||
@staticmethod
|
||||
def register_subcommand(parser: ArgumentParser):
|
||||
conversion_parser = parser.add_parser("fp16_safetensors")
|
||||
conversion_parser.add_argument(
|
||||
"--ckpt_id",
|
||||
type=str,
|
||||
help="Repo id of the checkpoints on which to run the conversion. Example: 'openai/shap-e'.",
|
||||
)
|
||||
conversion_parser.add_argument(
|
||||
"--fp16", action="store_true", help="If serializing the variables in FP16 precision."
|
||||
)
|
||||
conversion_parser.add_argument(
|
||||
"--use_safetensors", action="store_true", help="If serializing in the safetensors format."
|
||||
)
|
||||
conversion_parser.add_argument(
|
||||
"--use_auth_token",
|
||||
action="store_true",
|
||||
help="When working with checkpoints having private visibility. When used `huggingface-cli login` needs to be run beforehand.",
|
||||
)
|
||||
conversion_parser.set_defaults(func=conversion_command_factory)
|
||||
|
||||
def __init__(self, ckpt_id: str, fp16: bool, use_safetensors: bool):
|
||||
self.logger = logging.get_logger("diffusers-cli/fp16_safetensors")
|
||||
self.ckpt_id = ckpt_id
|
||||
self.local_ckpt_dir = f"/tmp/{ckpt_id}"
|
||||
self.fp16 = fp16
|
||||
|
||||
self.use_safetensors = use_safetensors
|
||||
|
||||
if not self.use_safetensors and not self.fp16:
|
||||
raise NotImplementedError(
|
||||
"When `use_safetensors` and `fp16` both are False, then this command is of no use."
|
||||
)
|
||||
|
||||
def run(self):
|
||||
if version.parse(huggingface_hub.__version__) < version.parse("0.9.0"):
|
||||
raise ImportError(
|
||||
"The huggingface_hub version must be >= 0.9.0 to use this command. Please update your huggingface_hub"
|
||||
" installation."
|
||||
)
|
||||
else:
|
||||
from huggingface_hub import create_commit
|
||||
from huggingface_hub._commit_api import CommitOperationAdd
|
||||
|
||||
model_index = hf_hub_download(repo_id=self.ckpt_id, filename="model_index.json")
|
||||
with open(model_index, "r") as f:
|
||||
pipeline_class_name = json.load(f)["_class_name"]
|
||||
pipeline_class = getattr(import_module("diffusers"), pipeline_class_name)
|
||||
self.logger.info(f"Pipeline class imported: {pipeline_class_name}.")
|
||||
|
||||
# Load the appropriate pipeline. We could have use `DiffusionPipeline`
|
||||
# here, but just to avoid any rough edge cases.
|
||||
pipeline = pipeline_class.from_pretrained(
|
||||
self.ckpt_id, torch_dtype=torch.float16 if self.fp16 else torch.float32
|
||||
)
|
||||
pipeline.save_pretrained(
|
||||
self.local_ckpt_dir,
|
||||
safe_serialization=True if self.use_safetensors else False,
|
||||
variant="fp16" if self.fp16 else None,
|
||||
)
|
||||
self.logger.info(f"Pipeline locally saved to {self.local_ckpt_dir}.")
|
||||
|
||||
# Fetch all the paths.
|
||||
if self.fp16:
|
||||
modified_paths = glob.glob(f"{self.local_ckpt_dir}/*/*.fp16.*")
|
||||
elif self.use_safetensors:
|
||||
modified_paths = glob.glob(f"{self.local_ckpt_dir}/*/*.safetensors")
|
||||
|
||||
# Prepare for the PR.
|
||||
commit_message = f"Serialize variables with FP16: {self.fp16} and safetensors: {self.use_safetensors}."
|
||||
operations = []
|
||||
for path in modified_paths:
|
||||
operations.append(CommitOperationAdd(path_in_repo="/".join(path.split("/")[4:]), path_or_fileobj=path))
|
||||
|
||||
# Open the PR.
|
||||
commit_description = (
|
||||
"Variables converted by the [`diffusers`' `fp16_safetensors`"
|
||||
" CLI](https://github.com/huggingface/diffusers/blob/main/src/diffusers/commands/fp16_safetensors.py)."
|
||||
)
|
||||
hub_pr_url = create_commit(
|
||||
repo_id=self.ckpt_id,
|
||||
operations=operations,
|
||||
commit_message=commit_message,
|
||||
commit_description=commit_description,
|
||||
repo_type="model",
|
||||
create_pr=True,
|
||||
).pr_url
|
||||
self.logger.info(f"PR created here: {hub_pr_url}.")
|
||||
@@ -1,720 +0,0 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2024 The HuggingFace Inc. team.
|
||||
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""ConfigMixin base class and utilities."""
|
||||
|
||||
import dataclasses
|
||||
import functools
|
||||
import importlib
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from collections import OrderedDict
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
from huggingface_hub import create_repo, hf_hub_download
|
||||
from huggingface_hub.utils import (
|
||||
EntryNotFoundError,
|
||||
RepositoryNotFoundError,
|
||||
RevisionNotFoundError,
|
||||
validate_hf_hub_args,
|
||||
)
|
||||
from requests import HTTPError
|
||||
|
||||
from . import __version__
|
||||
from .utils import (
|
||||
HUGGINGFACE_CO_RESOLVE_ENDPOINT,
|
||||
DummyObject,
|
||||
deprecate,
|
||||
extract_commit_hash,
|
||||
http_user_agent,
|
||||
logging,
|
||||
)
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
_re_configuration_file = re.compile(r"config\.(.*)\.json")
|
||||
|
||||
|
||||
class FrozenDict(OrderedDict):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
for key, value in self.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
self.__frozen = True
|
||||
|
||||
def __delitem__(self, *args, **kwargs):
|
||||
raise Exception(f"You cannot use ``__delitem__`` on a {self.__class__.__name__} instance.")
|
||||
|
||||
def setdefault(self, *args, **kwargs):
|
||||
raise Exception(f"You cannot use ``setdefault`` on a {self.__class__.__name__} instance.")
|
||||
|
||||
def pop(self, *args, **kwargs):
|
||||
raise Exception(f"You cannot use ``pop`` on a {self.__class__.__name__} instance.")
|
||||
|
||||
def update(self, *args, **kwargs):
|
||||
raise Exception(f"You cannot use ``update`` on a {self.__class__.__name__} instance.")
|
||||
|
||||
def __setattr__(self, name, value):
|
||||
if hasattr(self, "__frozen") and self.__frozen:
|
||||
raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.")
|
||||
super().__setattr__(name, value)
|
||||
|
||||
def __setitem__(self, name, value):
|
||||
if hasattr(self, "__frozen") and self.__frozen:
|
||||
raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.")
|
||||
super().__setitem__(name, value)
|
||||
|
||||
|
||||
class ConfigMixin:
|
||||
r"""
|
||||
Base class for all configuration classes. All configuration parameters are stored under `self.config`. Also
|
||||
provides the [`~ConfigMixin.from_config`] and [`~ConfigMixin.save_config`] methods for loading, downloading, and
|
||||
saving classes that inherit from [`ConfigMixin`].
|
||||
|
||||
Class attributes:
|
||||
- **config_name** (`str`) -- A filename under which the config should stored when calling
|
||||
[`~ConfigMixin.save_config`] (should be overridden by parent class).
|
||||
- **ignore_for_config** (`List[str]`) -- A list of attributes that should not be saved in the config (should be
|
||||
overridden by subclass).
|
||||
- **has_compatibles** (`bool`) -- Whether the class has compatible classes (should be overridden by subclass).
|
||||
- **_deprecated_kwargs** (`List[str]`) -- Keyword arguments that are deprecated. Note that the `init` function
|
||||
should only have a `kwargs` argument if at least one argument is deprecated (should be overridden by
|
||||
subclass).
|
||||
"""
|
||||
|
||||
config_name = None
|
||||
ignore_for_config = []
|
||||
has_compatibles = False
|
||||
|
||||
_deprecated_kwargs = []
|
||||
|
||||
def register_to_config(self, **kwargs):
|
||||
if self.config_name is None:
|
||||
raise NotImplementedError(f"Make sure that {self.__class__} has defined a class name `config_name`")
|
||||
# Special case for `kwargs` used in deprecation warning added to schedulers
|
||||
# TODO: remove this when we remove the deprecation warning, and the `kwargs` argument,
|
||||
# or solve in a more general way.
|
||||
kwargs.pop("kwargs", None)
|
||||
|
||||
if not hasattr(self, "_internal_dict"):
|
||||
internal_dict = kwargs
|
||||
else:
|
||||
previous_dict = dict(self._internal_dict)
|
||||
internal_dict = {**self._internal_dict, **kwargs}
|
||||
logger.debug(f"Updating config from {previous_dict} to {internal_dict}")
|
||||
|
||||
self._internal_dict = FrozenDict(internal_dict)
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
"""The only reason we overwrite `getattr` here is to gracefully deprecate accessing
|
||||
config attributes directly. See https://github.com/huggingface/diffusers/pull/3129
|
||||
|
||||
This function is mostly copied from PyTorch's __getattr__ overwrite:
|
||||
https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module
|
||||
"""
|
||||
|
||||
is_in_config = "_internal_dict" in self.__dict__ and hasattr(self.__dict__["_internal_dict"], name)
|
||||
is_attribute = name in self.__dict__
|
||||
|
||||
if is_in_config and not is_attribute:
|
||||
deprecation_message = f"Accessing config attribute `{name}` directly via '{type(self).__name__}' object attribute is deprecated. Please access '{name}' over '{type(self).__name__}'s config object instead, e.g. 'scheduler.config.{name}'."
|
||||
deprecate("direct config name access", "1.0.0", deprecation_message, standard_warn=False)
|
||||
return self._internal_dict[name]
|
||||
|
||||
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
|
||||
|
||||
def save_config(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs):
|
||||
"""
|
||||
Save a configuration object to the directory specified in `save_directory` so that it can be reloaded using the
|
||||
[`~ConfigMixin.from_config`] class method.
|
||||
|
||||
Args:
|
||||
save_directory (`str` or `os.PathLike`):
|
||||
Directory where the configuration JSON file is saved (will be created if it does not exist).
|
||||
push_to_hub (`bool`, *optional*, defaults to `False`):
|
||||
Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
|
||||
repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
|
||||
namespace).
|
||||
kwargs (`Dict[str, Any]`, *optional*):
|
||||
Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
|
||||
"""
|
||||
if os.path.isfile(save_directory):
|
||||
raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
|
||||
|
||||
os.makedirs(save_directory, exist_ok=True)
|
||||
|
||||
# If we save using the predefined names, we can load using `from_config`
|
||||
output_config_file = os.path.join(save_directory, self.config_name)
|
||||
|
||||
self.to_json_file(output_config_file)
|
||||
logger.info(f"Configuration saved in {output_config_file}")
|
||||
|
||||
if push_to_hub:
|
||||
commit_message = kwargs.pop("commit_message", None)
|
||||
private = kwargs.pop("private", False)
|
||||
create_pr = kwargs.pop("create_pr", False)
|
||||
token = kwargs.pop("token", None)
|
||||
repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
|
||||
repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id
|
||||
|
||||
self._upload_folder(
|
||||
save_directory,
|
||||
repo_id,
|
||||
token=token,
|
||||
commit_message=commit_message,
|
||||
create_pr=create_pr,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Union[FrozenDict, Dict[str, Any]] = None, return_unused_kwargs=False, **kwargs):
|
||||
r"""
|
||||
Instantiate a Python class from a config dictionary.
|
||||
|
||||
Parameters:
|
||||
config (`Dict[str, Any]`):
|
||||
A config dictionary from which the Python class is instantiated. Make sure to only load configuration
|
||||
files of compatible classes.
|
||||
return_unused_kwargs (`bool`, *optional*, defaults to `False`):
|
||||
Whether kwargs that are not consumed by the Python class should be returned or not.
|
||||
kwargs (remaining dictionary of keyword arguments, *optional*):
|
||||
Can be used to update the configuration object (after it is loaded) and initiate the Python class.
|
||||
`**kwargs` are passed directly to the underlying scheduler/model's `__init__` method and eventually
|
||||
overwrite the same named arguments in `config`.
|
||||
|
||||
Returns:
|
||||
[`ModelMixin`] or [`SchedulerMixin`]:
|
||||
A model or scheduler object instantiated from a config dictionary.
|
||||
|
||||
Examples:
|
||||
|
||||
```python
|
||||
>>> from diffusers import DDPMScheduler, DDIMScheduler, PNDMScheduler
|
||||
|
||||
>>> # Download scheduler from huggingface.co and cache.
|
||||
>>> scheduler = DDPMScheduler.from_pretrained("google/ddpm-cifar10-32")
|
||||
|
||||
>>> # Instantiate DDIM scheduler class with same config as DDPM
|
||||
>>> scheduler = DDIMScheduler.from_config(scheduler.config)
|
||||
|
||||
>>> # Instantiate PNDM scheduler class with same config as DDPM
|
||||
>>> scheduler = PNDMScheduler.from_config(scheduler.config)
|
||||
```
|
||||
"""
|
||||
# <===== TO BE REMOVED WITH DEPRECATION
|
||||
# TODO(Patrick) - make sure to remove the following lines when config=="model_path" is deprecated
|
||||
if "pretrained_model_name_or_path" in kwargs:
|
||||
config = kwargs.pop("pretrained_model_name_or_path")
|
||||
|
||||
if config is None:
|
||||
raise ValueError("Please make sure to provide a config as the first positional argument.")
|
||||
# ======>
|
||||
|
||||
if not isinstance(config, dict):
|
||||
deprecation_message = "It is deprecated to pass a pretrained model name or path to `from_config`."
|
||||
if "Scheduler" in cls.__name__:
|
||||
deprecation_message += (
|
||||
f"If you were trying to load a scheduler, please use {cls}.from_pretrained(...) instead."
|
||||
" Otherwise, please make sure to pass a configuration dictionary instead. This functionality will"
|
||||
" be removed in v1.0.0."
|
||||
)
|
||||
elif "Model" in cls.__name__:
|
||||
deprecation_message += (
|
||||
f"If you were trying to load a model, please use {cls}.load_config(...) followed by"
|
||||
f" {cls}.from_config(...) instead. Otherwise, please make sure to pass a configuration dictionary"
|
||||
" instead. This functionality will be removed in v1.0.0."
|
||||
)
|
||||
deprecate("config-passed-as-path", "1.0.0", deprecation_message, standard_warn=False)
|
||||
config, kwargs = cls.load_config(pretrained_model_name_or_path=config, return_unused_kwargs=True, **kwargs)
|
||||
|
||||
init_dict, unused_kwargs, hidden_dict = cls.extract_init_dict(config, **kwargs)
|
||||
|
||||
# Allow dtype to be specified on initialization
|
||||
if "dtype" in unused_kwargs:
|
||||
init_dict["dtype"] = unused_kwargs.pop("dtype")
|
||||
|
||||
# add possible deprecated kwargs
|
||||
for deprecated_kwarg in cls._deprecated_kwargs:
|
||||
if deprecated_kwarg in unused_kwargs:
|
||||
init_dict[deprecated_kwarg] = unused_kwargs.pop(deprecated_kwarg)
|
||||
|
||||
# Return model and optionally state and/or unused_kwargs
|
||||
model = cls(**init_dict)
|
||||
|
||||
# make sure to also save config parameters that might be used for compatible classes
|
||||
# update _class_name
|
||||
if "_class_name" in hidden_dict:
|
||||
hidden_dict["_class_name"] = cls.__name__
|
||||
|
||||
model.register_to_config(**hidden_dict)
|
||||
|
||||
# add hidden kwargs of compatible classes to unused_kwargs
|
||||
unused_kwargs = {**unused_kwargs, **hidden_dict}
|
||||
|
||||
if return_unused_kwargs:
|
||||
return (model, unused_kwargs)
|
||||
else:
|
||||
return model
|
||||
|
||||
@classmethod
|
||||
def get_config_dict(cls, *args, **kwargs):
|
||||
deprecation_message = (
|
||||
f" The function get_config_dict is deprecated. Please use {cls}.load_config instead. This function will be"
|
||||
" removed in version v1.0.0"
|
||||
)
|
||||
deprecate("get_config_dict", "1.0.0", deprecation_message, standard_warn=False)
|
||||
return cls.load_config(*args, **kwargs)
|
||||
|
||||
@classmethod
|
||||
@validate_hf_hub_args
|
||||
def load_config(
|
||||
cls,
|
||||
pretrained_model_name_or_path: Union[str, os.PathLike],
|
||||
return_unused_kwargs=False,
|
||||
return_commit_hash=False,
|
||||
**kwargs,
|
||||
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
|
||||
r"""
|
||||
Load a model or scheduler configuration.
|
||||
|
||||
Parameters:
|
||||
pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
|
||||
Can be either:
|
||||
|
||||
- A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
|
||||
the Hub.
|
||||
- A path to a *directory* (for example `./my_model_directory`) containing model weights saved with
|
||||
[`~ConfigMixin.save_config`].
|
||||
|
||||
cache_dir (`Union[str, os.PathLike]`, *optional*):
|
||||
Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
|
||||
is not used.
|
||||
force_download (`bool`, *optional*, defaults to `False`):
|
||||
Whether or not to force the (re-)download of the model weights and configuration files, overriding the
|
||||
cached versions if they exist.
|
||||
proxies (`Dict[str, str]`, *optional*):
|
||||
A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
|
||||
'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
|
||||
output_loading_info(`bool`, *optional*, defaults to `False`):
|
||||
Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
|
||||
local_files_only (`bool`, *optional*, defaults to `False`):
|
||||
Whether to only load local model weights and configuration files or not. If set to `True`, the model
|
||||
won't be downloaded from the Hub.
|
||||
token (`str` or *bool*, *optional*):
|
||||
The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
|
||||
`diffusers-cli login` (stored in `~/.huggingface`) is used.
|
||||
revision (`str`, *optional*, defaults to `"main"`):
|
||||
The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
|
||||
allowed by Git.
|
||||
subfolder (`str`, *optional*, defaults to `""`):
|
||||
The subfolder location of a model file within a larger model repository on the Hub or locally.
|
||||
return_unused_kwargs (`bool`, *optional*, defaults to `False):
|
||||
Whether unused keyword arguments of the config are returned.
|
||||
return_commit_hash (`bool`, *optional*, defaults to `False):
|
||||
Whether the `commit_hash` of the loaded configuration are returned.
|
||||
|
||||
Returns:
|
||||
`dict`:
|
||||
A dictionary of all the parameters stored in a JSON configuration file.
|
||||
|
||||
"""
|
||||
cache_dir = kwargs.pop("cache_dir", None)
|
||||
local_dir = kwargs.pop("local_dir", None)
|
||||
local_dir_use_symlinks = kwargs.pop("local_dir_use_symlinks", "auto")
|
||||
force_download = kwargs.pop("force_download", False)
|
||||
proxies = kwargs.pop("proxies", None)
|
||||
token = kwargs.pop("token", None)
|
||||
local_files_only = kwargs.pop("local_files_only", False)
|
||||
revision = kwargs.pop("revision", None)
|
||||
_ = kwargs.pop("mirror", None)
|
||||
subfolder = kwargs.pop("subfolder", None)
|
||||
user_agent = kwargs.pop("user_agent", {})
|
||||
|
||||
user_agent = {**user_agent, "file_type": "config"}
|
||||
user_agent = http_user_agent(user_agent)
|
||||
|
||||
pretrained_model_name_or_path = str(pretrained_model_name_or_path)
|
||||
|
||||
if cls.config_name is None:
|
||||
raise ValueError(
|
||||
"`self.config_name` is not defined. Note that one should not load a config from "
|
||||
"`ConfigMixin`. Please make sure to define `config_name` in a class inheriting from `ConfigMixin`"
|
||||
)
|
||||
|
||||
if os.path.isfile(pretrained_model_name_or_path):
|
||||
config_file = pretrained_model_name_or_path
|
||||
elif os.path.isdir(pretrained_model_name_or_path):
|
||||
if subfolder is not None and os.path.isfile(
|
||||
os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
|
||||
):
|
||||
config_file = os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
|
||||
elif os.path.isfile(os.path.join(pretrained_model_name_or_path, cls.config_name)):
|
||||
# Load from a PyTorch checkpoint
|
||||
config_file = os.path.join(pretrained_model_name_or_path, cls.config_name)
|
||||
else:
|
||||
raise EnvironmentError(
|
||||
f"Error no file named {cls.config_name} found in directory {pretrained_model_name_or_path}."
|
||||
)
|
||||
else:
|
||||
try:
|
||||
# Load from URL or cache if already cached
|
||||
config_file = hf_hub_download(
|
||||
pretrained_model_name_or_path,
|
||||
filename=cls.config_name,
|
||||
cache_dir=cache_dir,
|
||||
force_download=force_download,
|
||||
proxies=proxies,
|
||||
local_files_only=local_files_only,
|
||||
token=token,
|
||||
user_agent=user_agent,
|
||||
subfolder=subfolder,
|
||||
revision=revision,
|
||||
local_dir=local_dir,
|
||||
local_dir_use_symlinks=local_dir_use_symlinks,
|
||||
)
|
||||
except RepositoryNotFoundError:
|
||||
raise EnvironmentError(
|
||||
f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier"
|
||||
" listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a"
|
||||
" token having permission to this repo with `token` or log in with `huggingface-cli login`."
|
||||
)
|
||||
except RevisionNotFoundError:
|
||||
raise EnvironmentError(
|
||||
f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for"
|
||||
" this model name. Check the model page at"
|
||||
f" 'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions."
|
||||
)
|
||||
except EntryNotFoundError:
|
||||
raise EnvironmentError(
|
||||
f"{pretrained_model_name_or_path} does not appear to have a file named {cls.config_name}."
|
||||
)
|
||||
except HTTPError as err:
|
||||
raise EnvironmentError(
|
||||
"There was a specific connection error when trying to load"
|
||||
f" {pretrained_model_name_or_path}:\n{err}"
|
||||
)
|
||||
except ValueError:
|
||||
raise EnvironmentError(
|
||||
f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it"
|
||||
f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a"
|
||||
f" directory containing a {cls.config_name} file.\nCheckout your internet connection or see how to"
|
||||
" run the library in offline mode at"
|
||||
" 'https://huggingface.co/docs/diffusers/installation#offline-mode'."
|
||||
)
|
||||
except EnvironmentError:
|
||||
raise EnvironmentError(
|
||||
f"Can't load config for '{pretrained_model_name_or_path}'. If you were trying to load it from "
|
||||
"'https://huggingface.co/models', make sure you don't have a local directory with the same name. "
|
||||
f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory "
|
||||
f"containing a {cls.config_name} file"
|
||||
)
|
||||
|
||||
try:
|
||||
# Load config dict
|
||||
config_dict = cls._dict_from_json_file(config_file)
|
||||
|
||||
commit_hash = extract_commit_hash(config_file)
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
raise EnvironmentError(f"It looks like the config file at '{config_file}' is not a valid JSON file.")
|
||||
|
||||
if not (return_unused_kwargs or return_commit_hash):
|
||||
return config_dict
|
||||
|
||||
outputs = (config_dict,)
|
||||
|
||||
if return_unused_kwargs:
|
||||
outputs += (kwargs,)
|
||||
|
||||
if return_commit_hash:
|
||||
outputs += (commit_hash,)
|
||||
|
||||
return outputs
|
||||
|
||||
@staticmethod
|
||||
def _get_init_keys(input_class):
|
||||
return set(dict(inspect.signature(input_class.__init__).parameters).keys())
|
||||
|
||||
@classmethod
|
||||
def extract_init_dict(cls, config_dict, **kwargs):
|
||||
# Skip keys that were not present in the original config, so default __init__ values were used
|
||||
used_defaults = config_dict.get("_use_default_values", [])
|
||||
config_dict = {k: v for k, v in config_dict.items() if k not in used_defaults and k != "_use_default_values"}
|
||||
|
||||
# 0. Copy origin config dict
|
||||
original_dict = dict(config_dict.items())
|
||||
|
||||
# 1. Retrieve expected config attributes from __init__ signature
|
||||
expected_keys = cls._get_init_keys(cls)
|
||||
expected_keys.remove("self")
|
||||
# remove general kwargs if present in dict
|
||||
if "kwargs" in expected_keys:
|
||||
expected_keys.remove("kwargs")
|
||||
# remove flax internal keys
|
||||
if hasattr(cls, "_flax_internal_args"):
|
||||
for arg in cls._flax_internal_args:
|
||||
expected_keys.remove(arg)
|
||||
|
||||
# 2. Remove attributes that cannot be expected from expected config attributes
|
||||
# remove keys to be ignored
|
||||
if len(cls.ignore_for_config) > 0:
|
||||
expected_keys = expected_keys - set(cls.ignore_for_config)
|
||||
|
||||
# load diffusers library to import compatible and original scheduler
|
||||
diffusers_library = importlib.import_module(__name__.split(".")[0])
|
||||
|
||||
if cls.has_compatibles:
|
||||
compatible_classes = [c for c in cls._get_compatibles() if not isinstance(c, DummyObject)]
|
||||
else:
|
||||
compatible_classes = []
|
||||
|
||||
expected_keys_comp_cls = set()
|
||||
for c in compatible_classes:
|
||||
expected_keys_c = cls._get_init_keys(c)
|
||||
expected_keys_comp_cls = expected_keys_comp_cls.union(expected_keys_c)
|
||||
expected_keys_comp_cls = expected_keys_comp_cls - cls._get_init_keys(cls)
|
||||
config_dict = {k: v for k, v in config_dict.items() if k not in expected_keys_comp_cls}
|
||||
|
||||
# remove attributes from orig class that cannot be expected
|
||||
orig_cls_name = config_dict.pop("_class_name", cls.__name__)
|
||||
if (
|
||||
isinstance(orig_cls_name, str)
|
||||
and orig_cls_name != cls.__name__
|
||||
and hasattr(diffusers_library, orig_cls_name)
|
||||
):
|
||||
orig_cls = getattr(diffusers_library, orig_cls_name)
|
||||
unexpected_keys_from_orig = cls._get_init_keys(orig_cls) - expected_keys
|
||||
config_dict = {k: v for k, v in config_dict.items() if k not in unexpected_keys_from_orig}
|
||||
elif not isinstance(orig_cls_name, str) and not isinstance(orig_cls_name, (list, tuple)):
|
||||
raise ValueError(
|
||||
"Make sure that the `_class_name` is of type string or list of string (for custom pipelines)."
|
||||
)
|
||||
|
||||
# remove private attributes
|
||||
config_dict = {k: v for k, v in config_dict.items() if not k.startswith("_")}
|
||||
|
||||
# 3. Create keyword arguments that will be passed to __init__ from expected keyword arguments
|
||||
init_dict = {}
|
||||
for key in expected_keys:
|
||||
# if config param is passed to kwarg and is present in config dict
|
||||
# it should overwrite existing config dict key
|
||||
if key in kwargs and key in config_dict:
|
||||
config_dict[key] = kwargs.pop(key)
|
||||
|
||||
if key in kwargs:
|
||||
# overwrite key
|
||||
init_dict[key] = kwargs.pop(key)
|
||||
elif key in config_dict:
|
||||
# use value from config dict
|
||||
init_dict[key] = config_dict.pop(key)
|
||||
|
||||
# 4. Give nice warning if unexpected values have been passed
|
||||
if len(config_dict) > 0:
|
||||
logger.warning(
|
||||
f"The config attributes {config_dict} were passed to {cls.__name__}, "
|
||||
"but are not expected and will be ignored. Please verify your "
|
||||
f"{cls.config_name} configuration file."
|
||||
)
|
||||
|
||||
# 5. Give nice info if config attributes are initialized to default because they have not been passed
|
||||
passed_keys = set(init_dict.keys())
|
||||
if len(expected_keys - passed_keys) > 0:
|
||||
logger.info(
|
||||
f"{expected_keys - passed_keys} was not found in config. Values will be initialized to default values."
|
||||
)
|
||||
|
||||
# 6. Define unused keyword arguments
|
||||
unused_kwargs = {**config_dict, **kwargs}
|
||||
|
||||
# 7. Define "hidden" config parameters that were saved for compatible classes
|
||||
hidden_config_dict = {k: v for k, v in original_dict.items() if k not in init_dict}
|
||||
|
||||
return init_dict, unused_kwargs, hidden_config_dict
|
||||
|
||||
@classmethod
|
||||
def _dict_from_json_file(cls, json_file: Union[str, os.PathLike]):
|
||||
with open(json_file, "r", encoding="utf-8") as reader:
|
||||
text = reader.read()
|
||||
return json.loads(text)
|
||||
|
||||
def __repr__(self):
|
||||
return f"{self.__class__.__name__} {self.to_json_string()}"
|
||||
|
||||
@property
|
||||
def config(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Returns the config of the class as a frozen dictionary
|
||||
|
||||
Returns:
|
||||
`Dict[str, Any]`: Config of the class.
|
||||
"""
|
||||
return self._internal_dict
|
||||
|
||||
def to_json_string(self) -> str:
|
||||
"""
|
||||
Serializes the configuration instance to a JSON string.
|
||||
|
||||
Returns:
|
||||
`str`:
|
||||
String containing all the attributes that make up the configuration instance in JSON format.
|
||||
"""
|
||||
config_dict = self._internal_dict if hasattr(self, "_internal_dict") else {}
|
||||
config_dict["_class_name"] = self.__class__.__name__
|
||||
config_dict["_diffusers_version"] = __version__
|
||||
|
||||
def to_json_saveable(value):
|
||||
if isinstance(value, np.ndarray):
|
||||
value = value.tolist()
|
||||
elif isinstance(value, Path):
|
||||
value = value.as_posix()
|
||||
return value
|
||||
|
||||
config_dict = {k: to_json_saveable(v) for k, v in config_dict.items()}
|
||||
# Don't save "_ignore_files" or "_use_default_values"
|
||||
config_dict.pop("_ignore_files", None)
|
||||
config_dict.pop("_use_default_values", None)
|
||||
|
||||
return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
|
||||
|
||||
def to_json_file(self, json_file_path: Union[str, os.PathLike]):
|
||||
"""
|
||||
Save the configuration instance's parameters to a JSON file.
|
||||
|
||||
Args:
|
||||
json_file_path (`str` or `os.PathLike`):
|
||||
Path to the JSON file to save a configuration instance's parameters.
|
||||
"""
|
||||
with open(json_file_path, "w", encoding="utf-8") as writer:
|
||||
writer.write(self.to_json_string())
|
||||
|
||||
|
||||
def register_to_config(init):
|
||||
r"""
|
||||
Decorator to apply on the init of classes inheriting from [`ConfigMixin`] so that all the arguments are
|
||||
automatically sent to `self.register_for_config`. To ignore a specific argument accepted by the init but that
|
||||
shouldn't be registered in the config, use the `ignore_for_config` class variable
|
||||
|
||||
Warning: Once decorated, all private arguments (beginning with an underscore) are trashed and not sent to the init!
|
||||
"""
|
||||
|
||||
@functools.wraps(init)
|
||||
def inner_init(self, *args, **kwargs):
|
||||
# Ignore private kwargs in the init.
|
||||
init_kwargs = {k: v for k, v in kwargs.items() if not k.startswith("_")}
|
||||
config_init_kwargs = {k: v for k, v in kwargs.items() if k.startswith("_")}
|
||||
if not isinstance(self, ConfigMixin):
|
||||
raise RuntimeError(
|
||||
f"`@register_for_config` was applied to {self.__class__.__name__} init method, but this class does "
|
||||
"not inherit from `ConfigMixin`."
|
||||
)
|
||||
|
||||
ignore = getattr(self, "ignore_for_config", [])
|
||||
# Get positional arguments aligned with kwargs
|
||||
new_kwargs = {}
|
||||
signature = inspect.signature(init)
|
||||
parameters = {
|
||||
name: p.default for i, (name, p) in enumerate(signature.parameters.items()) if i > 0 and name not in ignore
|
||||
}
|
||||
for arg, name in zip(args, parameters.keys()):
|
||||
new_kwargs[name] = arg
|
||||
|
||||
# Then add all kwargs
|
||||
new_kwargs.update(
|
||||
{
|
||||
k: init_kwargs.get(k, default)
|
||||
for k, default in parameters.items()
|
||||
if k not in ignore and k not in new_kwargs
|
||||
}
|
||||
)
|
||||
|
||||
# Take note of the parameters that were not present in the loaded config
|
||||
if len(set(new_kwargs.keys()) - set(init_kwargs)) > 0:
|
||||
new_kwargs["_use_default_values"] = list(set(new_kwargs.keys()) - set(init_kwargs))
|
||||
|
||||
new_kwargs = {**config_init_kwargs, **new_kwargs}
|
||||
getattr(self, "register_to_config")(**new_kwargs)
|
||||
init(self, *args, **init_kwargs)
|
||||
|
||||
return inner_init
|
||||
|
||||
|
||||
def flax_register_to_config(cls):
|
||||
original_init = cls.__init__
|
||||
|
||||
@functools.wraps(original_init)
|
||||
def init(self, *args, **kwargs):
|
||||
if not isinstance(self, ConfigMixin):
|
||||
raise RuntimeError(
|
||||
f"`@register_for_config` was applied to {self.__class__.__name__} init method, but this class does "
|
||||
"not inherit from `ConfigMixin`."
|
||||
)
|
||||
|
||||
# Ignore private kwargs in the init. Retrieve all passed attributes
|
||||
init_kwargs = dict(kwargs.items())
|
||||
|
||||
# Retrieve default values
|
||||
fields = dataclasses.fields(self)
|
||||
default_kwargs = {}
|
||||
for field in fields:
|
||||
# ignore flax specific attributes
|
||||
if field.name in self._flax_internal_args:
|
||||
continue
|
||||
if type(field.default) == dataclasses._MISSING_TYPE:
|
||||
default_kwargs[field.name] = None
|
||||
else:
|
||||
default_kwargs[field.name] = getattr(self, field.name)
|
||||
|
||||
# Make sure init_kwargs override default kwargs
|
||||
new_kwargs = {**default_kwargs, **init_kwargs}
|
||||
# dtype should be part of `init_kwargs`, but not `new_kwargs`
|
||||
if "dtype" in new_kwargs:
|
||||
new_kwargs.pop("dtype")
|
||||
|
||||
# Get positional arguments aligned with kwargs
|
||||
for i, arg in enumerate(args):
|
||||
name = fields[i].name
|
||||
new_kwargs[name] = arg
|
||||
|
||||
# Take note of the parameters that were not present in the loaded config
|
||||
if len(set(new_kwargs.keys()) - set(init_kwargs)) > 0:
|
||||
new_kwargs["_use_default_values"] = list(set(new_kwargs.keys()) - set(init_kwargs))
|
||||
|
||||
getattr(self, "register_to_config")(**new_kwargs)
|
||||
original_init(self, *args, **kwargs)
|
||||
|
||||
cls.__init__ = init
|
||||
return cls
|
||||
|
||||
|
||||
class LegacyConfigMixin(ConfigMixin):
|
||||
r"""
|
||||
A subclass of `ConfigMixin` to resolve class mapping from legacy classes (like `Transformer2DModel`) to more
|
||||
pipeline-specific classes (like `DiTTransformer2DModel`).
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Union[FrozenDict, Dict[str, Any]] = None, return_unused_kwargs=False, **kwargs):
|
||||
# To prevent dependency import problem.
|
||||
from .models.model_loading_utils import _fetch_remapped_cls_from_config
|
||||
|
||||
# resolve remapping
|
||||
remapped_class = _fetch_remapped_cls_from_config(config, cls)
|
||||
|
||||
return remapped_class.from_config(config, return_unused_kwargs, **kwargs)
|
||||
@@ -1,34 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from .dependency_versions_table import deps
|
||||
from .utils.versions import require_version, require_version_core
|
||||
|
||||
|
||||
# define which module versions we always want to check at run time
|
||||
# (usually the ones defined in `install_requires` in setup.py)
|
||||
#
|
||||
# order specific notes:
|
||||
# - tqdm must be checked before tokenizers
|
||||
|
||||
pkgs_to_check_at_runtime = "python requests filelock numpy".split()
|
||||
for pkg in pkgs_to_check_at_runtime:
|
||||
if pkg in deps:
|
||||
require_version_core(deps[pkg])
|
||||
else:
|
||||
raise ValueError(f"can't find {pkg} in {deps.keys()}, check dependency_versions_table.py")
|
||||
|
||||
|
||||
def dep_version_check(pkg, hint=None):
|
||||
require_version(deps[pkg], hint)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user