RFDiffusion3 `DesignInputSpecification` object cannot be passed directly to the inference engine

[kjczarne]

Hi, I have noticed that when using DesignInputSpecification as an input fed directly to the run method on the RFD3InferenceEngine, the run always crashes when _multiply_specifications is invoked. This is due to the fact that DesignInputSpecification object is not subscriptable, yet example_spec is attempted to be set:

    def _multiply_specifications(
        self, inputs: Dict[str, dict | DesignInputSpecification], n_batches=None
    ) -> Dict[str, Dict[str, Any]]:
        ...
        # Based on inputs, construct the specifications to loop through
        design_specifications = {}
        for prefix, example_spec in inputs.items():
            # Record task name in the specification
            example_spec["extra"]["task_name"] = prefix

It seems like this does not happen when running RFD3 using JSON inputs, since the process_input function in engine.py returns a dictionary. In this case accessing "extra" key works without issues. However, when I attempt to feed a dictionary into .run() method on the engine instance, I get _canonicalize_inputs throwing an error complaining that a dictionary input is not supported.

This issue can be easily mitigated by using a dictionary example_spec_ in _multiply_specifications:

    def _multiply_specifications(
        self, inputs: Dict[str, dict | DesignInputSpecification], n_batches=None
    ) -> Dict[str, Dict[str, Any]]:
        # Find existing example IDS in output directory
        if exists(self.out_dir):
            existing_example_ids = set(
                extract_example_id_from_path(path, CIF_LIKE_EXTENSIONS)
                for path in find_files_with_extension(self.out_dir, CIF_LIKE_EXTENSIONS)
            )
            ranked_logger.info(
                f"Found {len(existing_example_ids)} existing example IDs in the output directory."
            )

        # Based on inputs, construct the specifications to loop through
        design_specifications = {}
        for prefix, example_spec in inputs.items():
            # Record task name in the specification
            example_spec_ = example_spec.model_dump()
            example_spec_["extra"]["task_name"] = prefix

            # ... Create n_batches for example
            for batch_id in range((n_batches) if exists(n_batches) else 1):
                # ... Example ID
                example_id = f"{prefix}_{batch_id}" if exists(n_batches) else prefix

                if (
                    self.skip_existing
                    and exists(self.out_dir)
                    and example_id in existing_example_ids
                ):
                    ranked_logger.info(
                        f"Skipping design specification for example {example_id} | Already exists."
                    )
                    continue
                design_specifications[example_id] = example_spec_
        return design_specifications

Please let me know if such a fix is OK for you, I could submit a PR for this. Or feel free to fix it on your end if you prefer no external contributions 🙂

Many thanks for providing this amazing repository, it's truly impressive how RFDiffusion has evolved!

[Ubiquinone-dot]

Thanks Chris!
I presume - perhaps for your use-case - that you don't want to have to dump the specification only to re-make the specification? I.e. that you might want to build the spec first, modify in-place or similar, and then pass it to run()? In that case it seems dumping the model would likely miss your custom attributes.

Will take a look at your other issues first and then see if i can find a solution

[Ubiquinone-dot]

Can properly discuss this later but have submitted a preliminary pr #118 with some of your suggested changes

[kjczarne]

Hi Jasper! Thanks for the fix! Seems like #117 is fixed (as long as I use absolute path on my end as a part of the DesignInputSpec), so is #116.

However, I'm still encountering a crash when using the DesignInputSpec object.

Here's what's happening.

L394:401 in the PR #118 leverages type-checking:

        for prefix, example_spec in inputs.items():
            # Record task name in the specification
            if isinstance(example_spec, DesignInputSpecification):
                example_spec.extra = example_spec.extra or {}
                example_spec.extra["task_name"] = prefix
            else:
                if "extra" not in example_spec:
                    example_spec["extra"] = {}
                example_spec["extra"]["task_name"] = prefix

This is in the _multiply_specification method in RFD3InferenceEngine. Return type hint on that function is Dict[str, Dict[str, Any]], meanwhile the example_spec object in the body of that function is of type DesignInputSpecification. Normally wouldn't be a problem but the output from _multiply_specification is passed directly to ensure_inference_sampler_matches_design_spec function (L291 in engine.py).

ensure_inference_sampler_matches_design_spec function has this attribute check:

    has_symmetry_specification = [
        True
        if "symmetry" in item.keys() and item.get("symmetry") is not None
        else False
        for item in design_spec.values()
    ]

But .items() does not exist for a DesignInputSpec instance! I think the easiest way to work around it is to just always return a dictionary from _multiply_specification method, i.e.

        for prefix, example_spec in inputs.items():
            # Record task name in the specification
            if isinstance(example_spec, DesignInputSpecification):
                example_spec = example_spec.model_dump()
            if "extra" not in example_spec:
                example_spec["extra"] = {}
            example_spec["extra"]["task_name"] = prefix

example_spec.model_dump() will return a dictionary which avoids having to juggle two different types here and has two primary benefits:

• ensure_inference_sampler_matches_design_spec doesn't crash
• Input type hints on ensure_inference_sampler_matches_design_spec are consistent, so are the output type hints of _multiply_specification

Let me know what are your thoughts on this. I've hotfixed my local copy of the repo with the snippet I've provided in the code block above.

[Ubiquinone-dot]

I think we'd still want to keep the original input specifciation intact - have merged #126 to handle the symmetry issue you mentioned

[Ubiquinone-dot]

Will close this as I believe it's solved but lmk!

[kjczarne]

Sweet, thanks Jasper, this works for me! 🚀

[Ubiquinone-dot]

Great!