Upload 17 files

.gitattributes

@@ -1,35 +1,28 @@
+*.png filter=lfs diff=lfs merge=lfs -text
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
 *.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+pretrained_models

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "Text2Human"]
+	path = Text2Human
+	url = https://github.com/yumingj/Text2Human

.pre-commit-config.yaml

@@ -0,0 +1,36 @@
+exclude: ^(Text2Human|patch)
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.2.0
+  hooks:
+  - id: check-executables-have-shebangs
+  - id: check-json
+  - id: check-merge-conflict
+  - id: check-shebang-scripts-are-executable
+  - id: check-toml
+  - id: check-yaml
+  - id: double-quote-string-fixer
+  - id: end-of-file-fixer
+  - id: mixed-line-ending
+    args: ['--fix=lf']
+  - id: requirements-txt-fixer
+  - id: trailing-whitespace
+- repo: https://github.com/myint/docformatter
+  rev: v1.4
+  hooks:
+  - id: docformatter
+    args: ['--in-place']
+- repo: https://github.com/pycqa/isort
+  rev: 5.12.0
+  hooks:
+    - id: isort
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v0.991
+  hooks:
+    - id: mypy
+      args: ['--ignore-missing-imports']
+- repo: https://github.com/google/yapf
+  rev: v0.32.0
+  hooks:
+  - id: yapf
+    args: ['--parallel', '--in-place']

.style.yapf

@@ -0,0 +1,5 @@
+[style]
+based_on_style = pep8
+blank_line_before_nested_class_or_def = false
+spaces_before_comment = 2
+split_before_logical_operator = true

README.md

@@ -1,12 +1,13 @@
 ---
-title: Outfits
-emoji: 📈
+title: Text2Human
+emoji: 🏃
 colorFrom: purple
-colorTo: blue
+colorTo: gray
 sdk: gradio
-sdk_version: 3.41.2
+sdk_version: 3.36.1
 app_file: app.py
 pinned: false
+suggested_hardware: t4-small
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference

app.py

