enable --eval for diffusion model (#1522)

Signed-off-by: Xin He <xin3.he@intel.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com>

Author: 3 people

auto_round/compressors/diffusion/eval.py

@@ -70,6 +70,7 @@ def diffusion_eval(
     image_save_dir,
     batch_size,
     gen_kwargs,
+    limit=None,
 ):
     if (
         not importlib.util.find_spec("clip")
@@ -82,6 +83,7 @@ def diffusion_eval(
     dataloader, _, _ = get_diffusion_dataloader(prompt_file, nsamples=-1, bs=batch_size)
     prompt_list = []
     image_list = []
+    num_samples = 0
     for image_ids, prompts in dataloader:
         prompt_list.extend(prompts)
 
@@ -95,6 +97,9 @@ def diffusion_eval(
                 continue
             new_ids.append(image_id)
             new_prompts.append(prompts[idx])
+            num_samples += 1
+            if limit is not None and num_samples >= limit > 0:
+                break
 
         if len(new_prompts) == 0:
             continue
@@ -103,6 +108,9 @@ def diffusion_eval(
         for idx, image_id in enumerate(new_ids):
             output.images[idx].save(os.path.join(image_save_dir, str(image_id) + ".png"))
 
+        if limit is not None and num_samples >= limit > 0:
+            break
+
     result = {}
     for metric in metrics:
         result.update(metric_map[metric](prompt_list, image_list, pipe.device))

auto_round/eval/eval_cli.py

@@ -24,6 +24,7 @@
     get_device_and_parallelism,
     get_device_str,
     get_model_dtype,
+    is_diffusion_model,
     set_cuda_visible_devices,
 )
 
@@ -113,6 +114,59 @@ def __init__(self, *args, **kwargs):
             help="(for vllm) Custom vllm arguments in format: 'arg1=value1,arg2=value2'. "
             "Example: 'tensor_parallel_size=2,gpu_memory_utilization=0.9'",
         )
+        # Diffusion model specific arguments
+        diffusion_args = self.add_argument_group("diffusion model arguments")
+        diffusion_args.add_argument(
+            "--prompt_file",
+            default=None,
+            type=str,
+            help="File containing prompts for evaluation, one per line. "
+            "Use this for batch evaluation with multiple prompts.",
+        )
+        diffusion_args.add_argument(
+            "--prompt",
+            default=None,
+            type=str,
+            help="Single prompt for quick testing. " "Overrides prompt_file if both are specified.",
+        )
+        diffusion_args.add_argument(
+            "--metrics",
+            "--metric",
+            default="clip",
+            help="Evaluation metrics for generated images. "
+            "'clip': CLIP score measuring text-image alignment. "
+            "'clip-iqa': CLIP-based image quality assessment. "
+            "'imagereward': Learned metric for image quality.",
+        )
+        diffusion_args.add_argument(
+            "--image_save_dir",
+            default="./tmp_image_save",
+            type=str,
+            help="Directory to save generated images during evaluation. " "Useful for visual inspection of results.",
+        )
+        diffusion_args.add_argument(
+            "--guidance_scale",
+            default=7.5,
+            type=float,
+            help="Classifier-free guidance scale for diffusion models. "
+            "Higher values (7-20) make the model follow the prompt more closely. "
+            "Lower values give more creative/random results.",
+        )
+        diffusion_args.add_argument(
+            "--num_inference_steps",
+            default=50,
+            type=int,
+            help="Number of denoising steps in the diffusion process. "
+            "More steps (50-100) usually give better quality but take longer. "
+            "Fewer steps (10-30) are faster but lower quality.",
+        )
+        diffusion_args.add_argument(
+            "--generator_seed",
+            default=None,
+            type=int,
+            help="Random seed for image generation reproducibility. "
+            "Using the same seed produces identical results across runs.",
+        )
 
 
 def _eval_init(tasks, model_path, device, disable_trust_remote_code=False, dtype="auto"):
@@ -134,6 +188,13 @@ def eval(args):
         "lm_eval>=0.4.2", "lm-eval is required for evaluation, please install it with `pip install 'lm-eval>=0.4.2'`"
     )
 
+    if is_diffusion_model(args.model):
+        from auto_round.eval.evaluation import evaluate_diffusion_model
+        from auto_round.utils import diffusion_load_model
+
+        pipe, _ = diffusion_load_model(args.model)
+        evaluate_diffusion_model(args, pipe=pipe)
+        return
     if args.eval_backend == "vllm":
         assert isinstance(args.model, str), "vllm evaluation only supports model name or path."
         eval_with_vllm(args)

auto_round/eval/evaluation.py

@@ -85,23 +85,28 @@ def simple_evaluate(
     )
 
 
-def evaluate_diffusion_model(autoround, model, args):
+def evaluate_diffusion_model(args, autoround=None, model=None, pipe=None):
     """
     Evaluate diffusion models.
 
     Args:
-        autoround: AutoRound instance
-        model: Quantized model
         args: Command line arguments
+        autoround: AutoRound instance (option 1)
+        model: Diffusion model instance (option 1)
+        pipe: Diffusion pipeline instance (option 2)
     """
+    if pipe is None and (autoround is None or model is None):
+        raise ValueError("Either 'pipe' must be provided, or both 'autoround' and 'model' must be provided.")
+
     import torch
 
     from auto_round.utils import detect_device, get_model_dtype, logger
 
     # Prepare inference pipeline
-    pipe = autoround.pipe
-    pipe.to(model.dtype)
-    pipe.transformer = model
+    if pipe is None:
+        pipe = autoround.pipe
+        pipe.to(model.dtype)
+        pipe.transformer = model
     device_str = detect_device(args.device_map if hasattr(args, "device_map") else "0")
     pipe = pipe.to(device_str)
 
@@ -134,11 +139,11 @@ def evaluate_diffusion_model(autoround, model, args):
         logger.info(f"Image generated with prompt {args.prompt} is saved as {save_path}")
 
     # Batch prompt evaluation
-    if args.prompt_file is not None:
+    elif args.prompt_file is not None:
         from auto_round.compressors.diffusion import diffusion_eval
 
         metrics = args.metrics.split(",")
-        diffusion_eval(pipe, args.prompt_file, metrics, args.image_save_dir, 1, gen_kwargs)
+        diffusion_eval(pipe, args.prompt_file, metrics, args.image_save_dir, 1, gen_kwargs, args.limit)
 
 
 def load_gguf_model_for_eval(eval_folder, formats, args):
@@ -375,7 +380,7 @@ def run_model_evaluation(model, tokenizer, autoround, folders, formats, device_s
 
     # Handle diffusion models separately
     if getattr(autoround, "diffusion", False):
-        evaluate_diffusion_model(autoround, model, args)
+        evaluate_diffusion_model(args, autoround=autoround, model=model)
         return
 
     # Check if evaluation is needed for language models