Fix #2

app.py

@@ -57,6 +57,7 @@ def main():
                     input_image = gr.Image(label='Input Pose Image',
                                            type='pil',
                                            elem_id='input-image')
+                    pose_data = gr.Variable()
                 with gr.Row():
                     paths = sorted(pathlib.Path('pose_images').glob('*.png'))
                     example_images = gr.Dataset(components=[input_image],
@@ -112,18 +113,22 @@ Note: Currently, only 5 types of textures are supported, i.e., pure color, strip
         gr.Markdown(FOOTER)
 
         input_image.change(fn=model.process_pose_image,
-                           inputs=[input_image],
-                           outputs=None)
+                           inputs=input_image,
+                           outputs=pose_data)
         generate_label_button.click(fn=model.generate_label_image,
-                                    inputs=[shape_text],
-                                    outputs=[label_image])
+                                    inputs=[
+                                        pose_data,
+                                        shape_text,
+                                    ],
+                                    outputs=label_image)
         generate_human_button.click(fn=model.generate_human,
                                     inputs=[
+                                        label_image,
                                         texture_text,
                                         sample_steps,
                                         seed,
                                     ],
-                                    outputs=[result])
+                                    outputs=result)
         example_images.click(fn=set_example_image,
                              inputs=example_images,
                              outputs=example_images.components)

model.py

@@ -98,29 +98,32 @@ class Model:
         result = np.asarray(result[0, :, :, :], dtype=np.uint8)
         return result
 
-    def process_pose_image(self, pose_image: PIL.Image.Image) -> None:
+    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, shape_text: str) -> np.ndarray:
+    def generate_label_image(self, pose_data: torch.Tensor,
+                             shape_text: str) -> np.ndarray:
+        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
 
-        mask = colored_segm.copy()
+    def generate_human(self, label_image: np.ndarray, texture_text: str,
+                       sample_steps: int, seed: int) -> np.ndarray:
+        mask = label_image.copy()
         seg_map = self.process_mask(mask)
         self.model.segm = torch.from_numpy(seg_map).unsqueeze(0).unsqueeze(
             0).to(self.model.device)
         self.model.generate_quantized_segm()
-        return colored_segm
 
-    def generate_human(self, texture_text: str, sample_steps: int,
-                       seed: int) -> np.ndarray:
         set_random_seed(seed)
 
         texture_attributes = generate_texture_attributes(texture_text)