debugging

models/text_diffusion_pipeline.py

@@ -22,8 +22,14 @@ class PipelineOutput(NamedTuple):
 # Create a custom pipeline for text-conditional generation
 class TextConditionalDDPMPipeline(DDPMPipeline):
     def __init__(self, unet, scheduler, text_encoder=None, tokenizer=None, supports_pretrained_split=False, block_embeddings=None):
-        # Don't call super().__init__() with arguments, call it without arguments first
-        super(DiffusionPipeline, self).__init__()
+        # Debug: Print what we're receiving
+        print(f"unet type: {type(unet)}, value: {unet}")
+        print(f"scheduler type: {type(scheduler)}, value: {scheduler}")
+        print(f"text_encoder type: {type(text_encoder)}, value: {text_encoder}")
+        print(f"tokenizer type: {type(tokenizer)}, value: {tokenizer}")
+        
+        # Call DiffusionPipeline.__init__() directly (skipping DDPMPipeline's init)
+        DiffusionPipeline.__init__(self)
         
         self.text_encoder = text_encoder
         self.tokenizer = tokenizer
@@ -35,13 +41,22 @@ class TextConditionalDDPMPipeline(DDPMPipeline):
             # Use the tokenizer from the text encoder if not provided
             self.tokenizer = self.text_encoder.tokenizer
         
-        # Register ALL modules at once, including the ones from the parent class
-        self.register_modules(
-            unet=unet,
-            scheduler=scheduler,
-            text_encoder=self.text_encoder,
-            tokenizer=self.tokenizer,
-        )
+        # Only register modules that are actual objects, not None or lists
+        modules_to_register = {}
+        
+        if unet is not None and not isinstance(unet, (list, tuple)):
+            modules_to_register['unet'] = unet
+        if scheduler is not None and not isinstance(scheduler, (list, tuple)):
+            modules_to_register['scheduler'] = scheduler
+        if self.text_encoder is not None and not isinstance(self.text_encoder, (list, tuple)):
+            modules_to_register['text_encoder'] = self.text_encoder
+        if self.tokenizer is not None and not isinstance(self.tokenizer, (list, tuple)):
+            modules_to_register['tokenizer'] = self.tokenizer
+        
+        print(f"Registering modules: {list(modules_to_register.keys())}")
+        
+        # Register ALL modules at once
+        self.register_modules(**modules_to_register)
     
     # Override the to() method to ensure text_encoder is moved to the correct device
     def to(self, device=None, dtype=None):