[Torch FX][SD3] Add Dynamic Shapes

[anzr299]

Add Dynamic shapes to export of transformer model in SD3 Torch FX notebook.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[eaidova]

@anzr299 looks like something went wrong

Cell In[11], line 48, in collect_calibration_data(pipe, calibration_dataset_size, num_inference_steps)
     41 dataset = datasets.load_dataset(
     42     "google-research-datasets/conceptual_captions",
     43     split="train",
     44     trust_remote_code=True,
     45 ).shuffle(seed=42)
     47 transformer_config = dict(pipe.transformer.config)
---> 48 del transformer_config["model"]
     49 wrapped_unet = UNetWrapper(pipe.transformer.model, transformer_config)
     50 pipe.transformer = wrapped_unet

KeyError: 'model'

could you please fix?

[eaidova]

@anzr299, I made rude fix in this commit: 894d859#diff-ced57002da68fe81391edf1eab6f32e3320f21f35656a2211871d180924b1f27

if "model" in config:
   del config["model"]

but not sure that it s correct fix, lack of expected key in config may hide some other issues (e.g. applying patching several times or changes in config structure).I think it will be better if you look on it

[eaidova]

Now, next error happens:

RuntimeError: Cannot access data pointer of Tensor (e.g. FakeTensor, FunctionalTensor). If you're using torch.compile/export/fx, it is likely that we are erroneously tracing into a custom kernel. To fix this, please wrap the custom kernel into an opaque custom op. Please see the following for details: https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html

[anzr299]

Sorry for being unresposnive, I am looking into this issue now.

[anzr299]

@anzr299, I made rude fix in this commit: 894d859#diff-ced57002da68fe81391edf1eab6f32e3320f21f35656a2211871d180924b1f27

if "model" in config:
   del config["model"]

but not sure that it s correct fix, lack of expected key in config may hide some other issues (e.g. applying patching several times or changes in config structure).I think it will be better if you look on it

It actually slipped in from an older version I had locally. I have removed it now.

Now, next error happens:

This looks like it is caused due to dynamic shapes in torch export