Upload src/app/pipeline.py with huggingface_hub

src/app/pipeline.py

@@ -0,0 +1,96 @@
+import os
+import sys
+import pickle
+import shutil
+from threading import Timer
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+app_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+
+from utils.inference_utils import gen_prog_ind
+from utils.constants import TO_24
+from inference import inference
+
+from omegaconf import OmegaConf
+from inference.joint2smplx import process_file
+
+def delete_folder(folder_path):
+    if os.path.exists(folder_path):
+        shutil.rmtree(folder_path)
+
+def rand_folder_name():
+    import time
+    return str(time.time()).replace('.', '')
+
+def pipeline(data, models, device, diffuser, **kwargs):
+    from app.process_data import get_a_sample
+    from app.setup_models import text_embedder, test_configs, normalize, denormalize
+
+    len_data = min(data['source']['transl'].shape[0]//((kwargs['SEQLEN']-2)*2), 4)
+    if len_data < 4:
+        return None # not enough data
+    
+    joints_orig = get_a_sample(data['source'],
+                                len_data,
+                                kwargs['SEQLEN'],
+                                smplx_pth=os.path.abspath(os.path.join(app_root, 'deps/smplx/models'))
+                                ).to(device)
+    
+    joints_orig = normalize(joints_orig)
+    
+    hint_text = data['text']
+
+    if data['prog_ind'] is None:
+        prog_ind = gen_prog_ind(num_cases=1, sublist_length=len_data)[0]
+    else:
+        prog_ind = data['prog_ind']
+    generated_samples, orig = inference.test_model(
+                                                models=models, 
+                                                diffuser=diffuser, 
+                                                normalizer=(normalize, denormalize), 
+                                                configs=test_configs, 
+                                                text_embedder=text_embedder, 
+                                                hint_text=hint_text, 
+                                                prog_ind=prog_ind, 
+                                                joint_orig=joints_orig,
+                                                All_one_model=data['All_one_model'],
+                                                model_type=data['model_type']
+                                            )
+    
+    generated_samples = generated_samples.reshape(1, -1, 28, 3)[..., TO_24, :].reshape(1, -1, 72)
+    orig = orig.reshape(1, -1, 28, 3)[..., TO_24, :].reshape(1, -1, 72)
+
+    combined_dict = {
+        'generated_samples': generated_samples,
+        'original_samples': orig, 
+        'text' : hint_text,
+    }
+    # return combined_dict
+
+    input_folder = os.path.join(app_root, rand_folder_name())
+    output_folder = os.path.join(app_root, rand_folder_name())
+
+    if not os.path.exists(input_folder):
+        os.makedirs(input_folder)
+    with open(os.path.join(input_folder, 'temp.pkl'), 'wb') as file:
+        pickle.dump(combined_dict, file)
+    
+    
+    j2s_config = OmegaConf.load(os.path.join(app_root, "configs/j2s.yaml"))
+
+    for file_name in os.listdir(input_folder):
+        if file_name.endswith('.pkl'):
+            process_file(file_path=input_folder, 
+                        file_name=file_name,
+                        save_path=output_folder,
+                        JointsToSMPLX_model_path=os.path.abspath(os.path.join(app_root, '..', j2s_config.JointsToSMPLX_model_path)),
+                        smplx_path=os.path.abspath(os.path.join(app_root, '..', j2s_config.smplx_path)),
+                        key_list = ['generated_samples'],
+                        # remenber to remove original samples when using app
+                        interp_s=j2s_config.interp_s, 
+                        )
+    
+    Timer(100, delete_folder, [input_folder]).start()
+    Timer(100, delete_folder, [output_folder]).start()
+    
+    return os.path.join(output_folder, 'generated_samples/temp.pkl')