Update app.py

app.py

@@ -12,8 +12,6 @@ from diffusion import generate_latent
 from vq_vae import create_mask
 from huggingface_hub import snapshot_download
 
-# model_dir = 'trained_models'
-
 from huggingface_hub import login
 login(token = os.getenv('HF_TOKEN'))
 
@@ -29,7 +27,11 @@ def create_image(cancer_type):
     os.makedirs(tmp_dir)
     generate_latent(model_dir, cancer_type, tmp_dir)
     create_mask(model_dir, "./tmp", "./tmp/test_masks")
-    os.system('python pix2pixhd_test.py --name diffusion_dp --dataroot ./tmp --label_nc 0 --results_dir ./tmp --no_instance --resize_or_crop none --checkpoints_dir ' + model_dir)
+    if(cancer_type == 'benign'):
+        diffusion_model_name = 'diffusion_dp_benign'
+    else:
+        diffusion_model_name = 'diffusion_dp'
+    os.system('python pix2pixhd_test.py --name ' + diffusion_model_name + ' --dataroot ./tmp --label_nc 0 --results_dir ./tmp --no_instance --resize_or_crop none --checkpoints_dir ' + model_dir)
     image_dir = "./tmp/diffusion_dp/test_latest/images"
     input_label_image = Image.open(os.path.join(image_dir, "sample_input_label.jpg"))
     synthesized_image = Image.open(os.path.join(image_dir, "sample_synthesized_image.jpg"))