Could we make `BaseOutput` picklable?

[larme]

Currently we can pickle but cannot recover from the pickled bytes:

In [40]: bs =pickle.dumps(output)

In [41]: pickle.loads(bs)
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ in <module>:1                                                                                    │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
TypeError: IFPipelineOutput.__init__() missing 3 required positional arguments: 'images', 'nsfw_detected', and 'watermark_detected'

The motive is that we want to pass outputs of pipelines between processes or even machines. We can use to_tuple but user will need to know the order of different outputs' keys to use the returned tuples.

System Info

• diffusers version: 0.16.1
• Platform: debian 11
• Python version: 3.10.6
• PyTorch version (GPU?): 2.0 (True)
• Using distributed or parallel set-up in script?: True

[patrickvonplaten]

What is the use case here exactly ? Can you just use tuples by adding return_dict=False to the forward pass of IF?

[larme]

Hi @patrickvonplaten , I’m improving support of diffusers in BentoML. We need to pickle the BaseOutput so that we can pass the output around machines on a network using runner architecture. A practical usage examples is that we can distribute each stage of IF’s pipeline to its own pods/machines and scale them independetly (for example more pods for stage1 pipeline because it take more time to inference).

We want our integration works for every pipeline that is a subclass of diffusers.DiffusionPipeline, so we cannot pass return_dict=False because not every pipeline implement this argument.

[patrickvonplaten]

I see! IMO we should try to add return_dict=False to all missing pipelines as a first fix. Otherwise, any idea on how we could make BaseOutput picklable?

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[baijumeswani]

Can something like this be done? huggingface/transformers#8989

The problem is that the BaseOutput is not copyable (mostly because of this problem: python/cpython#105736).

Probably the easiest fix is to make the dataclass members not required while constructing the class. This should make the class copyable at least.