just use snapshot download?

models/pipeline_loader.py

@@ -2,6 +2,7 @@ from models.text_diffusion_pipeline import TextConditionalDDPMPipeline
 from models.latent_diffusion_pipeline import UnconditionalDDPMPipeline
 import os
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from huggingface_hub import snapshot_download
 
 
 def get_pipeline(model_path):
@@ -16,24 +17,16 @@ def get_pipeline(model_path):
                 #If it has no text encoder, use the unconditional diffusion model
                 pipe = UnconditionalDDPMPipeline.from_pretrained(model_path)
     else:
-        # Assume it's a Hugging Face Hub model ID
-        # Try to load config to determine if it's text-conditional
-        config = DiffusionPipeline.load_config(model_path)
-        has_text_encoder = "text_encoder" in config
-            
+        # For HF Hub models, download first then load locally
+        print(f"Downloading model {model_path}...")
+        local_path = snapshot_download(repo_id=model_path, cache_dir="./temp_model_cache")
+        
+        # Check what components exist
+        has_text_encoder = os.path.exists(os.path.join(local_path, "text_encoder"))
+        
         if has_text_encoder:
-            # Use the local pipeline file for custom_pipeline
-            pipe = DiffusionPipeline.from_pretrained(
-                model_path,
-                custom_pipeline=model_path,
-                trust_remote_code=True,
-            )
+            pipe = TextConditionalDDPMPipeline.from_pretrained(local_path)
         else:
-            # Fallback: try unconditional 
-            pipe = DiffusionPipeline.from_pretrained(
-                model_path,
-                custom_pipeline=model_path,
-                trust_remote_code=True,
-            )
+            pipe = UnconditionalDDPMPipeline.from_pretrained(local_path)
 
     return pipe

models/text_diffusion_pipeline.py

@@ -18,18 +18,11 @@ class PipelineOutput(NamedTuple):
     images: torch.Tensor
     
 
-
 # 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):
-        # 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)
+        # Call parent class init normally
+        super().__init__(unet=unet, scheduler=scheduler)
         
         self.text_encoder = text_encoder
         self.tokenizer = tokenizer
@@ -39,24 +32,18 @@ class TextConditionalDDPMPipeline(DDPMPipeline):
         
         if self.tokenizer is None and self.text_encoder is not None:
             # Use the tokenizer from the text encoder if not provided
-            self.tokenizer = self.text_encoder.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
+            if hasattr(self.text_encoder, 'tokenizer'):
+                self.tokenizer = self.text_encoder.tokenizer
         
-        print(f"Registering modules: {list(modules_to_register.keys())}")
-        
-        # Register ALL modules at once
-        self.register_modules(**modules_to_register)
+        # Register additional modules if they exist
+        additional_modules = {}
+        if self.text_encoder is not None:
+            additional_modules['text_encoder'] = self.text_encoder
+        if self.tokenizer is not None:
+            additional_modules['tokenizer'] = self.tokenizer
+            
+        if additional_modules:
+            self.register_modules(**additional_modules)
     
     # Override the to() method to ensure text_encoder is moved to the correct device
     def to(self, device=None, dtype=None):