libcom 0.0.1.post10__tar.gz → 0.1.0__tar.gz

{libcom-0.0.1.post10 → libcom-0.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: libcom
-Version: 0.0.1.post10
+Version: 0.1.0
 Summary: Image Composition Toolbox
 Home-page: https://github.com/bcmi/libcom
 Author: BCMI Lab
@@ -31,12 +31,15 @@ License-File: LICENSE
 [![Downloads](https://static.pepy.tech/badge/libcom)](https://pepy.tech/project/libcom)
 [![Hits](https://hits.sh/github.com/bcmi/libcom.svg?label=views)](https://hits.sh/github.com/bcmi/libcom/)
 [![Static Badge](https://img.shields.io/badge/Image%20Composition%20Demo-Green)](https://bcmi.sjtu.edu.cn/home/niuli/demo_image_composition/)
+[![Static Badge](https://img.shields.io/badge/Resources-Awesome-green?style=flat&label=Resources)](https://github.com/bcmi/Awesome-Image-Composition)
 [![Static Badge](https://img.shields.io/badge/survey-arxiv%3A2106.14490-red)](https://arxiv.org/pdf/2106.14490.pdf)
-[![GitHub](https://img.shields.io/github/license/bcmi/libcom)](https://github.com/bcmi/libcom/blob/main/LICENSE)
 
+**We co-founded a startup company [miguo.ai](https://miguocomics.com/), dedicated to accelerating the production of comics and animations using AIGC technology. If you are seeking internship or full-time positions, please feel free to send your resume to hr@miguocomics.com.**
+
+---
 
 ## Introduction
-**_libcom_ (the library of image composition) is an image composition toolbox.** The goal of image composition is inserting one foreground into a background image to get a realistic composite image, by addressing the inconsistencies (appearance, geometry, and semantic inconsistency) between foreground and background. Generally speaking, image composition could be used to combine the visual elements from different images.
+**_libcom_ (the library of image composition) is an image composition toolbox.** The goal of image composition (object insertion) is inserting one foreground into a background image to get a realistic composite image, by addressing the inconsistencies (appearance, geometry, and semantic inconsistency) between foreground and background. Generally speaking, image composition could be used to combine the visual elements from different images.
 <div align="center">
 </br>
 <img src="https://raw.githubusercontent.com/bcmi/libcom/main/resources/image_composition_task.gif" width="600" />
@@ -49,18 +52,20 @@ License-File: LICENSE
 ### Main Functions
 
 - **get_composite_image** generates composite images using naive copy-and-paste followed by image blending.
-- **OPAScoreModel**  evaluates the rationality of foreground object placement in a composite image.
-- **FOPAHeatMapModel** can predict the rationality scores for all locations/scales given a background-foreground pair, and output the composite image with optimal location/scale.  
+- **OPAScoreModel** [[OPA]](https://github.com/bcmi/Object-Placement-Assessment-Dataset-OPA) evaluates the rationality of foreground object placement in a composite image.
+- **FOPAHeatMapModel** [[FOPA]](https://github.com/bcmi/FOPA-Fast-Object-Placement-Assessment) can predict the rationality scores for all locations/scales given a background-foreground pair, and output the composite image with optimal location/scale.  
 - **color_transfer** adjusts the foreground color to match the background using traditional color transfer method.
-- **ImageHarmonizationModel** adjusts the foreground illumination to be compatible the background given photorealistic background and photorealistic foreground.
-- **PainterlyHarmonizationModel** adjusts the foreground style to be compatible with the background given artistic background and photorealistic foreground.
-- **HarmonyScoreModel** evaluates the harmony level between foreground and background in a composite image.
-- **InharmoniousLocalizationModel** localizes the inharmonious region in a synthetic image.
-- **FOSScoreModel** evaluates the compatibility between foreground and background in a composite image in terms of geometry and semantics.
-- **ControlComModel** is a generative image composition model, which unifies image blending and image harmonization in one diffusion model. 
-- **ShadowGenerationModel** generates plausible shadow for the inserted object in a composite image. 
+- **ImageHarmonizationModel** [[CDTNet]](https://github.com/bcmi/CDTNet-High-Resolution-Image-Harmonization)  [[PCTNet]](https://github.com/rakutentech/PCT-Net-Image-Harmonization) adjusts the foreground illumination to be compatible the background given photorealistic background and photorealistic foreground.
+- **PainterlyHarmonizationModel** [[PHDNet]](https://github.com/bcmi/PHDNet-Painterly-Image-Harmonization)  [[PHDiffusion]](https://github.com/bcmi/PHDiffusion-Painterly-Image-Harmonization) adjusts the foreground style to be compatible with the background given artistic background and photorealistic foreground.
+- **HarmonyScoreModel** [[BargainNet]](https://github.com/bcmi/BargainNet-Image-Harmonization) evaluates the harmony level between foreground and background in a composite image.
+- **InharmoniousLocalizationModel** [[MadisNet]](https://github.com/bcmi/MadisNet-Inharmonious-Region-Localization) localizes the inharmonious region in a synthetic image.
+- **FOSScoreModel** [[DiscoFOS]](https://github.com/bcmi/Foreground-Object-Search-Dataset-FOSD) evaluates the compatibility between foreground and background in a composite image in terms of geometry and semantics.
+- **ShadowGenerationModel** [[GPSDiffusion]](https://github.com/bcmi/GPSDiffusion-Object-Shadow-Generation) generates plausible shadow for the inserted object in a composite image. This model is unstable and you can pick the most satisfactory one from multiple generated results. 
+- **ControlComModel** [[ControlCom]](https://github.com/bcmi/ControlCom-Image-Composition) is a generative image composition model which unifies image blending and image harmonization. The pose and view of foreground stay unchanged. 
+- **MureObjectStitchModel** [[MureObjectStitch]](https://github.com/bcmi/MureObjectStitch-Image-Composition) is another generative image composition model which can adjust the pose and view of foreground. It supports multiple reference images of one foreground object. If you have a few images containing the foreground object, we suggest finetuning MureObjectStitch using these images for better detail preservation. 
+
 
-**For the detailed method descriptions, code examples, visualization results, and performance comments, please refer to our [[documents]](https://libcom.readthedocs.io/en/latest/).**
+**For the detailed method descriptions, code examples, visualization results, and performance comments, please refer to our [[documents]](https://libcom.readthedocs.io/en/latest/).** If the model performance is not satisfactory, you can finetune the pretrained model on your own dataset using the source repository and replace the original model. 
 
 ## Requirements
 
@@ -72,9 +77,9 @@ Please refer to [[Installation]](https://github.com/bcmi/libcom/blob/main/docs/g
 ## Contributors
 - Institution: [Brain-like Computing and Machine Intelligence (BCMI) Lab](https://bcmi.sjtu.edu.cn/).
 - Project Initiator & Team Manager: [Li Niu](https://www.ustcnewly.com/index.html). 
-- Architect & Lead Developer: [Bo Zhang](https://bo-zhang-cs.github.io/).   
+- Architect & Lead Developer: [Bo Zhang](https://bo-zhang-cs.github.io/), [Yujie Zhou](https://github.com/YujieOuO).
 - Documentation Manager: [Jiacheng Sui](https://github.com/charlessjc).
-- Module Developers: [Jiacheng Sui](https://github.com/charlessjc), [Binjie Gao](https://github.com/WhynotHAHA), [Lingxiao Lu](https://github.com/pokaaa), [Xinhao Tao](https://github.com/taoxinhao13), [Junyan Cao](https://github.com/cjy-4), [Junqi You](https://github.com/dhmbb2).
+- Module Developers: [Jiacheng Sui](https://github.com/charlessjc), [Binjie Gao](https://github.com/WhynotHAHA), [Lingxiao Lu](https://github.com/pokaaa), [Xinhao Tao](https://github.com/taoxinhao13), [Junyan Cao](https://github.com/cjy-4), [Haonan Zhao](https://github.com/nononononno), [Jiaxuan Chen](https://github.com/csdahunzi), [Junqi You](https://github.com/dhmbb2).
 
 ## License
 

libcom-0.1.0/README.md

@@ -0,0 +1,80 @@
+<div align="center">
+</br>
+<img src="https://raw.githubusercontent.com/bcmi/libcom/main/resources/LOGO.png" width="200" />
+
+</div>
+
+<h1 align="center">libcom: everything about image composition</h1>
+
+</br>
+
+[![PyPI](https://img.shields.io/pypi/v/libcom)](https://pypi.org/project/libcom)
+[![Downloads](https://static.pepy.tech/badge/libcom)](https://pepy.tech/project/libcom)
+[![Hits](https://hits.sh/github.com/bcmi/libcom.svg?label=views)](https://hits.sh/github.com/bcmi/libcom/)
+[![Static Badge](https://img.shields.io/badge/Image%20Composition%20Demo-Green)](https://bcmi.sjtu.edu.cn/home/niuli/demo_image_composition/)
+[![Static Badge](https://img.shields.io/badge/Resources-Awesome-green?style=flat&label=Resources)](https://github.com/bcmi/Awesome-Image-Composition)
+[![Static Badge](https://img.shields.io/badge/survey-arxiv%3A2106.14490-red)](https://arxiv.org/pdf/2106.14490.pdf)
+
+**We co-founded a startup company [miguo.ai](https://miguocomics.com/), dedicated to accelerating the production of comics and animations using AIGC technology. If you are seeking internship or full-time positions, please feel free to send your resume to hr@miguocomics.com.**
+
+---
+
+## Introduction
+**_libcom_ (the library of image composition) is an image composition toolbox.** The goal of image composition (object insertion) is inserting one foreground into a background image to get a realistic composite image, by addressing the inconsistencies (appearance, geometry, and semantic inconsistency) between foreground and background. Generally speaking, image composition could be used to combine the visual elements from different images.
+<div align="center">
+</br>
+<img src="https://raw.githubusercontent.com/bcmi/libcom/main/resources/image_composition_task.gif" width="600" />
+</div>
+
+**_libcom_ covers a diversity of related tasks in the field of image composition**, including image blending, standard/painterly image harmonization, shadow generation, object placement, generative composition, quality evaluation, *etc*. For each task, we integrate one or two selected methods considering both efficiency and effectiveness. The selected methods will be continuously updated upon the emergence of better methods. 
+
+**The ultimate goal of this library is solving all the problems related to image composition with simple `import libcom`.**
+
+### Main Functions
+
+- **get_composite_image** generates composite images using naive copy-and-paste followed by image blending.
+- **OPAScoreModel** [[OPA]](https://github.com/bcmi/Object-Placement-Assessment-Dataset-OPA) evaluates the rationality of foreground object placement in a composite image.
+- **FOPAHeatMapModel** [[FOPA]](https://github.com/bcmi/FOPA-Fast-Object-Placement-Assessment) can predict the rationality scores for all locations/scales given a background-foreground pair, and output the composite image with optimal location/scale.  
+- **color_transfer** adjusts the foreground color to match the background using traditional color transfer method.
+- **ImageHarmonizationModel** [[CDTNet]](https://github.com/bcmi/CDTNet-High-Resolution-Image-Harmonization)  [[PCTNet]](https://github.com/rakutentech/PCT-Net-Image-Harmonization) adjusts the foreground illumination to be compatible the background given photorealistic background and photorealistic foreground.
+- **PainterlyHarmonizationModel** [[PHDNet]](https://github.com/bcmi/PHDNet-Painterly-Image-Harmonization)  [[PHDiffusion]](https://github.com/bcmi/PHDiffusion-Painterly-Image-Harmonization) adjusts the foreground style to be compatible with the background given artistic background and photorealistic foreground.
+- **HarmonyScoreModel** [[BargainNet]](https://github.com/bcmi/BargainNet-Image-Harmonization) evaluates the harmony level between foreground and background in a composite image.
+- **InharmoniousLocalizationModel** [[MadisNet]](https://github.com/bcmi/MadisNet-Inharmonious-Region-Localization) localizes the inharmonious region in a synthetic image.
+- **FOSScoreModel** [[DiscoFOS]](https://github.com/bcmi/Foreground-Object-Search-Dataset-FOSD) evaluates the compatibility between foreground and background in a composite image in terms of geometry and semantics.
+- **ShadowGenerationModel** [[GPSDiffusion]](https://github.com/bcmi/GPSDiffusion-Object-Shadow-Generation) generates plausible shadow for the inserted object in a composite image. This model is unstable and you can pick the most satisfactory one from multiple generated results. 
+- **ControlComModel** [[ControlCom]](https://github.com/bcmi/ControlCom-Image-Composition) is a generative image composition model which unifies image blending and image harmonization. The pose and view of foreground stay unchanged. 
+- **MureObjectStitchModel** [[MureObjectStitch]](https://github.com/bcmi/MureObjectStitch-Image-Composition) is another generative image composition model which can adjust the pose and view of foreground. It supports multiple reference images of one foreground object. If you have a few images containing the foreground object, we suggest finetuning MureObjectStitch using these images for better detail preservation. 
+
+
+**For the detailed method descriptions, code examples, visualization results, and performance comments, please refer to our [[documents]](https://libcom.readthedocs.io/en/latest/).** If the model performance is not satisfactory, you can finetune the pretrained model on your own dataset using the source repository and replace the original model. 
+
+## Requirements
+
+The main branch is built on the Linux system with **Python 3.8** and **PyTorch>=1.10.1**. For other dependencies, please refer to [[conda_env]](https://github.com/bcmi/libcom/blob/main/requirements/libcom.yaml) and [[runtime_dependencies]](https://github.com/bcmi/libcom/blob/main/requirements/runtime.txt).
+
+## Get Started
+Please refer to [[Installation]](https://github.com/bcmi/libcom/blob/main/docs/get_started.md) for installation instructions and [[documents]](https://libcom.readthedocs.io/en/latest/) for user guidance.
+
+## Contributors
+- Institution: [Brain-like Computing and Machine Intelligence (BCMI) Lab](https://bcmi.sjtu.edu.cn/).
+- Project Initiator & Team Manager: [Li Niu](https://www.ustcnewly.com/index.html). 
+- Architect & Lead Developer: [Bo Zhang](https://bo-zhang-cs.github.io/), [Yujie Zhou](https://github.com/YujieOuO).
+- Documentation Manager: [Jiacheng Sui](https://github.com/charlessjc).
+- Module Developers: [Jiacheng Sui](https://github.com/charlessjc), [Binjie Gao](https://github.com/WhynotHAHA), [Lingxiao Lu](https://github.com/pokaaa), [Xinhao Tao](https://github.com/taoxinhao13), [Junyan Cao](https://github.com/cjy-4), [Haonan Zhao](https://github.com/nononononno), [Jiaxuan Chen](https://github.com/csdahunzi), [Junqi You](https://github.com/dhmbb2).
+
+## License
+
+This project is released under the [Apache 2.0 license](https://github.com/bcmi/libcom/blob/main/LICENSE).
+
+## Bibtex
+
+If you use our toolbox, please cite our survey paper using the following BibTeX  [[arxiv](https://arxiv.org/pdf/2106.14490.pdf)]:
+
+```
+@article{niu2021making,
+  title={Making images real again: A comprehensive survey on deep image composition},
+  author={Niu, Li and Cong, Wenyan and Liu, Liu and Hong, Yan and Zhang, Bo and Liang, Jing and Zhang, Liqing},
+  journal={arXiv preprint arXiv:2106.14490},
+  year={2021}
+}
+```

{libcom-0.0.1.post10 → libcom-0.1.0}/libcom/__init__.py

@@ -9,11 +9,12 @@ from .fopa_heat_map import *
 from .fos_score import *
 from .controllable_composition import *
 from .shadow_generation import *
+from .objectstitch import *
 
 __all__ = [
     'color_transfer', 'get_composite_image', 'OPAScoreModel', 
     'HarmonyScoreModel', 'InharmoniousLocalizationModel',
     'ImageHarmonizationModel', 'PainterlyHarmonizationModel',
     'FOPAHeatMapModel', 'FOSScoreModel', 'ControlComModel',
-    'ShadowGenerationModel'
+    'ShadowGenerationModel', 'Mure_ObjectStitchModel'
 ]

{libcom-0.0.1.post10 → libcom-0.1.0}/libcom/controllable_composition/controllable_composition.py

@@ -29,8 +29,9 @@ from .source.ControlCom.ldm.models.diffusion.plms import PLMSSampler
 from .source.ControlCom.ldm.data.open_images_control import get_tensor, get_tensor_clip, get_bbox_tensor, bbox2mask, tensor2numpy
 
 cur_dir   = os.path.dirname(os.path.abspath(__file__))
-model_set = ['ControlCom'] 
-task_set  = ['blending', 'harmonization'] # 'viewsynthesis', 'composition'
+model_dir = os.environ.get('LIBCOM_MODEL_DIR',cur_dir)
+model_set = ['ControlCom', 'ControlCom_blend_harm', 'ControlCom_view_comp'] 
+task_set  = ['blending', 'harmonization', 'viewsynthesis', 'composition']
 
 class ControlComModel:
     """
@@ -38,7 +39,7 @@ class ControlComModel:
 
     Args:
         device (str | torch.device): gpu id
-        model_type (str): predefined model type
+        model_type (str): predefined model type. "ControlCom" refers to the version trained on all four tasks comprehensively, covering a wide range of domains or objectives. "ControlCom_blend_harm" fine-tunes the "full" version specifically for image blending and harmonization tasks. "ControlCom_view_comp" fine-tunes the "full" version to excel in view synthesis and generative composition tasks.
         kwargs (dict): sampler='ddim' (default) or 'plms', other parameters for building model
     
     Examples:
@@ -53,7 +54,7 @@ class ControlComModel:
         >>>     fg_img  = test_dir + 'foreground/' + img_names[i]
         >>>     bbox    = bboxes[i]
         >>>     mask    = test_dir + 'foreground_mask/' + img_names[i]
-        >>>     net     = ControlComModel(device=0)
+        >>>     net     = ControlComModel(device=0, model_type="ControlCom")
         >>>     comp    = net(bg_img, fg_img, bbox, mask, task=['blending', 'harmonization'])
         >>>     bg_img  = draw_bbox_on_image(bg_img, bbox)
         >>>     grid_img = make_image_grid([bg_img, fg_img, comp[0], comp[1]])
@@ -73,7 +74,7 @@ class ControlComModel:
         self.model_type = model_type
         self.option = kwargs
         
-        weight_path = os.path.join(cur_dir, 'pretrained_models', 'ControlCom.pth')
+        weight_path = os.path.join(cur_dir, 'pretrained_models', f'{self.model_type}.pth')
         download_pretrained_model(weight_path)
         
         self.device = check_gpu_device(device)
@@ -84,7 +85,7 @@ class ControlComModel:
         pl_sd  = torch.load(weight_path, map_location="cpu")
         sd     = pl_sd["state_dict"]
         config = OmegaConf.load(os.path.join(cur_dir, 'source/ControlCom/configs/controlcom.yaml'))
-        clip_path = os.path.join(cur_dir, '../shared_pretrained_models', 'openai-clip-vit-large-patch14')
+        clip_path = os.path.join(model_dir, '../shared_pretrained_models', 'openai-clip-vit-large-patch14')
         download_entire_folder(clip_path)
         config.model.params.cond_stage_config.params.version = clip_path
         model  = instantiate_from_config(config.model)
@@ -181,6 +182,10 @@ class ControlComModel:
                 indicator.append([0,0])
             elif t == 'harmonization':
                 indicator.append([1,0])
+            elif t == 'viewsynthesis':
+                indicator.append([0,1])
+            else:
+                indicator.append([1,1])
         return indicator
                 
     @torch.no_grad()
@@ -195,7 +200,7 @@ class ControlComModel:
             foreground_image (str | numpy.ndarray): The path to foreground image or the background image in ndarray form.
             bbox (list): The bounding box which indicates the foreground's location in the background. [x1, y1, x2, y2].
             foreground_mask (None | str | numpy.ndarray): Mask of foreground image which indicates the foreground object region in the foreground image. default: None.
-            task (str | list of str): 'blending', 'harmonization', ['blending', 'harmonization']: Task types, including image blending, image harmonization. default: 'blending'.
+            task (str | list of str): 'blending', 'harmonization', 'viewsynthesis', 'composition' | ['blending', 'harmonization']: Task types, including image blending, image harmonization, view synthesis, generative composition. default: 'blending'.
             num_samples (int): Number of images to be generated for each task. default: 1.
             sample_steps (int): Number of denoising steps. The recommended setting is 25 for PLMS sampler and 50 for DDIM sampler. default: 50.
             guidance_scale (int): Scale in classifier-free guidance (minimum: 1; maximum: 20). default: 5.

{libcom-0.0.1.post10 → libcom-0.1.0}/libcom/controllable_composition/source/ControlCom/configs/controlcom.yaml

@@ -1,6 +1,6 @@
 model:
   base_learning_rate: 1.0e-05
-  target: ldm.models.diffusion.ddpm.LatentDiffusion
+  target: libcom.controllable_composition.source.ControlCom.ldm.models.diffusion.ddpm.LatentDiffusion
   params:
     linear_start: 0.00085
     linear_end: 0.012
@@ -41,7 +41,7 @@ model:
     use_guidance: true
     local_uncond: same
     scheduler_config:
-      target: ldm.lr_scheduler.LambdaLinearScheduler
+      target: libcom.controllable_composition.source.ControlCom.ldm.lr_scheduler.LambdaLinearScheduler
       params:
         warm_up_steps:
         - 1000
@@ -54,7 +54,7 @@ model:
         f_min:
         - 1.0
     unet_config:
-      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      target: libcom.controllable_composition.source.ControlCom.ldm.modules.diffusionmodules.openaimodel.UNetModel
       params:
         image_size: 32
         in_channels: 11
@@ -78,7 +78,7 @@ model:
         legacy: false
         add_conv_in_front_of_unet: false
         local_encoder_config:
-          conditioning_key: ldm.modules.local_module.LocalRefineBlock
+          conditioning_key: libcom.controllable_composition.source.ControlCom.ldm.modules.local_module.LocalRefineBlock
           add_position_emb: false
           roi_size: 16
           context_dim: 1024
@@ -89,7 +89,7 @@ model:
           add_in_decoder: true
           add_before_crossattn: false
     first_stage_config:
-      target: ldm.models.autoencoder.AutoencoderKL
+      target: libcom.controllable_composition.source.ControlCom.ldm.models.autoencoder.AutoencoderKL
       params:
         embed_dim: 4
         monitor: val/rec_loss
@@ -111,7 +111,7 @@ model:
         lossconfig:
           target: torch.nn.Identity
     cond_stage_config:
-      target: ldm.modules.encoders.modules.FrozenCLIPImageEmbedder
+      target: libcom.controllable_composition.source.ControlCom.ldm.modules.encoders.modules.FrozenCLIPImageEmbedder
       params:
         version: openai-clip-vit-large-patch14
         local_hidden_index: 12

libcom-0.1.0/libcom/controllable_composition/source/ControlCom/ldm/data/datamodule.py

@@ -0,0 +1,128 @@
+import argparse, os, sys, datetime, glob, importlib, csv
+import numpy as np
+import time
+import torch
+import torchvision
+import pytorch_lightning as pl
+
+from packaging import version
+from omegaconf import OmegaConf
+from torch.utils.data import random_split, DataLoader, Dataset, Subset
+from functools import partial
+from PIL import Image
+
+from pytorch_lightning import seed_everything
+from pytorch_lightning.trainer import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint, Callback, LearningRateMonitor
+from pytorch_lightning.utilities.rank_zero import rank_zero_only
+from pytorch_lightning.utilities import rank_zero_info
+
+from libcom.controllable_composition.source.ControlCom.ldm.data.base import Txt2ImgIterableBaseDataset
+from libcom.controllable_composition.source.ControlCom.ldm.util import instantiate_from_config
+import socket
+from pytorch_lightning.plugins.environments import ClusterEnvironment,SLURMEnvironment
+
+class WrappedDataset(Dataset):
+    """Wraps an arbitrary object with __len__ and __getitem__ into a pytorch dataset"""
+
+    def __init__(self, dataset):
+        self.data = dataset
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        return self.data[idx]
+    
+def worker_init_fn(_):
+    worker_info = torch.utils.data.get_worker_info()
+
+    dataset = worker_info.dataset
+    worker_id = worker_info.id
+
+    if isinstance(dataset, Txt2ImgIterableBaseDataset):
+        split_size = dataset.num_records // worker_info.num_workers
+        # reset num_records to the true number to retain reliable length information
+        dataset.sample_ids = dataset.valid_ids[worker_id * split_size:(worker_id + 1) * split_size]
+        current_id = np.random.choice(len(np.random.get_state()[1]), 1)
+        return np.random.seed(np.random.get_state()[1][current_id] + worker_id)
+    else:
+        return np.random.seed(np.random.get_state()[1][0] + worker_id)
+
+class DataModuleFromConfig(pl.LightningDataModule):
+    def __init__(self, batch_size, num_workers=None, train=None, validation=None, test=None, predict=None,
+                 wrap=False, shuffle_test_loader=False, use_worker_init_fn=False,
+                 shuffle_val_dataloader=False):
+        super().__init__()
+        self.batch_size = batch_size
+        self.dataset_configs = dict()
+        self.num_workers = num_workers if num_workers is not None else batch_size
+        self.use_worker_init_fn = use_worker_init_fn
+        if train is not None:
+            self.dataset_configs["train"] = train
+            self.train_dataloader = self._train_dataloader
+        if validation is not None:
+            self.dataset_configs["validation"] = validation
+            self.val_dataloader = partial(self._val_dataloader, shuffle=shuffle_val_dataloader)
+        if test is not None:
+            self.dataset_configs["test"] = test
+            self.test_dataloader = partial(self._test_dataloader, shuffle=shuffle_test_loader)
+        if predict is not None:
+            self.dataset_configs["predict"] = predict
+            self.predict_dataloader = self._predict_dataloader
+        self.wrap = wrap
+
+    def prepare_data(self):
+        for data_cfg in self.dataset_configs.values():
+            instantiate_from_config(data_cfg)
+
+    def setup(self, stage=None):
+        self.datasets = dict(
+            (k, instantiate_from_config(self.dataset_configs[k]))
+            for k in self.dataset_configs)
+        if self.wrap:
+            for k in self.datasets:
+                self.datasets[k] = WrappedDataset(self.datasets[k])
+
+    def _train_dataloader(self):
+        is_iterable_dataset = isinstance(self.datasets['train'], Txt2ImgIterableBaseDataset)
+        if is_iterable_dataset or self.use_worker_init_fn:
+            init_fn = worker_init_fn
+        else:
+            init_fn = None
+        return DataLoader(self.datasets["train"], batch_size=self.batch_size,
+                          num_workers=self.num_workers, shuffle=False if is_iterable_dataset else True,
+                          worker_init_fn=init_fn)
+
+    def _val_dataloader(self, shuffle=False):
+        if isinstance(self.datasets['validation'], Txt2ImgIterableBaseDataset) or self.use_worker_init_fn:
+            init_fn = worker_init_fn
+        else:
+            init_fn = None
+        return DataLoader(self.datasets["validation"],
+                          batch_size=self.batch_size,
+                          num_workers=self.num_workers,
+                          worker_init_fn=init_fn,
+                          shuffle=shuffle)
+
+    def _test_dataloader(self, shuffle=False):
+        is_iterable_dataset = isinstance(self.datasets['train'], Txt2ImgIterableBaseDataset)
+        if is_iterable_dataset or self.use_worker_init_fn:
+            init_fn = worker_init_fn
+        else:
+            init_fn = None
+
+        # do not shuffle dataloader for iterable dataset
+        shuffle = shuffle and (not is_iterable_dataset)
+
+        return DataLoader(self.datasets["test"], batch_size=self.batch_size,
+                          num_workers=self.num_workers, worker_init_fn=init_fn, shuffle=shuffle)
+
+    def _predict_dataloader(self, shuffle=False):
+        if isinstance(self.datasets['predict'], Txt2ImgIterableBaseDataset) or self.use_worker_init_fn:
+            init_fn = worker_init_fn
+        else:
+            init_fn = None
+        return DataLoader(self.datasets["predict"], batch_size=self.batch_size,
+                          num_workers=self.num_workers, worker_init_fn=init_fn)
+