@@ -0,0 +1,140 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import os
+import pathlib
+import random
+import shlex
+import subprocess
+
+import gradio as gr
+import numpy as np
+
+if os.getenv('SYSTEM') == 'spaces':
+    import mim
+
+    mim.uninstall('mmcv-full', confirm_yes=True)
+    mim.install('mmcv-full==1.5.2', is_yes=True)
+
+    with open('patch') as f:
+        subprocess.run(shlex.split('patch -p1'), cwd='Text2Human', stdin=f)
+
+from model import Model
+
+DESCRIPTION = '''# [Text2Human](https://github.com/yumingj/Text2Human)
+
+You can modify sample steps and seeds. By varying seeds, you can sample different human images under the same pose, shape description, and texture description. The larger the sample steps, the better quality of the generated images. (The default value of sample steps is 256 in the original repo.)
+
+Label image generation step can be skipped. However, in that case, the input label image must be 512x256 in size and must contain only the specified colors.
+'''
+
+MAX_SEED = np.iinfo(np.int32).max
+
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+
+model = Model()
+
+with gr.Blocks(css='style.css') as demo:
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Row():
+        with gr.Column():
+            with gr.Row():
+                input_image = gr.Image(label='Input Pose Image',
+                                       type='pil',
+                                       elem_id='input-image')
+                pose_data = gr.State()
+            with gr.Row():
+                paths = sorted(pathlib.Path('pose_images').glob('*.png'))
+                gr.Examples(examples=[[path.as_posix()] for path in paths],
+                            inputs=input_image)
+
+            with gr.Row():
+                shape_text = gr.Textbox(
+                    label='Shape Description',
+                    placeholder=
+                    '''<gender>, <sleeve length>, <length of lower clothing>, <outer clothing type>, <other accessories1>, ...
+Note: The outer clothing type and accessories can be omitted.''')
+            with gr.Row():
+                gr.Examples(
+                    examples=[['man, sleeveless T-shirt, long pants'],
+                              ['woman, short-sleeve T-shirt, short jeans']],
+                    inputs=shape_text)
+            with gr.Row():
+                generate_label_button = gr.Button('Generate Label Image')
+
+        with gr.Column():
+            with gr.Row():
+                label_image = gr.Image(label='Label Image',
+                                       type='numpy',
+                                       elem_id='label-image')
+
+            with gr.Row():
+                texture_text = gr.Textbox(
+                    label='Texture Description',
+                    placeholder=
+                    '''<upper clothing texture>, <lower clothing texture>, <outer clothing texture>
+Note: Currently, only 5 types of textures are supported, i.e., pure color, stripe/spline, plaid/lattice, floral, denim.'''
+                )
+            with gr.Row():
+                gr.Examples(examples=[
+                    ['pure color, denim'],
+                    ['floral, stripe'],
+                ],
+                            inputs=texture_text)
+            with gr.Row():
+                sample_steps = gr.Slider(label='Sample Steps',
+                                         minimum=10,
+                                         maximum=300,
+                                         step=1,
+                                         value=256)
+            with gr.Row():
+                seed = gr.Slider(label='Seed',
+                                 minimum=0,
+                                 maximum=MAX_SEED,
+                                 step=1,
+                                 value=0)
+                randomize_seed = gr.Checkbox(label='Randomize seed',
+                                             value=True)
+            with gr.Row():
+                generate_human_button = gr.Button('Generate Human')
+
+        with gr.Column():
+            with gr.Row():
+                result = gr.Image(label='Result',
+                                  type='numpy',
+                                  elem_id='result-image')
+
+    input_image.change(
+        fn=model.process_pose_image,
+        inputs=input_image,
+        outputs=pose_data,
+    )
+    generate_label_button.click(
+        fn=model.generate_label_image,
+        inputs=[
+            pose_data,
+            shape_text,
+        ],
+        outputs=label_image,
+    )
+    generate_human_button.click(fn=randomize_seed_fn,
+                                inputs=[seed, randomize_seed],
+                                outputs=seed,
+                                queue=False).then(
+                                    fn=model.generate_human,
+                                    inputs=[
+                                        label_image,
+                                        texture_text,
+                                        sample_steps,
+                                        seed,
+                                    ],
+                                    outputs=result,
+                                )
+demo.queue(max_size=10).launch()

