Upload combined_stable_diffusion.py with huggingface_hub

combined_stable_diffusion.py

@@ -2,11 +2,11 @@ import math
 import random
 
 import torch
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, DDPMScheduler
 from diffusers.image_processor import VaeImageProcessor
 from huggingface_hub import PyTorchModelHubMixin
 from PIL import Image
-
+from transformers import CLIPTextModel, CLIPTokenizer
 
 class CombinedStableDiffusion(
     DiffusionPipeline,
@@ -20,21 +20,28 @@ class CombinedStableDiffusion(
         self,
         original_unet: torch.nn.Module,
         fine_tuned_unet: torch.nn.Module,
-        scheduler,
+        scheduler: DDPMScheduler,
         vae: torch.nn.Module,
-        tokenizer=None,
-        text_encoder=None,
+        tokenizer: CLIPTextModel,
+        text_encoder: CLIPTokenizer,
     ) -> None:
 
         super().__init__()
 
+        # print(tokenizer)
+        # print(text_encoder)
+        # print(original_unet)
+        # print(fine_tuned_unet)
+        # print(scheduler)
+        # print(vae)
+        #
         self.register_modules(
             tokenizer=tokenizer,
+            text_encoder=text_encoder,
             original_unet=original_unet,
             fine_tuned_unet=fine_tuned_unet,
             scheduler=scheduler,
             vae=vae,
-            text_encoder=text_encoder,
         )
 
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)