better logs and add helper for convert.py

cog.yaml

@@ -1,18 +1,12 @@
 build:
   gpu: true
-  cuda: "11.6"
+  cuda: "11.8"
   python_version: "3.10"
   python_packages:
     - "diffusers==0.16.1"
-      # - "opencv-contrib-python==4.6.0.66"
-    # - "torchvision==0.15.1"
-      # - "timm==0.6.12"
     - "torch==2.0.1"
-      # - "ftfy==6.1.1"
     - "scipy==1.10.1"
     - "transformers==4.29.2"
     - "accelerate==0.19.0"
-  run:
-    # - "pip install git+https://github.com/patrickvonplaten/controlnet_aux.git@dfb604120af6e5acf21936f79aecb046d0faaa29"
-    - "pip install compel==1.0.4"
+    - "safetensors==0.3.1"
 predict: "predict.py:Predictor"

convert.py

@@ -0,0 +1,65 @@
+import torch
+import os
+from collections import defaultdict
+from safetensors.torch import load_file, save_file
+
+
+def check_file_size(sf_filename: str, pt_filename: str):
+    sf_size = os.stat(sf_filename).st_size
+    pt_size = os.stat(pt_filename).st_size
+
+    if (sf_size - pt_size) / pt_size > 0.01:
+        raise RuntimeError(
+            f"""The file size different is more than 1%:
+         - {sf_filename}: {sf_size}
+         - {pt_filename}: {pt_size}
+         """
+        )
+
+
+def shared_pointers(tensors):
+    ptrs = defaultdict(list)
+    for k, v in tensors.items():
+        ptrs[v.data_ptr()].append(k)
+    failing = []
+    for ptr, names in ptrs.items():
+        if len(names) > 1:
+            failing.append(names)
+    return failing
+
+
+def convert_file(
+    pt_filename: str,
+    sf_filename: str,
+):
+    loaded = torch.load(pt_filename, map_location="cpu")
+    if "state_dict" in loaded:
+        loaded = loaded["state_dict"]
+    shared = shared_pointers(loaded)
+    for shared_weights in shared:
+        for name in shared_weights[1:]:
+            loaded.pop(name)
+
+    # For tensors to be contiguous
+    loaded = {k: v.contiguous() for k, v in loaded.items()}
+
+    dirname = os.path.dirname(sf_filename)
+    os.makedirs(dirname, exist_ok=True)
+    save_file(loaded, sf_filename, metadata={"format": "pt"})
+    check_file_size(sf_filename, pt_filename)
+    reloaded = load_file(sf_filename)
+    for k in loaded:
+        pt_tensor = loaded[k]
+        sf_tensor = reloaded[k]
+        if not torch.equal(pt_tensor, sf_tensor):
+            raise RuntimeError(f"The output tensors do not match for key {k}")
+
+
+if __name__ == "__main__":
+    import glob
+    for bin_fn in glob.glob('weights/**/*.bin'):
+        if 'text_encoder' in bin_fn:
+            continue
+        st_fn = bin_fn.replace('.bin', '.safetensors')
+        convert_file(bin_fn, st_fn)
+        os.remove(bin_fn)

predict.py

@@ -1,9 +1,20 @@
+import time
 import datetime
 import random
 
 MARKER = "anotherjesse-sd-timings-%s" % random.randint(0, 1000000)
 
-print(MARKER, "start:", datetime.datetime.now())
+last = start = time.time()
+
+
+def crappy_log(*args):
+    global last
+    t = time.time()
+    print(MARKER, *args, "%0.2f" % (t - start), "%0.2f" % (t - last))
+    last = t
+
+
+crappy_log("start")
 
 from typing import Iterator
 
@@ -22,32 +33,105 @@
     UniPCMultistepScheduler,
 )
 
-MODEL_CACHE = "diffusers-cache"
-BASE_MODEL_PATH = "./weights"
-
-print(MARKER, "finished imports:", datetime.datetime.now())
+crappy_log("finished imports")
 
 
 class Predictor(BasePredictor):
     def setup(self):
         """Load the model into memory to make running multiple predictions efficient"""
