update: config management for evals

hemm/eval_pipelines/eval_pipeline.py

@@ -110,7 +110,7 @@ def __call__(self, dataset: Union[List[Dict], str]) -> Dict[str, float]:
             dataset=dataset,
             scorers=[metric_fn.evaluate_async for metric_fn in self.metric_functions],
         )
-        with weave.attributes(self.evaluation_configs):
-            summary = asyncio.run(evaluation.evaluate(self.infer_async))
+        self.model.configs.update(self.evaluation_configs)
+        summary = asyncio.run(evaluation.evaluate(self.infer_async))
         self.log_summary(summary)
         return summary

hemm/eval_pipelines/model.py

@@ -21,6 +21,7 @@ class BaseDiffusionModel(weave.Model):
     image_height: int = 512
     image_width: int = 512
     disable_safety_checker: bool = True
+    configs: Dict[str, Any] = {}
     pipeline_configs: Dict[str, Any] = {}
     _torch_dtype: torch.dtype = torch.float16
     _pipeline: DiffusionPipeline = None
@@ -32,6 +33,7 @@ def __init__(
         image_height: int = 512,
         image_width: int = 512,
         disable_safety_checker: bool = True,
+        configs: Dict[str, Any] = {},
         pipeline_configs: Dict[str, Any] = {},
     ) -> None:
         super().__init__(
@@ -40,6 +42,7 @@ def __init__(
             image_height=image_height,
             image_width=image_width,
             disable_safety_checker=disable_safety_checker,
+            configs=configs,
             pipeline_configs=pipeline_configs,
         )
         pipeline_init_kwargs = {
@@ -65,12 +68,4 @@ def predict(self, prompt: str, seed: int) -> Dict[str, Any]:
             width=self.image_width,
             generator=torch.Generator(device="cuda").manual_seed(seed),
         )
-        result_dict = {
-            "image": pipeline_output.images[0],
-        }
-        result_dict["nsfw_content_detected"] = (
-            pipeline_output.nsfw_content_detected is not None
-            if hasattr(pipeline_output, "nsfw_content_detected")
-            else False
-        )
-        return result_dict
+        return {"image": pipeline_output.images[0]}