Upload matryoshka.py

matryoshka.py

@@ -3762,7 +3762,7 @@ class MatryoshkaPipeline(
         text_encoder: T5EncoderModel,
         tokenizer: T5TokenizerFast,
         scheduler: MatryoshkaDDIMScheduler,
-        unet: MatryoshkaUNet2DConditionModel,
+        unet: MatryoshkaUNet2DConditionModel = None,
         feature_extractor: CLIPImageProcessor = None,
         image_encoder: CLIPVisionModelWithProjection = None,
         trust_remote_code: bool = False,
@@ -3782,6 +3782,8 @@ class MatryoshkaPipeline(
         else:
             raise ValueError("Currently, nesting levels 0, 1, and 2 are supported.")
 
+        unet = unet.to(self.device)
+
         if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
             deprecation_message = (
                 f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"