Update generate.py

generate.py

@@ -7,6 +7,7 @@ from datetime import datetime
 from tqdm import tqdm
 
 def generate(args):
+    print("call generate")
     rank = int(os.getenv("RANK", 0))
     world_size = int(os.getenv("WORLD_SIZE", 1))
     local_rank = int(os.getenv("LOCAL_RANK", 0))
@@ -14,11 +15,11 @@ def generate(args):
     # Set device: use CPU if specified, else use GPU based on rank
     if args.t5_cpu or args.dit_fsdp:  # Use CPU if specified in arguments
         device = torch.device("cpu")
-        logging.info("Using CPU for model inference.")
+        print("Using CPU for model inference.")
     else:
         device = local_rank
         torch.cuda.set_device(local_rank)  # Ensure proper device assignment if using GPU
-        logging.info(f"Using GPU: {device}")
+        print(f"Using GPU: {device}")
 
     _init_logging(rank)
 
@@ -108,7 +109,7 @@ def generate(args):
             t5_cpu=args.t5_cpu,
         )
 
-        logging.info(f"Generating {'image' if 't2i' in args.task else 'video'} ...")
+        print(f"Generating {'image' if 't2i' in args.task else 'video'} ...")
         try:
             video = wan_t2v.generate(
                 args.prompt,
@@ -165,7 +166,7 @@ def generate(args):
             t5_cpu=args.t5_cpu,
         )
 
-        logging.info("Generating video ...")
+        print("Generating video ..6666666666666666")
         try:
             video = wan_i2v.generate(
                 args.prompt,