Update gradio_app.py

gradio_app.py

@@ -282,8 +282,9 @@ def inference(id_image, hair_image):
     # Balding
     bald_id_path = "gradio_outputs/bald_id.png"
     cv2.imwrite(bald_id_path, cv2.cvtColor(aligned_id, cv2.COLOR_RGB2BGR))
+    print("start bald", flush=True)
     bald_head(bald_id_path, bald_id_path)
-
+    print("done bald", flush=True)
     # Resolve trained model dir
     trained_model_dir = os.path.abspath("trained_model") if os.path.isdir("trained_model") else None
     if trained_model_dir is None and os.path.isdir("pretrain"):
@@ -365,14 +366,14 @@ def inference(id_image, hair_image):
     ).to(device)
     state_dict4 = torch.load(os.path.join(args.model_path, "pytorch_model_2.bin"), map_location="cpu")
     Hair_Encoder.load_state_dict(state_dict4, strict=False)
-
+    print("start sd", flush=True)
     # Run inference
     log_validation(
         vae, tokenizer, image_encoder, denoising_unet,
         args, device, logger,
         cc_projection, controlnet, Hair_Encoder
     )
-
+    print("done sd", flush=True)
     output_video = os.path.join(args.output_dir, "validation", "generated_video_0.mp4")
 
     # Extract frames for slider preview
@@ -565,6 +566,6 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate"),
 
 
 if __name__ == "__main__":
-    demo.queue().launch(server_name="0.0.0.0", server_port=7860)
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860,show_error=True)