model.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+import pathlib
+import sys
+import zipfile
+
+import huggingface_hub
+import numpy as np
+import PIL.Image
+import torch
+
+sys.path.insert(0, 'Text2Human')
+
+from models.sample_model import SampleFromPoseModel
+from utils.language_utils import (generate_shape_attributes,
+                                  generate_texture_attributes)
+from utils.options import dict_to_nonedict, parse
+from utils.util import set_random_seed
+
+COLOR_LIST = [
+    (0, 0, 0),
+    (255, 250, 250),
+    (220, 220, 220),
+    (250, 235, 215),
+    (255, 250, 205),
+    (211, 211, 211),
+    (70, 130, 180),
+    (127, 255, 212),
+    (0, 100, 0),
+    (50, 205, 50),
+    (255, 255, 0),
+    (245, 222, 179),
+    (255, 140, 0),
+    (255, 0, 0),
+    (16, 78, 139),
+    (144, 238, 144),
+    (50, 205, 174),
+    (50, 155, 250),
+    (160, 140, 88),
+    (213, 140, 88),
+    (90, 140, 90),
+    (185, 210, 205),
+    (130, 165, 180),
+    (225, 141, 151),
+]
+
+
+class Model:
+    def __init__(self):
+        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        self.config = self._load_config()
+        self.config['device'] = device.type
+        self._download_models()
+        self.model = SampleFromPoseModel(self.config)
+        self.model.batch_size = 1
+
+    def _load_config(self) -> dict:
+        path = 'Text2Human/configs/sample_from_pose.yml'
+        config = parse(path, is_train=False)
+        config = dict_to_nonedict(config)
+        return config
+
+    def _download_models(self) -> None:
+        model_dir = pathlib.Path('pretrained_models')
+        if model_dir.exists():
+            return
+        path = huggingface_hub.hf_hub_download('yumingj/Text2Human_SSHQ',
+                                               'pretrained_models.zip')
+        model_dir.mkdir()
+        with zipfile.ZipFile(path) as f:
+            f.extractall(model_dir)
+
+    @staticmethod
+    def preprocess_pose_image(image: PIL.Image.Image) -> torch.Tensor:
+        image = np.array(
+            image.resize(
+                size=(256, 512),
+                resample=PIL.Image.Resampling.LANCZOS))[:, :, 2:].transpose(
+                    2, 0, 1).astype(np.float32)
+        image = image / 12. - 1
+        data = torch.from_numpy(image).unsqueeze(1)
+        return data
+
+    @staticmethod
+    def process_mask(mask: np.ndarray) -> np.ndarray:
+        if mask.shape != (512, 256, 3):
+            return None
+        seg_map = np.full(mask.shape[:-1], -1)
+        for index, color in enumerate(COLOR_LIST):
+            seg_map[np.sum(mask == color, axis=2) == 3] = index
+        if not (seg_map != -1).all():
+            return None
+        return seg_map
+
+    @staticmethod
+    def postprocess(result: torch.Tensor) -> np.ndarray:
+        result = result.permute(0, 2, 3, 1)
+        result = result.detach().cpu().numpy()
+        result = result * 255
+        result = np.asarray(result[0, :, :, :], dtype=np.uint8)
+        return result
+
+    def process_pose_image(self, pose_image: PIL.Image.Image) -> torch.Tensor:
+        if pose_image is None:
+            return
+        data = self.preprocess_pose_image(pose_image)
+        self.model.feed_pose_data(data)
+        return data
+
+    def generate_label_image(self, pose_data: torch.Tensor,
+                             shape_text: str) -> np.ndarray:
+        if pose_data is None:
+            return
+        self.model.feed_pose_data(pose_data)
+        shape_attributes = generate_shape_attributes(shape_text)
+        shape_attributes = torch.LongTensor(shape_attributes).unsqueeze(0)
+        self.model.feed_shape_attributes(shape_attributes)
+        self.model.generate_parsing_map()
+        self.model.generate_quantized_segm()
+        colored_segm = self.model.palette_result(self.model.segm[0].cpu())
+        return colored_segm
+
+    def generate_human(self, label_image: np.ndarray, texture_text: str,
+                       sample_steps: int, seed: int) -> np.ndarray:
+        if label_image is None:
+            return
+        mask = label_image.copy()
+        seg_map = self.process_mask(mask)
+        if seg_map is None:
+            return
+        self.model.segm = torch.from_numpy(seg_map).unsqueeze(0).unsqueeze(
+            0).to(self.model.device)
+        self.model.generate_quantized_segm()
+
+        set_random_seed(seed)
+
+        texture_attributes = generate_texture_attributes(texture_text)
+        texture_attributes = torch.LongTensor(texture_attributes)
+        self.model.feed_texture_attributes(texture_attributes)
+        self.model.generate_texture_map()
+
+        self.model.sample_steps = sample_steps
+        out = self.model.sample_and_refine()
+        res = self.postprocess(out)
+        return res

patch

