can't deepcopy quantized pipeline

[noskill]

It seems that quantized models can't be deepcopied anymore

optimum.quanto.version
'0.2.6'

from diffusers import FluxPipeline
from optimum.quanto import freeze, qfloat8, quantize as _quantize


def quantize(pipe, dtype=qfloat8):
    components = ['unet', 'transformer', 'text_encoder', 'text_encoder_2', 'vae']
    
    for component in components:
        if hasattr(pipe, component):
            component_obj = getattr(pipe, component)
            _quantize(component_obj, weights=dtype)
            freeze(component_obj)
            # Add attributes to indicate quantization
            component_obj._is_quantized = True
            component_obj._quantization_dtype = dtype


pipe0 = FluxPipeline.from_pretrained("./models/flux-1-dev", torch_dtype=torch.bfloat16).to("cpu")
quantize(pipe0)
cp.deepcopy(pipe0)

error:

    cp.deepcopy(pipe0)
  File "/usr/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/usr/lib/python3.10/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
  File "/usr/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/usr/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/usr/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/usr/lib/python3.10/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
  File "/usr/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/usr/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/usr/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/usr/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/usr/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/usr/lib/python3.10/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
  File "/usr/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/usr/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/usr/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/usr/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/usr/lib/python3.10/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/usr/lib/python3.10/copy.py", line 271, in _reconstruct
    state = deepcopy(state, memo)
  File "/usr/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/usr/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/usr/lib/python3.10/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/usr/lib/python3.10/copy.py", line 231, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/usr/lib/python3.10/copy.py", line 153, in deepcopy
    y = copier(memo)
  File "/tester/_work/_tool/Python/3.10.15/x64/lib/python3.10/site-packages/torch/_tensor.py", line 83, in __deepcopy__
    return handle_torch_function(Tensor.__deepcopy__, (self,), self, memo)
  File "/tester/_work/_tool/Python/3.10.15/x64/lib/python3.10/site-packages/torch/overrides.py", line 1739, in handle_torch_function
    result = torch_func_method(public_api, types, args, kwargs)
  File "/tester/_work/_tool/Python/3.10.15/x64/lib/python3.10/site-packages/optimum/quanto/tensor/weights/qbytes.py", line 272, in __torch_function__
    return func(*args, **kwargs)
  File "/tester/_work/_tool/Python/3.10.15/x64/lib/python3.10/site-packages/torch/_tensor.py", line 112, in __deepcopy__
    raise RuntimeError(
RuntimeError: The default implementation of __deepcopy__() for wrapper subclasses only works for subclass types that implement clone() and for which cloning returns another instance of the same subclass. You should either properly implement clone() for your subclass or override __deepcopy__() if it is intended behavior for clone() to return an instance of a different type.

if i downgrade to 0.2.4 code works

[github-actions]

This issue is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 5 days.

[noskill]

not stale!