Update src/eval.py

src/eval.py

@@ -1,9 +1,7 @@
 import os
 
-#  external libraries
+# External libraries
 import torch
-import torch.utils.checkpoint
-import torch.utils.checkpoint
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from diffusers import AutoencoderKL, DDIMScheduler
@@ -11,16 +9,15 @@ from diffusers.utils import check_min_version
 from diffusers.utils.import_utils import is_xformers_available
 from transformers import CLIPTextModel, CLIPTokenizer
 
-# custom imports
+# Custom imports
 from src.datasets.dresscode import DressCodeDataset
 from src.datasets.vitonhd import VitonHDDataset
 from src.mgd_pipelines.mgd_pipe import MGDPipe
 from src.mgd_pipelines.mgd_pipe_disentangled import MGDPipeDisentangled
-from src.utils.arg_parser import eval_parse_args
 from src.utils.image_from_pipe import generate_images_from_mgd_pipe
 from src.utils.set_seeds import set_seed
 
-# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+# Ensure the minimum version of diffusers is installed
 check_min_version("0.10.0.dev0")
 
 logger = get_logger(__name__, log_level="INFO")
@@ -28,139 +25,139 @@ os.environ["TOKENIZERS_PARALLELISM"] = "true"
 os.environ["WANDB_START_METHOD"] = "thread"
 
 
-def main() -> None:
-    args = eval_parse_args()
-    accelerator = Accelerator(
-        mixed_precision=args.mixed_precision,
-    )
+def main(args):
+    # Initialize Accelerator
+    accelerator = Accelerator(mixed_precision=args.get("mixed_precision", "fp16"))
     device = accelerator.device
 
     # Set the training seed
-    if args.seed is not None:
-        set_seed(args.seed)
+    if args.get("seed") is not None:
+        set_seed(args["seed"])
 
     # Load scheduler, tokenizer, and models
-    val_scheduler = DDIMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
+    val_scheduler = DDIMScheduler.from_pretrained(args["pretrained_model_name_or_path"], subfolder="scheduler")
     val_scheduler.set_timesteps(50, device=device)
 
     tokenizer = CLIPTokenizer.from_pretrained(
-        args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision
+        args["pretrained_model_name_or_path"], subfolder="tokenizer", revision=args.get("revision", None)
     )
     text_encoder = CLIPTextModel.from_pretrained(
-        args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision
+        args["pretrained_model_name_or_path"], subfolder="text_encoder", revision=args.get("revision", None)
     )
-    vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision)
+    vae = AutoencoderKL.from_pretrained(args["pretrained_model_name_or_path"], subfolder="vae", revision=args.get("revision", None))
 
-    # Load unet
+    # Load UNet
     unet = torch.hub.load(
-        dataset=args.dataset,
         repo_or_dir="aimagelab/multimodal-garment-designer",
         source="github",
         model="mgd",
         pretrained=True,
     )
 
-    # Freeze vae and text_encoder
+    # Freeze models
     vae.requires_grad_(False)
     text_encoder.requires_grad_(False)
 
     # Enable memory efficient attention if requested
-    if args.enable_xformers_memory_efficient_attention:
+    if args.get("enable_xformers_memory_efficient_attention", False):
         if is_xformers_available():
             unet.enable_xformers_memory_efficient_attention()
         else:
-            raise ValueError("xformers is not available. Make sure it is installed correctly")
+            raise ValueError("xformers is not available. Install it to enable memory-efficient attention.")
 
-    # Set the dataset category
-    category = [args.category] if args.category else ["dresses", "upper_body", "lower_body"]
+    # Set dataset category
+    category = [args.get("category", "dresses")]
 
-    # Load the appropriate dataset
-    if args.dataset == "dresscode":
+    # Load dataset
+    if args["dataset"] == "dresscode":
         test_dataset = DressCodeDataset(
-            dataroot_path=args.dataset_path,
+            dataroot_path=args["dataset_path"],
             phase="test",
-            order=args.test_order,
+            order=args.get("test_order", 0),
             radius=5,
             sketch_threshold_range=(20, 20),
             tokenizer=tokenizer,
             category=category,
             size=(512, 384),
         )
-    elif args.dataset == "vitonhd":
+    elif args["dataset"] == "vitonhd":
         test_dataset = VitonHDDataset(
-            dataroot_path=args.dataset_path,
+            dataroot_path=args["dataset_path"],
             phase="test",
-            order=args.test_order,
+            order=args.get("test_order", 0),
             sketch_threshold_range=(20, 20),
             radius=5,
             tokenizer=tokenizer,
             size=(512, 384),
         )
     else:
