updates

app.py

@@ -22,23 +22,21 @@ with gr.Blocks() as demo:
             submit_btn = gr.Button(value="Generate")
         with gr.Column():
             animation = gr.Video(label="Result")
-            frames = gr.Gallery(type="pil", label="Frames")
+            frames = gr.Gallery(type="pil", label="Frames", format="png")
 
     submit_btn.click(
         run_app, inputs=[char_imgs, mocap, tr_steps, inf_steps, fps, remove_bg, resize_inputs], outputs=[animation, frames]
     )
     
     train_btn.click(
-        run_train, inputs=[char_imgs, tr_steps, remove_bg, resize_inputs, modelId], outputs=[]
+        run_train, inputs=[char_imgs, tr_steps, modelId, remove_bg, resize_inputs], outputs=[]
     )
     
     inference_btn.click(
-        run_inference, inputs=[char_imgs, mocap, inf_steps, fps, remove_bg, resize_inputs, modelId], outputs=[animation, frames]
+        run_inference, inputs=[char_imgs, mocap, tr_steps, inf_steps, fps, modelId, remove_bg, resize_inputs], outputs=[animation, frames]
     )
 
 
-
-
 demo.launch(share=True)
 
 

main.py

@@ -57,6 +57,8 @@ import rembg
 import uuid
 import gc
 from numba import cuda
+import requests
+import uuid
 
 from huggingface_hub import hf_hub_download
 
@@ -76,8 +78,34 @@ fps = 12
 
 debug = False
 save_model = True
+should_gen_vid = False
 max_batch_size = 8
 
+
+def save_temp_imgs(imgs):
+    for img in imgs:
+
+        img_name = str(uuid.uuid4())+'.png'
+        img.save(img_name)
+        print(img_name)
+
+        url = 'https://tmpfiles.org/api/v1/upload'
+        data_payload = {'file': img_name}
+
+        try:
+            response = requests.post(url, data=data_payload)
+
+            # Check for successful response (status code 200)
+            response.raise_for_status()
+
+            # Print the server's response
+            print("Status Code:", response.status_code)
+            print("Response JSON:", response.json())
+
+        except requests.exceptions.RequestException as e:
+            print(f"An error occurred: {e}")
+
+
 # Pose detection ==============================================================================================
 
 def load_models():
@@ -712,7 +740,11 @@ def train(modelId, in_image, in_pose, train_images, train_poses, train_steps, pc
     accelerator.wait_for_everyone()
     accelerator.end_training()
 
-    
+    sd_model.unet.cpu()
+    sd_model.cpu()
+    del vae
+    del image_encoder_p
+    del image_encoder_g
 
     if save_model: #if global_steps % args.checkpointing_steps == 0 or global_steps == args.max_train_steps:
         print('saving', modelId)
@@ -724,13 +756,14 @@ def train(modelId, in_image, in_pose, train_images, train_poses, train_steps, pc
         print(list(sd_model.state_dict().keys())[:20])
         torch.save(checkpoint_state_dict, modelId+".pt")
         
+        del sd_model
         gc.collect()
         torch.cuda.empty_cache()
-        #device = cuda.get_current_device()
-        #device.reset()
         print('done train')
+        print(torch.cuda.memory_allocated()/1024**2)
         return
     
+    del sd_model
     gc.collect()
     torch.cuda.empty_cache()    
     return {k: v.cpu() for k, v in sd_model.state_dict().items()}
@@ -953,8 +986,18 @@ def inference(modelId, in_image, in_pose, target_poses, inference_steps, finetun
         results.append(result)
         progress_bar.update(1)
 
+    del unet
+    del vae
+    del image_encoder
+    del image_proj_model
+    del pose_proj_model
+
+    if not save_model:
+        del finetuned_model
+
     gc.collect()
     torch.cuda.empty_cache()
+    print(torch.cuda.memory_allocated()/1024**2)
         
     return results
         
@@ -1006,7 +1049,7 @@ def run(images, video_path, train_steps=100, inference_steps=10, fps=12, bg_remo
     return results
 
 
-def run_train(images, train_steps=100, bg_remove=False, resize_inputs=True, modelId="fine_tuned_pcdms"):
+def run_train(images, train_steps=100, modelId="fine_tuned_pcdms", bg_remove=True, resize_inputs=True):
     finetune=True
     is_app=True
     images = [img[0] for img in images]
@@ -1023,21 +1066,29 @@ def run_train(images, train_steps=100, bg_remove=False, resize_inputs=True, mode
     train(modelId, in_img, in_pose, train_imgs, train_poses, train_steps, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet, finetune, is_app)
 
 
-def run_inference(images, video_path, inference_steps=10, fps=12, bg_remove=False, resize_inputs=True, modelId="fine_tuned_pcdms"):
+def run_inference(images, video_path, train_steps=100, inference_steps=10, fps=12, modelId="fine_tuned_pcdms", bg_remove=True, resize_inputs=True):
+    finetune=True
     is_app=True
-    images = [img[0] for img in images]
-    in_img = images[0]
+    
     
     dwpose, rembg_session, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet = load_models()
     
+    if not os.path.exists(modelId+".pt"):
+        run_train(images, train_steps, modelId, bg_remove, resize_inputs)
+    
+    images = [img[0] for img in images]
+    in_img = images[0]
+
     target_poses, in_pose = prepare_inputs_inference(in_img, video_path, fps, dwpose, 'target', is_app)
     
     results = inference(modelId, in_img, in_pose, target_poses, inference_steps, None, vae, unet, image_encoder_p, is_app)
-    
-    if debug:
-        gen_vid(results, out_vid+'.mp4', fps, 'mp4')
-    else:
-        gen_vid(results, out_vid+'.webm', fps, 'webm')
+    save_temp_imgs(results)
+
+    if should_gen_vid:
+        if debug:
+            gen_vid(results, out_vid+'.mp4', fps, 'mp4')
+        else:
+            gen_vid(results, out_vid+'.webm', fps, 'webm')
     
     print("Done!")