0712-1532

diffusion.py

@@ -622,8 +622,11 @@ def generate_samples(model, num_new_img, max_num_img_per_gpu, rank, world_size):
         return None
 
 
-def sample(rank, world_size, model, num_new_img, max_num_img_per_gpu, return_dict):
-    samples = generate_samples(model, num_new_img, max_num_img_per_gpu, rank, world_size)
+def sample(rank, world_size, config, num_new_img, max_num_img_per_gpu, return_dict):
+    ddp_setup(rank, world_size)
+    ddpm21cm = DDPM21CM(config)
+
+    samples = generate_samples(ddpm21cm, num_new_img, max_num_img_per_gpu, rank, world_size)
 
     if rank == 0:
         return_dict['samples'] = samples
@@ -648,14 +651,11 @@ if __name__ == "__main__":
         filename = f"./outputs/model_state-N{num_image}-epoch9-device0"
         config.num_image = num_image
 
-        ddp_setup(rank, world_size)
-        ddpm21cm = DDPM21CM(config)
         # print("ddpm21cm = DDPM21CM(config)")
-
         manager = np.Manager()
         return_dict = manager.dict()
 
-        mp.spawn(sample, args=(world_size, ddpm21cm, num_new_img, max_num_img_per_gpu, return_dict), nprocs=world_size, join=True)
+        mp.spawn(sample, args=(world_size, config, num_new_img, max_num_img_per_gpu, return_dict), nprocs=world_size, join=True)
 
         if "samples" in return_dict:
             samples = return_dict["samples"]