Update implementation so that it actually uses the custom classes (#3)

- Update implementation so that it actually uses the custom classes (a28892bf8c0de64975befa77ab58389e58b090cc)


Co-authored-by: Sayak Paul <[email protected]>

my_pipeline.py

@@ -24,7 +24,9 @@ from diffusers.loaders import (
     StableDiffusionXLLoraLoaderMixin,
     TextualInversionLoaderMixin,
 )
-from diffusers.models import AutoencoderKL, UNet2DConditionModel
+from diffusers.models import AutoencoderKL
+from .scheduler.my_scheduler import MyScheduler
+from .unet.my_unet_model import MyUNetModel
 from diffusers.models.attention_processor import (
     AttnProcessor2_0,
     LoRAAttnProcessor2_0,
@@ -32,7 +34,6 @@ from diffusers.models.attention_processor import (
     XFormersAttnProcessor,
 )
 from diffusers.models.lora import adjust_lora_scale_text_encoder
-from diffusers.schedulers import KarrasDiffusionSchedulers
 from diffusers.utils import (
     USE_PEFT_BACKEND,
     is_invisible_watermark_available,
@@ -145,8 +146,8 @@ class MyPipeline(
         text_encoder_2: CLIPTextModelWithProjection,
         tokenizer: CLIPTokenizer,
         tokenizer_2: CLIPTokenizer,
-        unet: UNet2DConditionModel,
-        scheduler: KarrasDiffusionSchedulers,
+        unet: MyUNetModel,
+        scheduler: MyScheduler,
         force_zeros_for_empty_prompt: bool = True,
         add_watermarker: Optional[bool] = None,
     ):