-        print(MARKER, "Loading pipeline...", datetime.datetime.now())
+        crappy_log("starting setup")
 
-        if not os.path.exists(BASE_MODEL_PATH):
+        if not os.path.exists("./weights"):
             self.real = False
             return
 
-        print(MARKER, "Loading txt2img...")
-        self.txt2img_pipe = StableDiffusionPipeline.from_pretrained(
-            BASE_MODEL_PATH,
-            torch_dtype=torch.float16,
-            local_files_only=True,
-        ).to("cuda")
+        if True:
+            self.txt2img_pipe = StableDiffusionPipeline.from_pretrained(
+                "./weights",
+                torch_dtype=torch.float16,
+                local_files_only=True,
+                safety_checker=None,
+            )
+            crappy_log("finished importing pipeline")
+        else:
+            from diffusers.models import AutoencoderKL, UNet2DConditionModel
+            from transformers import CLIPTextModel, CLIPTokenizer
+            from diffusers.pipelines.stable_diffusion.safety_checker import (
+                StableDiffusionSafetyChecker,
+            )
+            from transformers import CLIPFeatureExtractor
+
+            vae = AutoencoderKL.from_pretrained(
+                "./weights",
+                torch_dtype=torch.float16,
+                local_files_only=True,
+                subfolder="vae",
+            )
+            crappy_log("loaded vae")
+
+            unet = UNet2DConditionModel.from_pretrained(
+                "./weights",
+                torch_dtype=torch.float16,
+                local_files_only=True,
+                subfolder="unet",
+            )
+            crappy_log("loaded unet")
+
+            text_encoder = CLIPTextModel.from_pretrained(
+                "./weights",
+                torch_dtype=torch.float16,
+                local_files_only=True,
+                subfolder="text_encoder",
+            )
+            crappy_log("loaded text_encoder")
+
+            tokenizer = CLIPTokenizer.from_pretrained(
+                "./weights",
+                torch_dtype=torch.float16,
+                local_files_only=True,
+                subfolder="tokenizer",
+            )
+            crappy_log("loaded tokenizer")
+
+            scheduler = DDIMScheduler.from_pretrained(
+                "./weights/scheduler/scheduler_config.json"
+            )
+            crappy_log("loaded scheduler")
+
+            safety_checker = StableDiffusionSafetyChecker.from_pretrained(
+                "./weights",
+                torch_dtype=torch.float16,
+                local_files_only=True,
+                subfolder="safety_checker",
+            )
+            crappy_log("loaded safety_checker")
+
+            feature_extractor = CLIPFeatureExtractor.from_pretrained(
+                "./weights",
+                torch_dtype=torch.float16,
+                local_files_only=True,
+                subfolder="feature_extractor",
+            )
+            crappy_log("loaded feature_extractor")
+
+            self.txt2img_pipe = StableDiffusionPipeline(
+                vae=vae,
+                unet=unet,
+                text_encoder=text_encoder,
+                tokenizer=tokenizer,
+                scheduler=scheduler,
+                safety_checker=safety_checker,
+                feature_extractor=feature_extractor,
+            )
+            crappy_log("manually created pipeline")
+
+        self.txt2img_pipe.to("cuda")
+        crappy_log("moved to cuda")
 
         self.safety_checker = self.txt2img_pipe.safety_checker
         self.real = True
 
-        print(MARKER, "loaded  pipeline...", datetime.datetime.now())
+        crappy_log("finished setup")
 
     def make_scheduler(self, name, config):
         return {
@@ -123,7 +207,7 @@ def predict(
         if not self.real:
             raise RuntimeError("This is a template, not a real model - add weights")
 
-        print(MARKER, "Using txt2img pipeline")
+        crappy_log("Using txt2img pipeline")
         pipe = self.txt2img_pipe
         extra_kwargs = {
             "width": width,
@@ -132,7 +216,8 @@ def predict(
 
         if seed is None:
             seed = int.from_bytes(os.urandom(2), "big")
-        print(MARKER, f"Using seed: {seed}")
+
+        crappy_log("Using seed: ", seed)
 
         if width * height > 786432:
             raise ValueError(