@@ -0,0 +1,169 @@
+diff --git a/models/hierarchy_inference_model.py b/models/hierarchy_inference_model.py
+index 3116307..5de661d 100644
+--- a/models/hierarchy_inference_model.py
++++ b/models/hierarchy_inference_model.py
+@@ -21,7 +21,7 @@ class VQGANTextureAwareSpatialHierarchyInferenceModel():
+ 
+     def __init__(self, opt):
+         self.opt = opt
+-        self.device = torch.device('cuda')
++        self.device = torch.device(opt['device'])
+         self.is_train = opt['is_train']
+ 
+         self.top_encoder = Encoder(
+diff --git a/models/hierarchy_vqgan_model.py b/models/hierarchy_vqgan_model.py
+index 4b0d657..0bf4712 100644
+--- a/models/hierarchy_vqgan_model.py
++++ b/models/hierarchy_vqgan_model.py
+@@ -20,7 +20,7 @@ class HierarchyVQSpatialTextureAwareModel():
+ 
+     def __init__(self, opt):
+         self.opt = opt
+-        self.device = torch.device('cuda')
++        self.device = torch.device(opt['device'])
+         self.top_encoder = Encoder(
+             ch=opt['top_ch'],
+             num_res_blocks=opt['top_num_res_blocks'],
+diff --git a/models/parsing_gen_model.py b/models/parsing_gen_model.py
+index 9440345..15a1ecb 100644
+--- a/models/parsing_gen_model.py
++++ b/models/parsing_gen_model.py
+@@ -22,7 +22,7 @@ class ParsingGenModel():
+ 
+     def __init__(self, opt):
+         self.opt = opt
+-        self.device = torch.device('cuda')
++        self.device = torch.device(opt['device'])
+         self.is_train = opt['is_train']
+ 
+         self.attr_embedder = ShapeAttrEmbedding(
+diff --git a/models/sample_model.py b/models/sample_model.py
+index 4c60e3f..5265cd0 100644
+--- a/models/sample_model.py
++++ b/models/sample_model.py
+@@ -23,7 +23,7 @@ class BaseSampleModel():
+ 
+     def __init__(self, opt):
+         self.opt = opt
+-        self.device = torch.device('cuda')
++        self.device = torch.device(opt['device'])
+ 
+         # hierarchical VQVAE
+         self.decoder = Decoder(
+@@ -123,7 +123,7 @@ class BaseSampleModel():
+ 
+     def load_top_pretrain_models(self):
+         # load pretrained vqgan
+-        top_vae_checkpoint = torch.load(self.opt['top_vae_path'])
++        top_vae_checkpoint = torch.load(self.opt['top_vae_path'], map_location=self.device)
+ 
+         self.decoder.load_state_dict(
+             top_vae_checkpoint['decoder'], strict=True)
+@@ -137,7 +137,7 @@ class BaseSampleModel():
+         self.top_post_quant_conv.eval()
+ 
+     def load_bot_pretrain_network(self):
+-        checkpoint = torch.load(self.opt['bot_vae_path'])
++        checkpoint = torch.load(self.opt['bot_vae_path'], map_location=self.device)
+         self.bot_decoder_res.load_state_dict(
+             checkpoint['bot_decoder_res'], strict=True)
+         self.decoder.load_state_dict(checkpoint['decoder'], strict=True)
+@@ -153,7 +153,7 @@ class BaseSampleModel():
+ 
+     def load_pretrained_segm_token(self):
+         # load pretrained vqgan for segmentation mask
+-        segm_token_checkpoint = torch.load(self.opt['segm_token_path'])
++        segm_token_checkpoint = torch.load(self.opt['segm_token_path'], map_location=self.device)
+         self.segm_encoder.load_state_dict(
+             segm_token_checkpoint['encoder'], strict=True)
+         self.segm_quantizer.load_state_dict(
+@@ -166,7 +166,7 @@ class BaseSampleModel():
+         self.segm_quant_conv.eval()
+ 
+     def load_index_pred_network(self):
+-        checkpoint = torch.load(self.opt['pretrained_index_network'])
++        checkpoint = torch.load(self.opt['pretrained_index_network'], map_location=self.device)
+         self.index_pred_guidance_encoder.load_state_dict(
+             checkpoint['guidance_encoder'], strict=True)
+         self.index_pred_decoder.load_state_dict(
+@@ -176,7 +176,7 @@ class BaseSampleModel():
+         self.index_pred_decoder.eval()
+ 
+     def load_sampler_pretrained_network(self):
+-        checkpoint = torch.load(self.opt['pretrained_sampler'])
++        checkpoint = torch.load(self.opt['pretrained_sampler'], map_location=self.device)
+         self.sampler_fn.load_state_dict(checkpoint, strict=True)
+         self.sampler_fn.eval()
+ 
+@@ -397,7 +397,7 @@ class SampleFromPoseModel(BaseSampleModel):
+                         [185, 210, 205], [130, 165, 180], [225, 141, 151]]
+ 
+     def load_shape_generation_models(self):
+-        checkpoint = torch.load(self.opt['pretrained_parsing_gen'])
++        checkpoint = torch.load(self.opt['pretrained_parsing_gen'], map_location=self.device)
+ 
+         self.shape_attr_embedder.load_state_dict(
+             checkpoint['embedder'], strict=True)
+diff --git a/models/transformer_model.py b/models/transformer_model.py
+index 7db0f3e..4523d17 100644
+--- a/models/transformer_model.py
++++ b/models/transformer_model.py
+@@ -21,7 +21,7 @@ class TransformerTextureAwareModel():
+ 
+     def __init__(self, opt):
+         self.opt = opt
+-        self.device = torch.device('cuda')
++        self.device = torch.device(opt['device'])
+         self.is_train = opt['is_train']
+ 
+         # VQVAE for image
+@@ -317,10 +317,10 @@ class TransformerTextureAwareModel():
+     def sample_fn(self, temp=1.0, sample_steps=None):
+         self._denoise_fn.eval()
+ 
+-        b, device = self.image.size(0), 'cuda'
++        b = self.image.size(0)
+         x_t = torch.ones(
+-            (b, np.prod(self.shape)), device=device).long() * self.mask_id
+-        unmasked = torch.zeros_like(x_t, device=device).bool()
++            (b, np.prod(self.shape)), device=self.device).long() * self.mask_id
++        unmasked = torch.zeros_like(x_t, device=self.device).bool()
+         sample_steps = list(range(1, sample_steps + 1))
+ 
+         texture_mask_flatten = self.texture_tokens.view(-1)
+@@ -336,11 +336,11 @@ class TransformerTextureAwareModel():
+ 
+         for t in reversed(sample_steps):
+             print(f'Sample timestep {t:4d}', end='\r')
+-            t = torch.full((b, ), t, device=device, dtype=torch.long)
++            t = torch.full((b, ), t, device=self.device, dtype=torch.long)
+ 
+             # where to unmask
+             changes = torch.rand(
+-                x_t.shape, device=device) < 1 / t.float().unsqueeze(-1)
++                x_t.shape, device=self.device) < 1 / t.float().unsqueeze(-1)
+             # don't unmask somewhere already unmasked
+             changes = torch.bitwise_xor(changes,
+                                         torch.bitwise_and(changes, unmasked))
+diff --git a/models/vqgan_model.py b/models/vqgan_model.py
+index 13a2e70..9c840f1 100644
+--- a/models/vqgan_model.py
++++ b/models/vqgan_model.py
+@@ -20,7 +20,7 @@ class VQModel():
+     def __init__(self, opt):
+         super().__init__()
+         self.opt = opt
+-        self.device = torch.device('cuda')
++        self.device = torch.device(opt['device'])
+         self.encoder = Encoder(
+             ch=opt['ch'],
+             num_res_blocks=opt['num_res_blocks'],
+@@ -390,7 +390,7 @@ class VQImageSegmTextureModel(VQImageModel):
+ 
+     def __init__(self, opt):
+         self.opt = opt
+-        self.device = torch.device('cuda')
++        self.device = torch.device(opt['device'])
+         self.encoder = Encoder(
+             ch=opt['ch'],
+             num_res_blocks=opt['num_res_blocks'],

requirements.txt

@@ -0,0 +1,12 @@
+einops==0.6.1
+lpips==0.1.4
+mmcv-full==1.5.2
+mmsegmentation==0.24.1
+numpy==1.23.5
+openmim==0.1.5
+Pillow==9.5.0
+sentence-transformers==2.2.2
+tokenizers==0.13.3
+torch==1.11.0
+torchvision==0.12.0
+transformers==4.30.2

style.css

@@ -0,0 +1,16 @@
+h1 {
+  text-align: center;
+}
+#input-image {
+  max-height: 300px;
+}
+#label-image {
+  height: 300px;
+}
+#result-image {
+  height: 300px;
+}
+img#visitor-badge {
+  display: block;
+  margin: auto;
+}