-        raise NotImplementedError(f"Dataset {args.dataset} is not supported.")
+        raise NotImplementedError(f"Dataset {args['dataset']} is not supported.")
 
-    # Prepare the dataloader
+    # Prepare dataloader
     test_dataloader = torch.utils.data.DataLoader(
         test_dataset,
         shuffle=False,
-        batch_size=args.batch_size,
-        num_workers=args.num_workers_test,
+        batch_size=args.get("batch_size", 1),
+        num_workers=args.get("num_workers_test", 4),
     )
 
-    # Cast text_encoder and vae to half-precision for mixed precision training
-    weight_dtype = torch.float32 if args.mixed_precision != "fp16" else torch.float16
+    # Cast models to appropriate precision
+    weight_dtype = torch.float32 if args.get("mixed_precision") != "fp16" else torch.float16
     text_encoder.to(device, dtype=weight_dtype)
     vae.to(device, dtype=weight_dtype)
-
-    # Ensure unet is in eval mode
     unet.eval()
 
-    # Select the appropriate pipeline
+    # Select pipeline
     with torch.inference_mode():
-        if args.disentagle:
-            val_pipe = MGDPipeDisentangled(
-                text_encoder=text_encoder,
-                vae=vae,
-                unet=unet.to(vae.dtype),
-                tokenizer=tokenizer,
-                scheduler=val_scheduler,
-            ).to(device)
-        else:
-            val_pipe = MGDPipe(
-                text_encoder=text_encoder,
-                vae=vae,
-                unet=unet.to(vae.dtype),
-                tokenizer=tokenizer,
-                scheduler=val_scheduler,
-            ).to(device)
-
-        # Debugging: Ensure val_pipe is callable
-        assert callable(val_pipe), "The pipeline object (val_pipe) is not callable. Check MGDPipe implementation."
-
-        # Enable attention slicing for memory efficiency
+        pipeline_class = MGDPipeDisentangled if args.get("disentagle", False) else MGDPipe
+        val_pipe = pipeline_class(
+            text_encoder=text_encoder,
+            vae=vae,
+            unet=unet.to(vae.dtype),
+            tokenizer=tokenizer,
+            scheduler=val_scheduler,
+        ).to(device)
+
         val_pipe.enable_attention_slicing()
 
         # Prepare dataloader with accelerator
         test_dataloader = accelerator.prepare(test_dataloader)
 
-        # Call the image generation function
+        # Generate images
+        output_path = os.path.join(args["output_dir"], args.get("save_name", "generated_image.png"))
         generate_images_from_mgd_pipe(
-            test_order=args.test_order,
+            test_order=args.get("test_order", 0),
             pipe=val_pipe,
             test_dataloader=test_dataloader,
-            save_name=args.save_name,
-            dataset=args.dataset,
-            output_dir=args.output_dir,
-            guidance_scale=args.guidance_scale,
-            guidance_scale_pose=args.guidance_scale_pose,
-            guidance_scale_sketch=args.guidance_scale_sketch,
-            sketch_cond_rate=args.sketch_cond_rate,
-            start_cond_rate=args.start_cond_rate,
+            save_name=args.get("save_name", "generated_image"),
+            dataset=args["dataset"],
+            output_dir=args["output_dir"],
+            guidance_scale=args.get("guidance_scale", 7.5),
+            guidance_scale_pose=args.get("guidance_scale_pose", 0.5),
+            guidance_scale_sketch=args.get("guidance_scale_sketch", 7.5),
+            sketch_cond_rate=args.get("sketch_cond_rate", 1.0),
+            start_cond_rate=args.get("start_cond_rate", 0.0),
             no_pose=False,
-            disentagle=args.disentagle,
-            seed=args.seed,
+            disentagle=args.get("disentagle", False),
+            seed=args.get("seed", None),
         )
 
+    # Return the output image path for verification
+    return output_path
+
 
 if __name__ == "__main__":
-    main()
+    # Example usage for debugging
+    example_args = {
+        "pretrained_model_name_or_path": "./models",
+        "dataset": "dresscode",
+        "dataset_path": "./datasets/dresscode",
+        "output_dir": "./outputs",
+        "guidance_scale": 7.5,
+        "guidance_scale_sketch": 7.5,
+        "mixed_precision": "fp16",
+        "batch_size": 1,
+        "seed": 42,
+    }
+    output_image = main(example_args)
+    print(f"Image generated at: {output_image}")