add more

main.py

@@ -94,7 +94,7 @@ global auto_pipeline, text2img_pipeline, s3_client
 
 if __name__ == "__main__":
     auto_pipeline = AudioPipeline(audio_text_path='/home/user/app/dedup_audio_text_80.json',
-                                  audio_text_embeddings_path='/home/user/app/audio_text_embeddings.safetensors')
+                                  audio_text_embeddings_path='/home/user/app/audio_text_embeddings_cpu.safetensors')
     text2img_pipeline = init_text2img_pipe()
     s3_client = Minio(
         os.environ.get("S3_ENDPOINT"),

test.py

@@ -1,7 +1,15 @@
-from optimum.onnxruntime import ORTStableDiffusionPipeline
+from safetensors import safe_open
+from safetensors.torch import save_file
 
-model_id = "eatdianatoday/yiwu"
-pipe = ORTStableDiffusionPipeline.from_pretrained(model_id, export=True)
-prompt = "a photo of an astronaut riding a horse on mars"
-images = pipe(prompt).images[0]
-pipe.save_pretrained("/root/autodl-tmp/onnx-novelai-diffusion")
+tensors = {}
+
+with safe_open('audio_text_embeddings.safetensors', framework="pt", device=0) as f:
+    for k in f.keys():
+        tensors[k] = f.get_tensor(k)
+
+text_embed = tensors["text_embed"].detach().cpu()
+
+tensors = {
+    "text_embed": text_embed,
+}
+save_file(tensors, "/root/autodl-tmp/audio_text_embeddings_ cpu.safetensors")

warm_up.py

@@ -11,7 +11,7 @@ if __name__ == "__main__":
         "(bad proportions:1.331)",
     ]
     pipeline = AudioPipeline(audio_text_path='/home/user/app/dedup_audio_text_80.json',
-                             audio_text_embeddings_path='/home/user/app/audio_text_embeddings.safetensors')
+                             audio_text_embeddings_path='/home/user/app/audio_text_embeddings_cpu.safetensors')
     text = pipeline.audio2txt('/home/user/app/demo.wav')
 
     text2img_pipeline = init_text2img_pipe()