Upload 13 files

model_index.json

@@ -1,15 +1,24 @@
 {
   "_class_name": "SuperDiffPipeline",
   "_diffusers_version": "0.31.0",
-  "batch_size": null,
-  "device": "cuda",
-  "guidance_scale": null,
-  "lift": null,
-  "num_inference_steps": null,
-  "scheduler": "EulerDiscreteScheduler",
-  "seed": null,
-  "text_encoder": "CLIPTextModel",
-  "tokenizer": "CLIPTokenizer",
-  "unet": "UNet2DConditionModel",
-  "vae": "AutoencoderKL"
+  "scheduler": [
+    "diffusers",
+    "EulerDiscreteScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "CLIPTextModel"
+  ],
+  "tokenizer": [
+    "transformers",
+    "CLIPTokenizer"
+  ],
+  "unet": [
+    "diffusers",
+    "UNet2DConditionModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ]
 }

pipeline.py

@@ -1,4 +1,5 @@
 import random
+from tqdm import tqdm
 from typing import Callable, Dict, List, Optional
 
 import torch
@@ -33,31 +34,38 @@ class SuperDiffPipeline(DiffusionPipeline, ConfigMixin):
 
         """
         super().__init__()
-        self.unet = unet
-        self.vae = vae
-        self.text_encoder = text_encoder
-        self.tokenizer = tokenizer
-        self.scheduler = scheduler
+        # Register additional parameters for flexibility
+        # Explicitly assign required components
+        #self.unet = unet
+        #self.vae = vae
+        #self.text_encoder = text_encoder
+        #self.tokenizer = tokenizer
+        #self.scheduler = scheduler
 
         device = "cuda" if torch.cuda.is_available() else "cpu"
 
-        self.vae.to(device)
-        self.unet.to(device)
-        self.text_encoder.to(device)
-
-        self.register_to_config(
-            vae=vae.__class__.__name__,
-            scheduler=scheduler.__class__.__name__,
-            tokenizer=tokenizer.__class__.__name__,
-            unet=unet.__class__.__name__,
-            text_encoder=text_encoder.__class__.__name__,
-            device=device,
-            batch_size=None,
-            num_inference_steps=None,
-            guidance_scale=None,
-            lift=None,
-            seed=None,
-        )
+        vae.to(device)
+        unet.to(device)
+        text_encoder.to(device)
+        self.register_modules(unet=unet, 
+                              scheduler=scheduler, 
+                              vae=vae, 
+                              text_encoder=text_encoder,
+                              tokenizer=tokenizer,)
+
+        #self.register_to_config(
+        #    vae=vae.__class__.__name__,
+        #    scheduler=scheduler.__class__.__name__,
+        #    tokenizer=tokenizer.__class__.__name__,
+        #    unet=unet.__class__.__name__,
+        #    text_encoder=text_encoder.__class__.__name__,
+        #    device=device,
+        #    batch_size=None,
+        #    num_inference_steps=None,
+        #    guidance_scale=None,
+        #    lift=None,
+        #    seed=None,
+        #)
 
     @torch.no_grad
     def get_batch(self, latents: Callable, nrow: int, ncol: int) -> Callable:
@@ -241,7 +249,7 @@ class SuperDiffPipeline(DiffusionPipeline, ConfigMixin):
             (self.num_inference_steps + 1, self.batch_size), device=self.device
         )
         with torch.no_grad():
-            for i, t in enumerate(self.scheduler.timesteps):
+            for i, t in tqdm(enumerate(self.scheduler.timesteps)):
                 dsigma = self.scheduler.sigmas[i +
                                                1] - self.scheduler.sigmas[i]
                 sigma = self.scheduler.sigmas[i]