chore(deps): bump actions/checkout from 5 to 6

Bumps [actions/checkout](https://github.com/actions/checkout) from 5 to 6.
- [Release notes](https://github.com/actions/checkout/releases)
- [Changelog](https://github.com/actions/checkout/blob/main/CHANGELOG.md)
- [Commits](https://github.com/actions/checkout/compare/v5...v6)

---
updated-dependencies:
- dependency-name: actions/checkout
  dependency-version: '6'
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>

.github/CODEOWNERS

@@ -1 +1 @@
-*       @mashb1t
+*       @lllyasviel

.github/ISSUE_TEMPLATE/bug_report.yml

@@ -16,9 +16,9 @@ body:
       description: |
         Please perform basic debugging to see if your configuration is the cause of the issue.
         Basic debug procedure
-        　2. Update Fooocus - sometimes things just need to be updated
-        　3. Backup and remove your config.txt - check if the issue is caused by bad configuration
-        　5. Try a fresh installation of Fooocus in a different directory - see if a clean installation solves the issue
+        　1. Update Fooocus - sometimes things just need to be updated
+        　2. Backup and remove your config.txt - check if the issue is caused by bad configuration
+        　3. Try a fresh installation of Fooocus in a different directory - see if a clean installation solves the issue
         Before making a issue report please, check that the issue hasn't been reported recently.
       options:
         - label: The issue has not been resolved by following the [troubleshooting guide](https://github.com/lllyasviel/Fooocus/blob/main/troubleshoot.md)

.github/workflows/build_container.yml

@@ -17,7 +17,7 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Log in to the Container registry
         uses: docker/login-action@v3
@@ -38,7 +38,7 @@ jobs:
             type=edge,branch=main
 
       - name: Build and push Docker image
-        uses: docker/build-push-action@v5
+        uses: docker/build-push-action@v6
         with:
           context: .
           file: ./Dockerfile

.gitignore

@@ -10,6 +10,7 @@ __pycache__
 *.partial
 *.onnx
 sorted_styles.json
+hash_cache.txt
 /input
 /cache
 /language/default.json

args_manager.py

@@ -1,7 +1,4 @@
 import ldm_patched.modules.args_parser as args_parser
-import os
-
-from tempfile import gettempdir
 
 args_parser.parser.add_argument("--share", action='store_true', help="Set whether to share on Gradio.")
 
@@ -10,7 +7,7 @@ args_parser.parser.add_argument("--disable-preset-selection", action='store_true
                                 help="Disables preset selection in Gradio.")
 
 args_parser.parser.add_argument("--language", type=str, default='default',
-                                help="Translate UI using json files in [language] folder."
+                                help="Translate UI using json files in [language] folder. "
                                   "For example, [--language example] will use [language/example.json] for translation.")
 
 # For example, https://github.com/lllyasviel/Fooocus/issues/849
@@ -18,26 +15,33 @@ args_parser.parser.add_argument("--disable-offload-from-vram", action="store_tru
                                 help="Force loading models to vram when the unload can be avoided. "
                                   "Some Mac users may need this.")
 
-args_parser.parser.add_argument("--theme", type=str, help="launches the UI with light or dark theme.", default=None)
+args_parser.parser.add_argument("--theme", type=str, help="launches the UI with light or dark theme", default=None)
 args_parser.parser.add_argument("--disable-image-log", action='store_true',
-                                help="Prevent writing images and logs to hard drive.")
+                                help="Prevent writing images and logs to the outputs folder.")
 
 args_parser.parser.add_argument("--disable-analytics", action='store_true',
                                 help="Disables analytics for Gradio.")
 
-args_parser.parser.add_argument("--disable-preset-download", action='store_true',
-                                help="Disables downloading models for presets.", default=False)
+args_parser.parser.add_argument("--disable-metadata", action='store_true',
+                                help="Disables saving metadata to images.")
 
-args_parser.parser.add_argument("--enable-describe-uov-image", action='store_true',
-                                help="Disables automatic description of uov images when prompt is empty.", default=False)
+args_parser.parser.add_argument("--disable-preset-download", action='store_true',
+                                help="Disables downloading models for presets", default=False)
+
+args_parser.parser.add_argument("--disable-enhance-output-sorting", action='store_true',
+                                help="Disables enhance output sorting for final image gallery.")
+
+args_parser.parser.add_argument("--enable-auto-describe-image", action='store_true',
+                                help="Enables automatic description of uov and enhance image when prompt is empty", default=False)
 
 args_parser.parser.add_argument("--always-download-new-model", action='store_true',
-                                help="Always download newer models.", default=False)
+                                help="Always download newer models", default=False)
 
-args_parser.parser.add_argument("--favicon-path", type=str, default=None, help="Set the favicon filepath.")
-args_parser.parser.add_argument("--auth-message", type=str, default=None, help="Message to show for auth.")
+args_parser.parser.add_argument("--rebuild-hash-cache", help="Generates missing model and LoRA hashes.",
+                                type=int, nargs="?", metavar="CPU_NUM_THREADS", const=-1)
 
 args_parser.parser.set_defaults(
+    disable_cuda_malloc=True,
     in_browser=True,
     port=None
 )

css/style.css

@@ -99,7 +99,7 @@ div:has(> #positive_prompt) {
 }
 
 .advanced_check_row {
-  width: 250px !important;
+  width: 330px !important;
 }
 
 .min_check {

experiments_mask_generation.py

@@ -0,0 +1,24 @@
+# https://github.com/sail-sg/EditAnything/blob/main/sam2groundingdino_edit.py
+
+import numpy as np
+from PIL import Image
+
+from extras.inpaint_mask import SAMOptions, generate_mask_from_image
+
+original_image = Image.open('cat.webp')
+image = np.array(original_image, dtype=np.uint8)
+
+sam_options = SAMOptions(
+    dino_prompt='eye',
+    dino_box_threshold=0.3,
+    dino_text_threshold=0.25,
+    dino_erode_or_dilate=0,
+    dino_debug=False,
+    max_detections=2,
+    model_type='vit_b'
+)
+
+mask_image, _, _, _ = generate_mask_from_image(image, sam_options=sam_options)
+
+merged_masks_img = Image.fromarray(mask_image)
+merged_masks_img.show()

extras/BLIP/models/blip.py

@@ -216,9 +216,9 @@ def is_url(url_or_filename):
 def load_checkpoint(model,url_or_filename):
     if is_url(url_or_filename):
         cached_file = download_cached_file(url_or_filename, check_hash=False, progress=True)
-        checkpoint = torch.load(cached_file, map_location='cpu') 
+        checkpoint = torch.load(cached_file, map_location='cpu', weights_only=True) 
     elif os.path.isfile(url_or_filename):        
-        checkpoint = torch.load(url_or_filename, map_location='cpu') 
+        checkpoint = torch.load(url_or_filename, map_location='cpu', weights_only=True) 
     else:
         raise RuntimeError('checkpoint url or path is invalid')
         

extras/BLIP/models/blip_nlvr.py

@@ -78,9 +78,9 @@ def blip_nlvr(pretrained='',**kwargs):
 def load_checkpoint(model,url_or_filename):
     if is_url(url_or_filename):
         cached_file = download_cached_file(url_or_filename, check_hash=False, progress=True)
-        checkpoint = torch.load(cached_file, map_location='cpu') 
+        checkpoint = torch.load(cached_file, map_location='cpu', weights_only=True) 
     elif os.path.isfile(url_or_filename):        
-        checkpoint = torch.load(url_or_filename, map_location='cpu') 
+        checkpoint = torch.load(url_or_filename, map_location='cpu', weights_only=True) 
     else:
         raise RuntimeError('checkpoint url or path is invalid')
     state_dict = checkpoint['model']

extras/GroundingDINO/config/GroundingDINO_SwinT_OGC.py

@@ -0,0 +1,43 @@
+batch_size = 1
+modelname = "groundingdino"
+backbone = "swin_T_224_1k"
+position_embedding = "sine"
+pe_temperatureH = 20
+pe_temperatureW = 20
+return_interm_indices = [1, 2, 3]
+backbone_freeze_keywords = None
+enc_layers = 6
+dec_layers = 6
+pre_norm = False
+dim_feedforward = 2048
+hidden_dim = 256
+dropout = 0.0
+nheads = 8
+num_queries = 900
+query_dim = 4
+num_patterns = 0
+num_feature_levels = 4
+enc_n_points = 4
+dec_n_points = 4
+two_stage_type = "standard"
+two_stage_bbox_embed_share = False
+two_stage_class_embed_share = False
+transformer_activation = "relu"
+dec_pred_bbox_embed_share = True
+dn_box_noise_scale = 1.0
+dn_label_noise_ratio = 0.5
+dn_label_coef = 1.0
+dn_bbox_coef = 1.0
+embed_init_tgt = True
+dn_labelbook_size = 2000
+max_text_len = 256
+text_encoder_type = "bert-base-uncased"
+use_text_enhancer = True
+use_fusion_layer = True
+use_checkpoint = True
+use_transformer_ckpt = True
+use_text_cross_attention = True
+text_dropout = 0.0
+fusion_dropout = 0.0
+fusion_droppath = 0.1
+sub_sentence_present = True

extras/GroundingDINO/util/inference.py

@@ -0,0 +1,100 @@
+from typing import Tuple, List
+
+import ldm_patched.modules.model_management as model_management
+from ldm_patched.modules.model_patcher import ModelPatcher
+from modules.config import path_inpaint
+from modules.model_loader import load_file_from_url
+
+import numpy as np
+import supervision as sv
+import torch
+from groundingdino.util.inference import Model
+from groundingdino.util.inference import load_model, preprocess_caption, get_phrases_from_posmap
+
+
+class GroundingDinoModel(Model):
+    def __init__(self):
+        self.config_file = 'extras/GroundingDINO/config/GroundingDINO_SwinT_OGC.py'
+        self.model = None
+        self.load_device = torch.device('cpu')
+        self.offload_device = torch.device('cpu')
+
+    @torch.no_grad()
+    @torch.inference_mode()
+    def predict_with_caption(
+            self,
+            image: np.ndarray,
+            caption: str,
+            box_threshold: float = 0.35,
+            text_threshold: float = 0.25
+    ) -> Tuple[sv.Detections, torch.Tensor, torch.Tensor, List[str]]:
+        if self.model is None:
+            filename = load_file_from_url(
+                url="https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth",
+                file_name='groundingdino_swint_ogc.pth',
+                model_dir=path_inpaint)
+            model = load_model(model_config_path=self.config_file, model_checkpoint_path=filename)
+
+            self.load_device = model_management.text_encoder_device()
+            self.offload_device = model_management.text_encoder_offload_device()
+
+            model.to(self.offload_device)
+
+            self.model = ModelPatcher(model, load_device=self.load_device, offload_device=self.offload_device)
+
+        model_management.load_model_gpu(self.model)
+
+        processed_image = GroundingDinoModel.preprocess_image(image_bgr=image).to(self.load_device)
+        boxes, logits, phrases = predict(
+            model=self.model,
+            image=processed_image,
+            caption=caption,
+            box_threshold=box_threshold,
+            text_threshold=text_threshold,
+            device=self.load_device)
+        source_h, source_w, _ = image.shape
+        detections = GroundingDinoModel.post_process_result(
+            source_h=source_h,
+            source_w=source_w,
+            boxes=boxes,
+            logits=logits)
+        return detections, boxes, logits, phrases
+
+
+def predict(
+        model,
+        image: torch.Tensor,
+        caption: str,
+        box_threshold: float,
+        text_threshold: float,
+        device: str = "cuda"
+) -> Tuple[torch.Tensor, torch.Tensor, List[str]]:
+    caption = preprocess_caption(caption=caption)
+
+    # override to use model wrapped by patcher
+    model = model.model.to(device)
+    image = image.to(device)
+
+    with torch.no_grad():
+        outputs = model(image[None], captions=[caption])
+
+    prediction_logits = outputs["pred_logits"].cpu().sigmoid()[0]  # prediction_logits.shape = (nq, 256)
+    prediction_boxes = outputs["pred_boxes"].cpu()[0]  # prediction_boxes.shape = (nq, 4)
+
+    mask = prediction_logits.max(dim=1)[0] > box_threshold
+    logits = prediction_logits[mask]  # logits.shape = (n, 256)
+    boxes = prediction_boxes[mask]  # boxes.shape = (n, 4)
+
+    tokenizer = model.tokenizer
+    tokenized = tokenizer(caption)
+
+    phrases = [
+        get_phrases_from_posmap(logit > text_threshold, tokenized, tokenizer).replace('.', '')
+        for logit
+        in logits
+    ]
+
+    return boxes, logits.max(dim=1)[0], phrases
+
+
+default_groundingdino = GroundingDinoModel().predict_with_caption

extras/censor.py

@@ -41,7 +41,7 @@ class Censor:
         model_management.load_model_gpu(self.safety_checker_model)
 
         single = False
-        if not isinstance(images, list) or isinstance(images, np.ndarray):
+        if not isinstance(images, (list, np.ndarray)):
             images = [images]
             single = True
 

extras/facexlib/detection/__init__.py

@@ -19,7 +19,7 @@ def init_detection_model(model_name, half=False, device='cuda', model_rootpath=N
         url=model_url, model_dir='facexlib/weights', progress=True, file_name=None, save_dir=model_rootpath)
 
     # TODO: clean pretrained model
-    load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
+    load_net = torch.load(model_path, map_location=lambda storage, loc: storage, weights_only=True)
     # remove unnecessary 'module.'
     for k, v in deepcopy(load_net).items():
         if k.startswith('module.'):

extras/facexlib/parsing/__init__.py

@@ -17,7 +17,7 @@ def init_parsing_model(model_name='bisenet', half=False, device='cuda', model_ro
 
     model_path = load_file_from_url(
         url=model_url, model_dir='facexlib/weights', progress=True, file_name=None, save_dir=model_rootpath)
-    load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
+    load_net = torch.load(model_path, map_location=lambda storage, loc: storage, weights_only=True)
     model.load_state_dict(load_net, strict=True)
     model.eval()
     model = model.to(device)

extras/inpaint_mask.py

@@ -0,0 +1,130 @@
+import sys
+
+import modules.config
+import numpy as np
+import torch
+from extras.GroundingDINO.util.inference import default_groundingdino
+from extras.sam.predictor import SamPredictor
+from rembg import remove, new_session
+from segment_anything import sam_model_registry
+from segment_anything.utils.amg import remove_small_regions
+
+
+class SAMOptions:
+    def __init__(self,
+                 # GroundingDINO
+                 dino_prompt: str = '',
+                 dino_box_threshold=0.3,
+                 dino_text_threshold=0.25,
+                 dino_erode_or_dilate=0,
+                 dino_debug=False,
+
+                 # SAM
+                 max_detections=2,
+                 model_type='vit_b'
+                 ):
+        self.dino_prompt = dino_prompt
+        self.dino_box_threshold = dino_box_threshold
+        self.dino_text_threshold = dino_text_threshold
+        self.dino_erode_or_dilate = dino_erode_or_dilate
+        self.dino_debug = dino_debug
+        self.max_detections = max_detections
+        self.model_type = model_type
+
+
+def optimize_masks(masks: torch.Tensor) -> torch.Tensor:
+    """
+    removes small disconnected regions and holes
+    """
+    fine_masks = []
+    for mask in masks.to('cpu').numpy():  # masks: [num_masks, 1, h, w]
+        fine_masks.append(remove_small_regions(mask[0], 400, mode="holes")[0])
+    masks = np.stack(fine_masks, axis=0)[:, np.newaxis]
+    return torch.from_numpy(masks)
+
+
+def generate_mask_from_image(image: np.ndarray, mask_model: str = 'sam', extras=None,
+                             sam_options: SAMOptions | None = SAMOptions) -> tuple[np.ndarray | None, int | None, int | None, int | None]:
+    dino_detection_count = 0
+    sam_detection_count = 0
+    sam_detection_on_mask_count = 0
+
+    if image is None:
+        return None, dino_detection_count, sam_detection_count, sam_detection_on_mask_count
+
+    if extras is None:
+        extras = {}
+
+    if 'image' in image:
+        image = image['image']
+
+    if mask_model != 'sam' or sam_options is None:
+        result = remove(
+            image,
+            session=new_session(mask_model, **extras),
+            only_mask=True,
+            **extras
+        )
+
+        return result, dino_detection_count, sam_detection_count, sam_detection_on_mask_count
+
+    detections, boxes, logits, phrases = default_groundingdino(
+        image=image,
+        caption=sam_options.dino_prompt,
+        box_threshold=sam_options.dino_box_threshold,
+        text_threshold=sam_options.dino_text_threshold
+    )
+
+    H, W = image.shape[0], image.shape[1]
+    boxes = boxes * torch.Tensor([W, H, W, H])
+    boxes[:, :2] = boxes[:, :2] - boxes[:, 2:] / 2
+    boxes[:, 2:] = boxes[:, 2:] + boxes[:, :2]
+
+    sam_checkpoint = modules.config.download_sam_model(sam_options.model_type)
+    sam = sam_model_registry[sam_options.model_type](checkpoint=sam_checkpoint)
+
+    sam_predictor = SamPredictor(sam)
+    final_mask_tensor = torch.zeros((image.shape[0], image.shape[1]))
+    dino_detection_count = boxes.size(0)
+
+    if dino_detection_count > 0:
+        sam_predictor.set_image(image)
+
+        if sam_options.dino_erode_or_dilate != 0:
+            for index in range(boxes.size(0)):
+                assert boxes.size(1) == 4
+                boxes[index][0] -= sam_options.dino_erode_or_dilate
+                boxes[index][1] -= sam_options.dino_erode_or_dilate
+                boxes[index][2] += sam_options.dino_erode_or_dilate
+                boxes[index][3] += sam_options.dino_erode_or_dilate
+
+        if sam_options.dino_debug:
+            from PIL import ImageDraw, Image
+            debug_dino_image = Image.new("RGB", (image.shape[1], image.shape[0]), color="black")
+            draw = ImageDraw.Draw(debug_dino_image)
+            for box in boxes.numpy():
+                draw.rectangle(box.tolist(), fill="white")
+            return np.array(debug_dino_image), dino_detection_count, sam_detection_count, sam_detection_on_mask_count
+
+        transformed_boxes = sam_predictor.transform.apply_boxes_torch(boxes, image.shape[:2])
+        masks, _, _ = sam_predictor.predict_torch(
+            point_coords=None,
+            point_labels=None,
+            boxes=transformed_boxes,
+            multimask_output=False,
+        )
+
+        masks = optimize_masks(masks)
+        sam_detection_count = len(masks)
+        if sam_options.max_detections == 0:
+            sam_options.max_detections = sys.maxsize
+        sam_objects = min(len(logits), sam_options.max_detections)
+        for obj_ind in range(sam_objects):
+            mask_tensor = masks[obj_ind][0]
+            final_mask_tensor += mask_tensor
+            sam_detection_on_mask_count += 1
+
+    final_mask_tensor = (final_mask_tensor > 0).to('cpu').numpy()
+    mask_image = np.dstack((final_mask_tensor, final_mask_tensor, final_mask_tensor)) * 255
+    mask_image = np.array(mask_image, dtype=np.uint8)
+    return mask_image, dino_detection_count, sam_detection_count, sam_detection_on_mask_count

extras/ip_adapter.py

@@ -104,7 +104,7 @@ def load_ip_adapter(clip_vision_path, ip_negative_path, ip_adapter_path):
     offload_device = torch.device('cpu')
 
     use_fp16 = model_management.should_use_fp16(device=load_device)
-    ip_state_dict = torch.load(ip_adapter_path, map_location="cpu")
+    ip_state_dict = torch.load(ip_adapter_path, map_location="cpu", weights_only=True)
     plus = "latents" in ip_state_dict["image_proj"]
     cross_attention_dim = ip_state_dict["ip_adapter"]["1.to_k_ip.weight"].shape[1]
     sdxl = cross_attention_dim == 2048

extras/sam/predictor.py

@@ -0,0 +1,288 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import numpy as np
+import torch
+from ldm_patched.modules import model_management
+from ldm_patched.modules.model_patcher import ModelPatcher
+
+from segment_anything.modeling import Sam
+
+from typing import Optional, Tuple
+
+from segment_anything.utils.transforms import ResizeLongestSide
+
+
+class SamPredictor:
+    def __init__(
+        self,
+        model: Sam,
+        load_device=model_management.text_encoder_device(),
+        offload_device=model_management.text_encoder_offload_device()
+    ) -> None:
+        """
+        Uses SAM to calculate the image embedding for an image, and then
+        allow repeated, efficient mask prediction given prompts.
+
+        Arguments:
+          model (Sam): The model to use for mask prediction.
+        """
+        super().__init__()
+
+        self.load_device = load_device
+        self.offload_device = offload_device
+        # can't use model.half() here as slow_conv2d_cpu is not implemented for half
+        model.to(self.offload_device)
+
+        self.patcher = ModelPatcher(model, load_device=self.load_device, offload_device=self.offload_device)
+
+        self.transform = ResizeLongestSide(model.image_encoder.img_size)
+        self.reset_image()
+
+    def set_image(
+        self,
+        image: np.ndarray,
+        image_format: str = "RGB",
+    ) -> None:
+        """
+        Calculates the image embeddings for the provided image, allowing
+        masks to be predicted with the 'predict' method.
+
+        Arguments:
+          image (np.ndarray): The image for calculating masks. Expects an
+            image in HWC uint8 format, with pixel values in [0, 255].
+          image_format (str): The color format of the image, in ['RGB', 'BGR'].
+        """
+        assert image_format in [
+            "RGB",
+            "BGR",
+        ], f"image_format must be in ['RGB', 'BGR'], is {image_format}."
+        if image_format != self.patcher.model.image_format:
+            image = image[..., ::-1]
+
+        # Transform the image to the form expected by the model
+        input_image = self.transform.apply_image(image)
+        input_image_torch = torch.as_tensor(input_image, device=self.load_device)
+        input_image_torch = input_image_torch.permute(2, 0, 1).contiguous()[None, :, :, :]
+
+        self.set_torch_image(input_image_torch, image.shape[:2])
+
+    @torch.no_grad()
+    def set_torch_image(
+        self,
+        transformed_image: torch.Tensor,
+        original_image_size: Tuple[int, ...],
+    ) -> None:
+        """
+        Calculates the image embeddings for the provided image, allowing
+        masks to be predicted with the 'predict' method. Expects the input
+        image to be already transformed to the format expected by the model.
+
+        Arguments:
+          transformed_image (torch.Tensor): The input image, with shape
+            1x3xHxW, which has been transformed with ResizeLongestSide.
+          original_image_size (tuple(int, int)): The size of the image
+            before transformation, in (H, W) format.
+        """
+        assert (
+            len(transformed_image.shape) == 4
+            and transformed_image.shape[1] == 3
+            and max(*transformed_image.shape[2:]) == self.patcher.model.image_encoder.img_size
+        ), f"set_torch_image input must be BCHW with long side {self.patcher.model.image_encoder.img_size}."
+        self.reset_image()
+
+        self.original_size = original_image_size
+        self.input_size = tuple(transformed_image.shape[-2:])
+        model_management.load_model_gpu(self.patcher)
+        input_image = self.patcher.model.preprocess(transformed_image.to(self.load_device))
+        self.features = self.patcher.model.image_encoder(input_image)
+        self.is_image_set = True
+
+    def predict(
+        self,
+        point_coords: Optional[np.ndarray] = None,
+        point_labels: Optional[np.ndarray] = None,
+        box: Optional[np.ndarray] = None,
+        mask_input: Optional[np.ndarray] = None,
+        multimask_output: bool = True,
+        return_logits: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Predict masks for the given input prompts, using the currently set image.
+
+        Arguments:
+          point_coords (np.ndarray or None): A Nx2 array of point prompts to the
+            model. Each point is in (X,Y) in pixels.
+          point_labels (np.ndarray or None): A length N array of labels for the
+            point prompts. 1 indicates a foreground point and 0 indicates a
+            background point.
+          box (np.ndarray or None): A length 4 array given a box prompt to the
+            model, in XYXY format.
+          mask_input (np.ndarray): A low resolution mask input to the model, typically
+            coming from a previous prediction iteration. Has form 1xHxW, where
+            for SAM, H=W=256.
+          multimask_output (bool): If true, the model will return three masks.
+            For ambiguous input prompts (such as a single click), this will often
+            produce better masks than a single prediction. If only a single
+            mask is needed, the model's predicted quality score can be used
+            to select the best mask. For non-ambiguous prompts, such as multiple
+            input prompts, multimask_output=False can give better results.
+          return_logits (bool): If true, returns un-thresholded masks logits
+            instead of a binary mask.
+
+        Returns:
+          (np.ndarray): The output masks in CxHxW format, where C is the
+            number of masks, and (H, W) is the original image size.
+          (np.ndarray): An array of length C containing the model's
+            predictions for the quality of each mask.
+          (np.ndarray): An array of shape CxHxW, where C is the number
+            of masks and H=W=256. These low resolution logits can be passed to
+            a subsequent iteration as mask input.
+        """
+        if not self.is_image_set:
+            raise RuntimeError("An image must be set with .set_image(...) before mask prediction.")
+
+        # Transform input prompts
+        coords_torch, labels_torch, box_torch, mask_input_torch = None, None, None, None
+        if point_coords is not None:
+            assert (
+                point_labels is not None
+            ), "point_labels must be supplied if point_coords is supplied."
+            point_coords = self.transform.apply_coords(point_coords, self.original_size)
+            coords_torch = torch.as_tensor(point_coords, dtype=torch.float, device=self.load_device)
+            labels_torch = torch.as_tensor(point_labels, dtype=torch.int, device=self.load_device)
+            coords_torch, labels_torch = coords_torch[None, :, :], labels_torch[None, :]
+        if box is not None:
+            box = self.transform.apply_boxes(box, self.original_size)
+            box_torch = torch.as_tensor(box, dtype=torch.float, device=self.load_device)
+            box_torch = box_torch[None, :]
+        if mask_input is not None:
+            mask_input_torch = torch.as_tensor(mask_input, dtype=torch.float, device=self.load_device)
+            mask_input_torch = mask_input_torch[None, :, :, :]
+
+        masks, iou_predictions, low_res_masks = self.predict_torch(
+            coords_torch,
+            labels_torch,
+            box_torch,
+            mask_input_torch,
+            multimask_output,
+            return_logits=return_logits,
+        )
+
+        masks = masks[0].detach().cpu().numpy()
+        iou_predictions = iou_predictions[0].detach().cpu().numpy()
+        low_res_masks = low_res_masks[0].detach().cpu().numpy()
+        return masks, iou_predictions, low_res_masks
+
+    @torch.no_grad()
+    def predict_torch(
+        self,
+        point_coords: Optional[torch.Tensor],
+        point_labels: Optional[torch.Tensor],
+        boxes: Optional[torch.Tensor] = None,
+        mask_input: Optional[torch.Tensor] = None,
+        multimask_output: bool = True,
+        return_logits: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Predict masks for the given input prompts, using the currently set image.
+        Input prompts are batched torch tensors and are expected to already be
+        transformed to the input frame using ResizeLongestSide.
+
+        Arguments:
+          point_coords (torch.Tensor or None): A BxNx2 array of point prompts to the
+            model. Each point is in (X,Y) in pixels.
+          point_labels (torch.Tensor or None): A BxN array of labels for the
+            point prompts. 1 indicates a foreground point and 0 indicates a
+            background point.
+          box (np.ndarray or None): A Bx4 array given a box prompt to the
+            model, in XYXY format.
+          mask_input (np.ndarray): A low resolution mask input to the model, typically
+            coming from a previous prediction iteration. Has form Bx1xHxW, where
+            for SAM, H=W=256. Masks returned by a previous iteration of the
+            predict method do not need further transformation.
+          multimask_output (bool): If true, the model will return three masks.
+            For ambiguous input prompts (such as a single click), this will often
+            produce better masks than a single prediction. If only a single
+            mask is needed, the model's predicted quality score can be used
+            to select the best mask. For non-ambiguous prompts, such as multiple
+            input prompts, multimask_output=False can give better results.
+          return_logits (bool): If true, returns un-thresholded masks logits
+            instead of a binary mask.
+
+        Returns:
+          (torch.Tensor): The output masks in BxCxHxW format, where C is the
+            number of masks, and (H, W) is the original image size.
+          (torch.Tensor): An array of shape BxC containing the model's
+            predictions for the quality of each mask.
+          (torch.Tensor): An array of shape BxCxHxW, where C is the number
+            of masks and H=W=256. These low res logits can be passed to
+            a subsequent iteration as mask input.
+        """
+        if not self.is_image_set:
+            raise RuntimeError("An image must be set with .set_image(...) before mask prediction.")
+
+        if point_coords is not None:
+            points = (point_coords.to(self.load_device), point_labels.to(self.load_device))
+        else:
+            points = None
+
+        # load
+        if boxes is not None:
+            boxes = boxes.to(self.load_device)
+        if mask_input is not None:
+            mask_input = mask_input.to(self.load_device)
+        model_management.load_model_gpu(self.patcher)
+
+        # Embed prompts
+        sparse_embeddings, dense_embeddings = self.patcher.model.prompt_encoder(
+            points=points,
+            boxes=boxes,
+            masks=mask_input,
+        )
+
+        # Predict masks
+        low_res_masks, iou_predictions = self.patcher.model.mask_decoder(
+            image_embeddings=self.features,
+            image_pe=self.patcher.model.prompt_encoder.get_dense_pe(),
+            sparse_prompt_embeddings=sparse_embeddings,
+            dense_prompt_embeddings=dense_embeddings,
+            multimask_output=multimask_output,
+        )
+
+        # Upscale the masks to the original image resolution
+        masks = self.patcher.model.postprocess_masks(low_res_masks, self.input_size, self.original_size)
+
+        if not return_logits:
+            masks = masks > self.patcher.model.mask_threshold
+
+        return masks, iou_predictions, low_res_masks
+
+    def get_image_embedding(self) -> torch.Tensor:
+        """
+        Returns the image embeddings for the currently set image, with
+        shape 1xCxHxW, where C is the embedding dimension and (H,W) are
+        the embedding spatial dimension of SAM (typically C=256, H=W=64).
+        """
+        if not self.is_image_set:
+            raise RuntimeError(
+                "An image must be set with .set_image(...) to generate an embedding."
+            )
+        assert self.features is not None, "Features must exist if an image has been set."
+        return self.features
+
+    @property
+    def device(self) -> torch.device:
+        return self.patcher.model.device
+
+    def reset_image(self) -> None:
+        """Resets the currently set image."""
+        self.is_image_set = False
+        self.features = None
+        self.orig_h = None
+        self.orig_w = None
+        self.input_h = None
+        self.input_w = None

fooocus_colab.ipynb

@@ -8,7 +8,7 @@
    },
    "outputs": [],
    "source": [
-    "!pip install pygit2==1.12.2\n",
+    "!pip install pygit2==1.15.1\n",
     "%cd /content\n",
     "!git clone https://github.com/lllyasviel/Fooocus.git\n",
     "%cd /content/Fooocus\n",

fooocus_version.py

@@ -1 +1 @@
-version = '2.4.1'
+version = '2.5.5'

javascript/zoom.js

@@ -642,4 +642,5 @@ onUiLoaded(async() => {
     }
 
     applyZoomAndPan("#inpaint_canvas");
+    applyZoomAndPan("#inpaint_mask_canvas");
 });

language/en.json

@@ -11,11 +11,13 @@
     "Image Prompt": "Image Prompt",
     "Inpaint or Outpaint": "Inpaint or Outpaint",
     "Outpaint Direction": "Outpaint Direction",
+    "Enable Advanced Masking Features": "Enable Advanced Masking Features",
     "Method": "Method",
     "Describe": "Describe",
     "Content Type": "Content Type",
     "Photograph": "Photograph",
     "Art/Anime": "Art/Anime",
+    "Apply Styles": "Apply Styles",
     "Describe this Image into Prompt": "Describe this Image into Prompt",
     "Image Size and Recommended Size": "Image Size and Recommended Size",
     "Upscale or Variation:": "Upscale or Variation:",
@@ -25,7 +27,7 @@
     "Upscale (1.5x)": "Upscale (1.5x)",
     "Upscale (2x)": "Upscale (2x)",
     "Upscale (Fast 2x)": "Upscale (Fast 2x)",
-    "\ud83d\udcd4 Document": "\uD83D\uDCD4 Document",
+    "\ud83d\udcd4 Documentation": "\uD83D\uDCD4 Documentation",
     "Image": "Image",
     "Stop At": "Stop At",
     "Weight": "Weight",
@@ -44,8 +46,11 @@
     "Top": "Top",
     "Bottom": "Bottom",
     "* \"Inpaint or Outpaint\" is powered by the sampler \"DPMPP Fooocus Seamless 2M SDE Karras Inpaint Sampler\" (beta)": "* \"Inpaint or Outpaint\" is powered by the sampler \"DPMPP Fooocus Seamless 2M SDE Karras Inpaint Sampler\" (beta)",
-    "Setting": "Setting",
+    "Advanced options": "Advanced options",
+    "Generate mask from image": "Generate mask from image",
+    "Settings": "Settings",
     "Style": "Style",
+    "Styles": "Styles",
     "Preset": "Preset",
     "Performance": "Performance",
     "Speed": "Speed",
@@ -64,6 +69,9 @@
     "Read wildcards in order": "Read wildcards in order",
     "Black Out NSFW": "Black Out NSFW",
     "Use black image if NSFW is detected.": "Use black image if NSFW is detected.",
+    "Save only final enhanced image": "Save only final enhanced image",
+    "Save Metadata to Images": "Save Metadata to Images",
+    "Adds parameters to generated images allowing manual regeneration.": "Adds parameters to generated images allowing manual regeneration.",
     "\ud83d\udcda History Log": "\uD83D\uDCDA History Log",
     "Image Style": "Image Style",
     "Fooocus V2": "Fooocus V2",
@@ -279,7 +287,7 @@
     "Volumetric Lighting": "Volumetric Lighting",
     "Watercolor 2": "Watercolor 2",
     "Whimsical And Playful": "Whimsical And Playful",
-    "Model": "Model",
+    "Models": "Models",
     "Base Model (SDXL only)": "Base Model (SDXL only)",
     "sd_xl_base_1.0_0.9vae.safetensors": "sd_xl_base_1.0_0.9vae.safetensors",
     "bluePencilXL_v009.safetensors": "bluePencilXL_v009.safetensors",
@@ -367,10 +375,14 @@
     "Disable preview during generation.": "Disable preview during generation.",
     "Disable Intermediate Results": "Disable Intermediate Results",
     "Disable intermediate results during generation, only show final gallery.": "Disable intermediate results during generation, only show final gallery.",
+    "Debug Inpaint Preprocessing": "Debug Inpaint Preprocessing",
+    "Debug GroundingDINO": "Debug GroundingDINO",
+    "Used for SAM object detection and box generation": "Used for SAM object detection and box generation",
+    "GroundingDINO Box Erode or Dilate": "GroundingDINO Box Erode or Dilate",
     "Inpaint Engine": "Inpaint Engine",
     "v1": "v1",
-    "Version of Fooocus inpaint model": "Version of Fooocus inpaint model",
     "v2.5": "v2.5",
+    "v2.6": "v2.6",
     "Control Debug": "Control Debug",
     "Debug Preprocessors": "Debug Preprocessors",
     "Mixing Image Prompt and Vary/Upscale": "Mixing Image Prompt and Vary/Upscale",
@@ -400,5 +412,74 @@
     "Image Prompt parameters are not included. Use png and a1111 for compatibility with Civitai.": "Image Prompt parameters are not included. Use png and a1111 for compatibility with Civitai.",
     "fooocus (json)": "fooocus (json)",
     "a1111 (plain text)": "a1111 (plain text)",
-    "Unsupported image type in input": "Unsupported image type in input"
+    "Unsupported image type in input": "Unsupported image type in input",
+    "Enhance": "Enhance",
+    "Detection prompt": "Detection prompt",
+    "Detection Prompt Quick List": "Detection Prompt Quick List",
+    "Maximum number of detections": "Maximum number of detections",
+    "Use with Enhance, skips image generation": "Use with Enhance, skips image generation",
+    "Order of Processing": "Order of Processing",
+    "Use before to enhance small details and after to enhance large areas.": "Use before to enhance small details and after to enhance large areas.",
+    "Before First Enhancement": "Before First Enhancement",
+    "After Last Enhancement": "After Last Enhancement",
+    "Prompt Type": "Prompt Type",
+    "Choose which prompt to use for Upscale or Variation.": "Choose which prompt to use for Upscale or Variation.",
+    "Original Prompts": "Original Prompts",
+    "Last Filled Enhancement Prompts": "Last Filled Enhancement Prompts",
+    "Enable": "Enable",
+    "Describe what you want to detect.": "Describe what you want to detect.",
+    "Enhancement positive prompt": "Enhancement positive prompt",
+    "Uses original prompt instead if empty.": "Uses original prompt instead if empty.",
+    "Enhancement negative prompt": "Enhancement negative prompt",
+    "Uses original negative prompt instead if empty.": "Uses original negative prompt instead if empty.",
+    "Detection": "Detection",
+    "u2net": "u2net",
+    "u2netp": "u2netp",
+    "u2net_human_seg": "u2net_human_seg",
+    "u2net_cloth_seg": "u2net_cloth_seg",
+    "silueta": "silueta",
+    "isnet-general-use": "isnet-general-use",
+    "isnet-anime": "isnet-anime",
+    "sam": "sam",
+    "Mask generation model": "Mask generation model",
+    "Cloth category": "Cloth category",
+    "Use singular whenever possible": "Use singular whenever possible",
+    "full": "full",
+    "upper": "upper",
+    "lower": "lower",
+    "SAM Options": "SAM Options",
+    "SAM model": "SAM model",
+    "vit_b": "vit_b",
+    "vit_l": "vit_l",
+    "vit_h": "vit_h",
+    "Box Threshold": "Box Threshold",
+    "Text Threshold": "Text Threshold",
+    "Set to 0 to detect all": "Set to 0 to detect all",
+    "Inpaint": "Inpaint",
+    "Inpaint or Outpaint (default)": "Inpaint or Outpaint (default)",
+    "Improve Detail (face, hand, eyes, etc.)": "Improve Detail (face, hand, eyes, etc.)",
+    "Modify Content (add objects, change background, etc.)": "Modify Content (add objects, change background, etc.)",
+    "Disable initial latent in inpaint": "Disable initial latent in inpaint",
+    "Version of Fooocus inpaint model. If set, use performance Quality or Speed (no performance LoRAs) for best results.": "Version of Fooocus inpaint model. If set, use performance Quality or Speed (no performance LoRAs) for best results.",
+    "Inpaint Denoising Strength": "Inpaint Denoising Strength",
+    "Same as the denoising strength in A1111 inpaint. Only used in inpaint, not used in outpaint. (Outpaint always use 1.0)": "Same as the denoising strength in A1111 inpaint. Only used in inpaint, not used in outpaint. (Outpaint always use 1.0)",
+    "Inpaint Respective Field": "Inpaint Respective Field",
+    "The area to inpaint. Value 0 is same as \"Only Masked\" in A1111. Value 1 is same as \"Whole Image\" in A1111. Only used in inpaint, not used in outpaint. (Outpaint always use 1.0)": "The area to inpaint. Value 0 is same as \"Only Masked\" in A1111. Value 1 is same as \"Whole Image\" in A1111. Only used in inpaint, not used in outpaint. (Outpaint always use 1.0)",
+    "Mask Erode or Dilate": "Mask Erode or Dilate",
+    "Positive value will make white area in the mask larger, negative value will make white area smaller. (default is 0, always processed before any mask invert)": "Positive value will make white area in the mask larger, negative value will make white area smaller. (default is 0, always processed before any mask invert)",
+    "Invert Mask When Generating": "Invert Mask When Generating",
+    "Debug Enhance Masks": "Debug Enhance Masks",
+    "Show enhance masks in preview and final results": "Show enhance masks in preview and final results",
+    "Use GroundingDINO boxes instead of more detailed SAM masks": "Use GroundingDINO boxes instead of more detailed SAM masks",
+    "highly detailed face": "highly detailed face",
+    "detailed girl face": "detailed girl face",
+    "detailed man face": "detailed man face",
+    "detailed hand": "detailed hand",
+    "beautiful eyes": "beautiful eyes",
+    "face": "face",
+    "eye": "eye",
+    "mouth": "mouth",
+    "hair": "hair",
+    "hand": "hand",
+    "body": "body"
 }

launch.py

@@ -80,11 +80,14 @@ if args.gpu_device_id is not None:
     os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_device_id)
     print("Set device to:", args.gpu_device_id)
 
-if args.hf_mirror is not None : 
+if args.hf_mirror is not None:
     os.environ['HF_MIRROR'] = str(args.hf_mirror)
     print("Set hf_mirror to:", args.hf_mirror)
 
 from modules import config
+from modules.hash_cache import init_cache
+
+os.environ["U2NET_HOME"] = config.path_inpaint
 
 os.environ['GRADIO_TEMP_DIR'] = config.temp_path
 
@@ -97,7 +100,9 @@ if config.temp_path_cleanup_on_launch:
         print(f"[Cleanup] Failed to delete content of temp dir.")
 
 
-def download_models(default_model, previous_default_models, checkpoint_downloads, embeddings_downloads, lora_downloads):
+def download_models(default_model, previous_default_models, checkpoint_downloads, embeddings_downloads, lora_downloads, vae_downloads):
+    from modules.util import get_file_from_folder_list
+
     for file_name, url in vae_approx_filenames:
         load_file_from_url(url=url, model_dir=config.path_vae_approx, file_name=file_name)
 
@@ -112,9 +117,9 @@ def download_models(default_model, previous_default_models, checkpoint_downloads
         return default_model, checkpoint_downloads
 
     if not args.always_download_new_model:
-        if not os.path.exists(os.path.join(config.paths_checkpoints[0], default_model)):
+        if not os.path.isfile(get_file_from_folder_list(default_model, config.paths_checkpoints)):
             for alternative_model_name in previous_default_models:
-                if os.path.exists(os.path.join(config.paths_checkpoints[0], alternative_model_name)):
+                if os.path.isfile(get_file_from_folder_list(alternative_model_name, config.paths_checkpoints)):
                     print(f'You do not have [{default_model}] but you have [{alternative_model_name}].')
                     print(f'Fooocus will use [{alternative_model_name}] to avoid downloading new models, '
                           f'but you are not using the latest models.')
@@ -124,17 +129,24 @@ def download_models(default_model, previous_default_models, checkpoint_downloads
                     break
 
     for file_name, url in checkpoint_downloads.items():
-        load_file_from_url(url=url, model_dir=config.paths_checkpoints[0], file_name=file_name)
+        model_dir = os.path.dirname(get_file_from_folder_list(file_name, config.paths_checkpoints))
+        load_file_from_url(url=url, model_dir=model_dir, file_name=file_name)
     for file_name, url in embeddings_downloads.items():
         load_file_from_url(url=url, model_dir=config.path_embeddings, file_name=file_name)
     for file_name, url in lora_downloads.items():
-        load_file_from_url(url=url, model_dir=config.paths_loras[0], file_name=file_name)
+        model_dir = os.path.dirname(get_file_from_folder_list(file_name, config.paths_loras))
+        load_file_from_url(url=url, model_dir=model_dir, file_name=file_name)
+    for file_name, url in vae_downloads.items():
+        load_file_from_url(url=url, model_dir=config.path_vae, file_name=file_name)
 
     return default_model, checkpoint_downloads
 
 
 config.default_base_model_name, config.checkpoint_downloads = download_models(
     config.default_base_model_name, config.previous_default_models, config.checkpoint_downloads,
-    config.embeddings_downloads, config.lora_downloads)
+    config.embeddings_downloads, config.lora_downloads, config.vae_downloads)
+
+config.update_files()
+init_cache(config.model_filenames, config.paths_checkpoints, config.lora_filenames, config.paths_loras)
 
 from webui import *

ldm_patched/contrib/external_advanced_samplers.py

@@ -1,61 +0,0 @@
-import ldm_patched.modules.samplers
-import ldm_patched.modules.utils
-import torch
-import numpy as np
-from tqdm.auto import trange, tqdm
-import math
-
-
-@torch.no_grad()
-def sample_lcm_upscale(model, x, sigmas, extra_args=None, callback=None, disable=None, total_upscale=2.0, upscale_method="bislerp", upscale_steps=None):
-    extra_args = {} if extra_args is None else extra_args
-
-    if upscale_steps is None:
-        upscale_steps = max(len(sigmas) // 2 + 1, 2)
-    else:
-        upscale_steps += 1
-        upscale_steps = min(upscale_steps, len(sigmas) + 1)
-
-    upscales = np.linspace(1.0, total_upscale, upscale_steps)[1:]
-
-    orig_shape = x.size()
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-
-        x = denoised
-        if i < len(upscales):
-            x = ldm_patched.modules.utils.common_upscale(x, round(orig_shape[-1] * upscales[i]), round(orig_shape[-2] * upscales[i]), upscale_method, "disabled")
-
-        if sigmas[i + 1] > 0:
-            x += sigmas[i + 1] * torch.randn_like(x)
-    return x
-
-
-class SamplerLCMUpscale:
-    upscale_methods = ["bislerp", "nearest-exact", "bilinear", "area", "bicubic"]
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"scale_ratio": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 20.0, "step": 0.01}),
-                     "scale_steps": ("INT", {"default": -1, "min": -1, "max": 1000, "step": 1}),
-                     "upscale_method": (s.upscale_methods,),
-                      }
-               }
-    RETURN_TYPES = ("SAMPLER",)
-    CATEGORY = "sampling/custom_sampling/samplers"
-
-    FUNCTION = "get_sampler"
-
-    def get_sampler(self, scale_ratio, scale_steps, upscale_method):
-        if scale_steps < 0:
-            scale_steps = None
-        sampler = ldm_patched.modules.samplers.KSAMPLER(sample_lcm_upscale, extra_options={"total_upscale": scale_ratio, "upscale_steps": scale_steps, "upscale_method": upscale_method})
-        return (sampler, )
-
-NODE_CLASS_MAPPINGS = {
-    "SamplerLCMUpscale": SamplerLCMUpscale,
-}

ldm_patched/contrib/external_align_your_steps.py

@@ -1,3 +1,5 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py
+
 #from: https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html
 import numpy as np
 import torch
@@ -50,4 +52,4 @@ class AlignYourStepsScheduler:
 
 NODE_CLASS_MAPPINGS = {
     "AlignYourStepsScheduler": AlignYourStepsScheduler,
-}
+}

ldm_patched/contrib/external_attention_multiply.py

@@ -1,120 +0,0 @@
-
-def attention_multiply(attn, model, q, k, v, out):
-    m = model.clone()
-    sd = model.model_state_dict()
-
-    for key in sd:
-        if key.endswith("{}.to_q.bias".format(attn)) or key.endswith("{}.to_q.weight".format(attn)):
-            m.add_patches({key: (None,)}, 0.0, q)
-        if key.endswith("{}.to_k.bias".format(attn)) or key.endswith("{}.to_k.weight".format(attn)):
-            m.add_patches({key: (None,)}, 0.0, k)
-        if key.endswith("{}.to_v.bias".format(attn)) or key.endswith("{}.to_v.weight".format(attn)):
-            m.add_patches({key: (None,)}, 0.0, v)
-        if key.endswith("{}.to_out.0.bias".format(attn)) or key.endswith("{}.to_out.0.weight".format(attn)):
-            m.add_patches({key: (None,)}, 0.0, out)
-
-    return m
-
-
-class UNetSelfAttentionMultiply:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "q": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              "k": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              "v": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              "out": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "_for_testing/attention_experiments"
-
-    def patch(self, model, q, k, v, out):
-        m = attention_multiply("attn1", model, q, k, v, out)
-        return (m, )
-
-class UNetCrossAttentionMultiply:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "q": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              "k": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              "v": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              "out": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "_for_testing/attention_experiments"
-
-    def patch(self, model, q, k, v, out):
-        m = attention_multiply("attn2", model, q, k, v, out)
-        return (m, )
-
-class CLIPAttentionMultiply:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "clip": ("CLIP",),
-                              "q": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              "k": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              "v": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              "out": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              }}
-    RETURN_TYPES = ("CLIP",)
-    FUNCTION = "patch"
-
-    CATEGORY = "_for_testing/attention_experiments"
-
-    def patch(self, clip, q, k, v, out):
-        m = clip.clone()
-        sd = m.patcher.model_state_dict()
-
-        for key in sd:
-            if key.endswith("self_attn.q_proj.weight") or key.endswith("self_attn.q_proj.bias"):
-                m.add_patches({key: (None,)}, 0.0, q)
-            if key.endswith("self_attn.k_proj.weight") or key.endswith("self_attn.k_proj.bias"):
-                m.add_patches({key: (None,)}, 0.0, k)
-            if key.endswith("self_attn.v_proj.weight") or key.endswith("self_attn.v_proj.bias"):
-                m.add_patches({key: (None,)}, 0.0, v)
-            if key.endswith("self_attn.out_proj.weight") or key.endswith("self_attn.out_proj.bias"):
-                m.add_patches({key: (None,)}, 0.0, out)
-        return (m, )
-
-class UNetTemporalAttentionMultiply:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "self_structural": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              "self_temporal": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              "cross_structural": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              "cross_temporal": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "_for_testing/attention_experiments"
-
-    def patch(self, model, self_structural, self_temporal, cross_structural, cross_temporal):
-        m = model.clone()
-        sd = model.model_state_dict()
-
-        for k in sd:
-            if (k.endswith("attn1.to_out.0.bias") or k.endswith("attn1.to_out.0.weight")):
-                if '.time_stack.' in k:
-                    m.add_patches({k: (None,)}, 0.0, self_temporal)
-                else:
-                    m.add_patches({k: (None,)}, 0.0, self_structural)
-            elif (k.endswith("attn2.to_out.0.bias") or k.endswith("attn2.to_out.0.weight")):
-                if '.time_stack.' in k:
-                    m.add_patches({k: (None,)}, 0.0, cross_temporal)
-                else:
-                    m.add_patches({k: (None,)}, 0.0, cross_structural)
-        return (m, )
-
-NODE_CLASS_MAPPINGS = {
-    "UNetSelfAttentionMultiply": UNetSelfAttentionMultiply,
-    "UNetCrossAttentionMultiply": UNetCrossAttentionMultiply,
-    "CLIPAttentionMultiply": CLIPAttentionMultiply,
-    "UNetTemporalAttentionMultiply": UNetTemporalAttentionMultiply,
-}

ldm_patched/contrib/external_canny.py

@@ -1,6 +1,282 @@
-from kornia.filters import canny
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
+#From https://github.com/kornia/kornia
+import math
+
+import torch
+import torch.nn.functional as F
 import ldm_patched.modules.model_management
 
+def get_canny_nms_kernel(device=None, dtype=None):
+    """Utility function that returns 3x3 kernels for the Canny Non-maximal suppression."""
+    return torch.tensor(
+        [
+            [[[0.0, 0.0, 0.0], [0.0, 1.0, -1.0], [0.0, 0.0, 0.0]]],
+            [[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, -1.0]]],
+            [[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, -1.0, 0.0]]],
+            [[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [-1.0, 0.0, 0.0]]],
+            [[[0.0, 0.0, 0.0], [-1.0, 1.0, 0.0], [0.0, 0.0, 0.0]]],
+            [[[-1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]]],
+            [[[0.0, -1.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]]],
+            [[[0.0, 0.0, -1.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]]],
+        ],
+        device=device,
+        dtype=dtype,
+    )
+
+
+def get_hysteresis_kernel(device=None, dtype=None):
+    """Utility function that returns the 3x3 kernels for the Canny hysteresis."""
+    return torch.tensor(
+        [
+            [[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 0.0]]],
+            [[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]],
+            [[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 1.0, 0.0]]],
+            [[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]],
+            [[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0]]],
+            [[[1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]],
+            [[[0.0, 1.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]],
+            [[[0.0, 0.0, 1.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]],
+        ],
+        device=device,
+        dtype=dtype,
+    )
+
+def gaussian_blur_2d(img, kernel_size, sigma):
+    ksize_half = (kernel_size - 1) * 0.5
+
+    x = torch.linspace(-ksize_half, ksize_half, steps=kernel_size)
+
+    pdf = torch.exp(-0.5 * (x / sigma).pow(2))
+
+    x_kernel = pdf / pdf.sum()
+    x_kernel = x_kernel.to(device=img.device, dtype=img.dtype)
+
+    kernel2d = torch.mm(x_kernel[:, None], x_kernel[None, :])
+    kernel2d = kernel2d.expand(img.shape[-3], 1, kernel2d.shape[0], kernel2d.shape[1])
+
+    padding = [kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2]
+
+    img = torch.nn.functional.pad(img, padding, mode="reflect")
+    img = torch.nn.functional.conv2d(img, kernel2d, groups=img.shape[-3])
+
+    return img
+
+def get_sobel_kernel2d(device=None, dtype=None):
+    kernel_x = torch.tensor([[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]], device=device, dtype=dtype)
+    kernel_y = kernel_x.transpose(0, 1)
+    return torch.stack([kernel_x, kernel_y])
+
+def spatial_gradient(input, normalized: bool = True):
+    r"""Compute the first order image derivative in both x and y using a Sobel operator.
+    .. image:: _static/img/spatial_gradient.png
+    Args:
+        input: input image tensor with shape :math:`(B, C, H, W)`.
+        mode: derivatives modality, can be: `sobel` or `diff`.
+        order: the order of the derivatives.
+        normalized: whether the output is normalized.
+    Return:
+        the derivatives of the input feature map. with shape :math:`(B, C, 2, H, W)`.
+    .. note::
+       See a working example `here <https://kornia.readthedocs.io/en/latest/
+       filtering_edges.html>`__.
+    Examples:
+        >>> input = torch.rand(1, 3, 4, 4)
+        >>> output = spatial_gradient(input)  # 1x3x2x4x4
+        >>> output.shape
+        torch.Size([1, 3, 2, 4, 4])
+    """
+    # KORNIA_CHECK_IS_TENSOR(input)
+    # KORNIA_CHECK_SHAPE(input, ['B', 'C', 'H', 'W'])
+
+    # allocate kernel
+    kernel = get_sobel_kernel2d(device=input.device, dtype=input.dtype)
+    if normalized:
+        kernel = normalize_kernel2d(kernel)
+
+    # prepare kernel
+    b, c, h, w = input.shape
+    tmp_kernel = kernel[:, None, ...]
+
+    # Pad with "replicate for spatial dims, but with zeros for channel
+    spatial_pad = [kernel.size(1) // 2, kernel.size(1) // 2, kernel.size(2) // 2, kernel.size(2) // 2]
+    out_channels: int = 2
+    padded_inp = torch.nn.functional.pad(input.reshape(b * c, 1, h, w), spatial_pad, 'replicate')
+    out = F.conv2d(padded_inp, tmp_kernel, groups=1, padding=0, stride=1)
+    return out.reshape(b, c, out_channels, h, w)
+
+def rgb_to_grayscale(image, rgb_weights = None):
+    r"""Convert a RGB image to grayscale version of image.
+
+    .. image:: _static/img/rgb_to_grayscale.png
+
+    The image data is assumed to be in the range of (0, 1).
+
+    Args:
+        image: RGB image to be converted to grayscale with shape :math:`(*,3,H,W)`.
+        rgb_weights: Weights that will be applied on each channel (RGB).
+            The sum of the weights should add up to one.
+    Returns:
+        grayscale version of the image with shape :math:`(*,1,H,W)`.
+
+    .. note::
+       See a working example `here <https://kornia.readthedocs.io/en/latest/
+       color_conversions.html>`__.
+
+    Example:
+        >>> input = torch.rand(2, 3, 4, 5)
+        >>> gray = rgb_to_grayscale(input) # 2x1x4x5
+    """
+
+    if len(image.shape) < 3 or image.shape[-3] != 3:
+        raise ValueError(f"Input size must have a shape of (*, 3, H, W). Got {image.shape}")
+
+    if rgb_weights is None:
+        # 8 bit images
+        if image.dtype == torch.uint8:
+            rgb_weights = torch.tensor([76, 150, 29], device=image.device, dtype=torch.uint8)
+        # floating point images
+        elif image.dtype in (torch.float16, torch.float32, torch.float64):
+            rgb_weights = torch.tensor([0.299, 0.587, 0.114], device=image.device, dtype=image.dtype)
+        else:
+            raise TypeError(f"Unknown data type: {image.dtype}")
+    else:
+        # is tensor that we make sure is in the same device/dtype
+        rgb_weights = rgb_weights.to(image)
+
+    # unpack the color image channels with RGB order
+    r: Tensor = image[..., 0:1, :, :]
+    g: Tensor = image[..., 1:2, :, :]
+    b: Tensor = image[..., 2:3, :, :]
+
+    w_r, w_g, w_b = rgb_weights.unbind()
+    return w_r * r + w_g * g + w_b * b
+
+def canny(
+    input,
+    low_threshold = 0.1,
+    high_threshold = 0.2,
+    kernel_size  = 5,
+    sigma = 1,
+    hysteresis = True,
+    eps = 1e-6,
+):
+    r"""Find edges of the input image and filters them using the Canny algorithm.
+    .. image:: _static/img/canny.png
+    Args:
+        input: input image tensor with shape :math:`(B,C,H,W)`.
+        low_threshold: lower threshold for the hysteresis procedure.
+        high_threshold: upper threshold for the hysteresis procedure.
+        kernel_size: the size of the kernel for the gaussian blur.
+        sigma: the standard deviation of the kernel for the gaussian blur.
+        hysteresis: if True, applies the hysteresis edge tracking.
+            Otherwise, the edges are divided between weak (0.5) and strong (1) edges.
+        eps: regularization number to avoid NaN during backprop.
+    Returns:
+        - the canny edge magnitudes map, shape of :math:`(B,1,H,W)`.
+        - the canny edge detection filtered by thresholds and hysteresis, shape of :math:`(B,1,H,W)`.
+    .. note::
+       See a working example `here <https://kornia.readthedocs.io/en/latest/
+       canny.html>`__.
+    Example:
+        >>> input = torch.rand(5, 3, 4, 4)
+        >>> magnitude, edges = canny(input)  # 5x3x4x4
+        >>> magnitude.shape
+        torch.Size([5, 1, 4, 4])
+        >>> edges.shape
+        torch.Size([5, 1, 4, 4])
+    """
+    # KORNIA_CHECK_IS_TENSOR(input)
+    # KORNIA_CHECK_SHAPE(input, ['B', 'C', 'H', 'W'])
+    # KORNIA_CHECK(
+    #     low_threshold <= high_threshold,
+    #     "Invalid input thresholds. low_threshold should be smaller than the high_threshold. Got: "
+    #     f"{low_threshold}>{high_threshold}",
+    # )
+    # KORNIA_CHECK(0 < low_threshold < 1, f'Invalid low threshold. Should be in range (0, 1). Got: {low_threshold}')
+    # KORNIA_CHECK(0 < high_threshold < 1, f'Invalid high threshold. Should be in range (0, 1). Got: {high_threshold}')
+
+    device = input.device
+    dtype = input.dtype
+
+    # To Grayscale
+    if input.shape[1] == 3:
+        input = rgb_to_grayscale(input)
+
+    # Gaussian filter
+    blurred: Tensor = gaussian_blur_2d(input, kernel_size, sigma)
+
+    # Compute the gradients
+    gradients: Tensor = spatial_gradient(blurred, normalized=False)
+
+    # Unpack the edges
+    gx: Tensor = gradients[:, :, 0]
+    gy: Tensor = gradients[:, :, 1]
+
+    # Compute gradient magnitude and angle
+    magnitude: Tensor = torch.sqrt(gx * gx + gy * gy + eps)
+    angle: Tensor = torch.atan2(gy, gx)
+
+    # Radians to Degrees
+    angle = 180.0 * angle / math.pi
+
+    # Round angle to the nearest 45 degree
+    angle = torch.round(angle / 45) * 45
+
+    # Non-maximal suppression
+    nms_kernels: Tensor = get_canny_nms_kernel(device, dtype)
+    nms_magnitude: Tensor = F.conv2d(magnitude, nms_kernels, padding=nms_kernels.shape[-1] // 2)
+
+    # Get the indices for both directions
+    positive_idx: Tensor = (angle / 45) % 8
+    positive_idx = positive_idx.long()
+
+    negative_idx: Tensor = ((angle / 45) + 4) % 8
+    negative_idx = negative_idx.long()
+
+    # Apply the non-maximum suppression to the different directions
+    channel_select_filtered_positive: Tensor = torch.gather(nms_magnitude, 1, positive_idx)
+    channel_select_filtered_negative: Tensor = torch.gather(nms_magnitude, 1, negative_idx)
+
+    channel_select_filtered: Tensor = torch.stack(
+        [channel_select_filtered_positive, channel_select_filtered_negative], 1
+    )
+
+    is_max: Tensor = channel_select_filtered.min(dim=1)[0] > 0.0
+
+    magnitude = magnitude * is_max
+
+    # Threshold
+    edges: Tensor = F.threshold(magnitude, low_threshold, 0.0)
+
+    low: Tensor = magnitude > low_threshold
+    high: Tensor = magnitude > high_threshold
+
+    edges = low * 0.5 + high * 0.5
+    edges = edges.to(dtype)
+
+    # Hysteresis
+    if hysteresis:
+        edges_old: Tensor = -torch.ones(edges.shape, device=edges.device, dtype=dtype)
+        hysteresis_kernels: Tensor = get_hysteresis_kernel(device, dtype)
+
+        while ((edges_old - edges).abs() != 0).any():
+            weak: Tensor = (edges == 0.5).float()
+            strong: Tensor = (edges == 1).float()
+
+            hysteresis_magnitude: Tensor = F.conv2d(
+                edges, hysteresis_kernels, padding=hysteresis_kernels.shape[-1] // 2
+            )
+            hysteresis_magnitude = (hysteresis_magnitude == 1).any(1, keepdim=True).to(dtype)
+            hysteresis_magnitude = hysteresis_magnitude * weak + strong
+
+            edges_old = edges.clone()
+            edges = hysteresis_magnitude + (hysteresis_magnitude == 0) * weak * 0.5
+
+        edges = hysteresis_magnitude
+
+    return magnitude, edges
+
 
 class Canny:
     @classmethod

ldm_patched/contrib/external_clip_sdxl.py

@@ -1,5 +1,7 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import torch
-from external import MAX_RESOLUTION
+from ldm_patched.contrib.external import MAX_RESOLUTION
 
 class CLIPTextEncodeSDXLRefiner:
     @classmethod
@@ -8,7 +10,7 @@ class CLIPTextEncodeSDXLRefiner:
             "ascore": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 1000.0, "step": 0.01}),
             "width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
             "height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
-            "text": ("STRING", {"multiline": True, "dynamicPrompts": True}), "clip": ("CLIP", ),
+            "text": ("STRING", {"multiline": True}), "clip": ("CLIP", ),
             }}
     RETURN_TYPES = ("CONDITIONING",)
     FUNCTION = "encode"
@@ -30,8 +32,8 @@ class CLIPTextEncodeSDXL:
             "crop_h": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}),
             "target_width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
             "target_height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
-            "text_g": ("STRING", {"multiline": True, "dynamicPrompts": True}), "clip": ("CLIP", ),
-            "text_l": ("STRING", {"multiline": True, "dynamicPrompts": True}), "clip": ("CLIP", ),
+            "text_g": ("STRING", {"multiline": True, "default": "CLIP_G"}), "clip": ("CLIP", ),
+            "text_l": ("STRING", {"multiline": True, "default": "CLIP_L"}), "clip": ("CLIP", ),
             }}
     RETURN_TYPES = ("CONDITIONING",)
     FUNCTION = "encode"

ldm_patched/contrib/external_compositing.py

@@ -1,3 +1,5 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import numpy as np
 import torch
 import ldm_patched.modules.utils

ldm_patched/contrib/external_cond.py

@@ -1,25 +0,0 @@
-
-
-class CLIPTextEncodeControlnet:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"clip": ("CLIP", ), "conditioning": ("CONDITIONING", ), "text": ("STRING", {"multiline": True, "dynamicPrompts": True})}}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "encode"
-
-    CATEGORY = "_for_testing/conditioning"
-
-    def encode(self, clip, conditioning, text):
-        tokens = clip.tokenize(text)
-        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
-        c = []
-        for t in conditioning:
-            n = [t[0], t[1].copy()]
-            n[1]['cross_attn_controlnet'] = cond
-            n[1]['pooled_output_controlnet'] = pooled
-            c.append(n)
-        return (c, )
-
-NODE_CLASS_MAPPINGS = {
-    "CLIPTextEncodeControlnet": CLIPTextEncodeControlnet
-}

ldm_patched/contrib/external_custom_sampler.py

@@ -1,10 +1,11 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import ldm_patched.modules.samplers
 import ldm_patched.modules.sample
 from ldm_patched.k_diffusion import sampling as k_diffusion_sampling
-import ldm_patched.utils.latent_visualization as latent_preview
+import ldm_patched.utils.latent_visualization
 import torch
 import ldm_patched.modules.utils
-import ldm_patched.utils.node_helpers as node_helpers
 
 
 class BasicScheduler:
@@ -25,11 +26,10 @@ class BasicScheduler:
     def get_sigmas(self, model, scheduler, steps, denoise):
         total_steps = steps
         if denoise < 1.0:
-            if denoise <= 0.0:
-                return (torch.FloatTensor([]),)
             total_steps = int(steps/denoise)
 
-        sigmas = ldm_patched.modules.samplers.calculate_sigmas(model.get_model_object("model_sampling"), scheduler, total_steps).cpu()
+        ldm_patched.modules.model_management.load_models_gpu([model])
+        sigmas = ldm_patched.modules.samplers.calculate_sigmas_scheduler(model.model, scheduler, total_steps).cpu()
         sigmas = sigmas[-(steps + 1):]
         return (sigmas, )
 
@@ -39,8 +39,8 @@ class KarrasScheduler:
     def INPUT_TYPES(s):
         return {"required":
                     {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                     "sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 5000.0, "step":0.01, "round": False}),
-                     "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 5000.0, "step":0.01, "round": False}),
+                     "sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
+                     "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
                      "rho": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
                     }
                }
@@ -58,8 +58,8 @@ class ExponentialScheduler:
     def INPUT_TYPES(s):
         return {"required":
                     {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                     "sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 5000.0, "step":0.01, "round": False}),
-                     "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 5000.0, "step":0.01, "round": False}),
+                     "sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
+                     "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
                     }
                }
     RETURN_TYPES = ("SIGMAS",)
@@ -76,8 +76,8 @@ class PolyexponentialScheduler:
     def INPUT_TYPES(s):
         return {"required":
                     {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                     "sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 5000.0, "step":0.01, "round": False}),
-                     "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 5000.0, "step":0.01, "round": False}),
+                     "sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
+                     "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
                      "rho": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
                     }
                }
@@ -107,7 +107,7 @@ class SDTurboScheduler:
     def get_sigmas(self, model, steps, denoise):
         start_step = 10 - int(10 * denoise)
         timesteps = torch.flip(torch.arange(1, 11) * 100 - 1, (0,))[start_step:start_step + steps]
-        sigmas = model.get_model_object("model_sampling").sigma(timesteps)
+        sigmas = model.model_sampling.sigma(timesteps)
         sigmas = torch.cat([sigmas, sigmas.new_zeros([1])])
         return (sigmas, )
 
@@ -116,8 +116,8 @@ class VPScheduler:
     def INPUT_TYPES(s):
         return {"required":
                     {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                     "beta_d": ("FLOAT", {"default": 19.9, "min": 0.0, "max": 5000.0, "step":0.01, "round": False}), #TODO: fix default values
-                     "beta_min": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 5000.0, "step":0.01, "round": False}),
+                     "beta_d": ("FLOAT", {"default": 19.9, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}), #TODO: fix default values
+                     "beta_min": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
                      "eps_s": ("FLOAT", {"default": 0.001, "min": 0.0, "max": 1.0, "step":0.0001, "round": False}),
                     }
                }
@@ -139,7 +139,6 @@ class SplitSigmas:
                      }
                 }
     RETURN_TYPES = ("SIGMAS","SIGMAS")
-    RETURN_NAMES = ("high_sigmas", "low_sigmas")
     CATEGORY = "sampling/custom_sampling/sigmas"
 
     FUNCTION = "get_sigmas"
@@ -149,27 +148,6 @@ class SplitSigmas:
         sigmas2 = sigmas[step:]
         return (sigmas1, sigmas2)
 
-class SplitSigmasDenoise:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"sigmas": ("SIGMAS", ),
-                    "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                     }
-                }
-    RETURN_TYPES = ("SIGMAS","SIGMAS")
-    RETURN_NAMES = ("high_sigmas", "low_sigmas")
-    CATEGORY = "sampling/custom_sampling/sigmas"
-
-    FUNCTION = "get_sigmas"
-
-    def get_sigmas(self, sigmas, denoise):
-        steps = max(sigmas.shape[-1] - 1, 0)
-        total_steps = round(steps * denoise)
-        sigmas1 = sigmas[:-(total_steps)]
-        sigmas2 = sigmas[-(total_steps + 1):]
-        return (sigmas1, sigmas2)
-
 class FlipSigmas:
     @classmethod
     def INPUT_TYPES(s):
@@ -183,9 +161,6 @@ class FlipSigmas:
     FUNCTION = "get_sigmas"
 
     def get_sigmas(self, sigmas):
-        if len(sigmas) == 0:
-            return (sigmas,)
-
         sigmas = sigmas.flip(0)
         if sigmas[0] == 0:
             sigmas[0] = 0.0001
@@ -207,28 +182,6 @@ class KSamplerSelect:
         sampler = ldm_patched.modules.samplers.sampler_object(sampler_name)
         return (sampler, )
 
-class SamplerDPMPP_3M_SDE:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "s_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "noise_device": (['gpu', 'cpu'], ),
-                      }
-               }
-    RETURN_TYPES = ("SAMPLER",)
-    CATEGORY = "sampling/custom_sampling/samplers"
-
-    FUNCTION = "get_sampler"
-
-    def get_sampler(self, eta, s_noise, noise_device):
-        if noise_device == 'cpu':
-            sampler_name = "dpmpp_3m_sde"
-        else:
-            sampler_name = "dpmpp_3m_sde_gpu"
-        sampler = ldm_patched.modules.samplers.ksampler(sampler_name, {"eta": eta, "s_noise": s_noise})
-        return (sampler, )
-
 class SamplerDPMPP_2M_SDE:
     @classmethod
     def INPUT_TYPES(s):
@@ -276,83 +229,24 @@ class SamplerDPMPP_SDE:
         sampler = ldm_patched.modules.samplers.ksampler(sampler_name, {"eta": eta, "s_noise": s_noise, "r": r})
         return (sampler, )
 
-class SamplerEulerAncestral:
+
+class SamplerTCD:
     @classmethod
     def INPUT_TYPES(s):
-        return {"required":
-                    {"eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "s_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                      }
-               }
+        return {
+            "required": {
+                "eta": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}),
+            }
+        }
     RETURN_TYPES = ("SAMPLER",)
     CATEGORY = "sampling/custom_sampling/samplers"
 
     FUNCTION = "get_sampler"
 
-    def get_sampler(self, eta, s_noise):
-        sampler = ldm_patched.modules.samplers.ksampler("euler_ancestral", {"eta": eta, "s_noise": s_noise})
+    def get_sampler(self, eta=0.3):
+        sampler = ldm_patched.modules.samplers.ksampler("tcd", {"eta": eta})
         return (sampler, )
 
-class SamplerLMS:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"order": ("INT", {"default": 4, "min": 1, "max": 100}),
-                      }
-               }
-    RETURN_TYPES = ("SAMPLER",)
-    CATEGORY = "sampling/custom_sampling/samplers"
-
-    FUNCTION = "get_sampler"
-
-    def get_sampler(self, order):
-        sampler = ldm_patched.modules.samplers.ksampler("lms", {"order": order})
-        return (sampler, )
-
-class SamplerDPMAdaptative:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"order": ("INT", {"default": 3, "min": 2, "max": 3}),
-                     "rtol": ("FLOAT", {"default": 0.05, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "atol": ("FLOAT", {"default": 0.0078, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "h_init": ("FLOAT", {"default": 0.05, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "pcoeff": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "icoeff": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "dcoeff": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "accept_safety": ("FLOAT", {"default": 0.81, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "eta": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "s_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                      }
-               }
-    RETURN_TYPES = ("SAMPLER",)
-    CATEGORY = "sampling/custom_sampling/samplers"
-
-    FUNCTION = "get_sampler"
-
-    def get_sampler(self, order, rtol, atol, h_init, pcoeff, icoeff, dcoeff, accept_safety, eta, s_noise):
-        sampler = ldm_patched.modules.samplers.ksampler("dpm_adaptive", {"order": order, "rtol": rtol, "atol": atol, "h_init": h_init, "pcoeff": pcoeff,
-                                                              "icoeff": icoeff, "dcoeff": dcoeff, "accept_safety": accept_safety, "eta": eta,
-                                                              "s_noise":s_noise })
-        return (sampler, )
-
-class Noise_EmptyNoise:
-    def __init__(self):
-        self.seed = 0
-
-    def generate_noise(self, input_latent):
-        latent_image = input_latent["samples"]
-        return torch.zeros(latent_image.shape, dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
-
-
-class Noise_RandomNoise:
-    def __init__(self, seed):
-        self.seed = seed
-
-    def generate_noise(self, input_latent):
-        latent_image = input_latent["samples"]
-        batch_inds = input_latent["batch_index"] if "batch_index" in input_latent else None
-        return ldm_patched.modules.sample.prepare_noise(latent_image, self.seed, batch_inds)
 
 class SamplerCustom:
     @classmethod
@@ -381,16 +275,17 @@ class SamplerCustom:
         latent = latent_image
         latent_image = latent["samples"]
         if not add_noise:
-            noise = Noise_EmptyNoise().generate_noise(latent)
+            noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
         else:
-            noise = Noise_RandomNoise(noise_seed).generate_noise(latent)
+            batch_inds = latent["batch_index"] if "batch_index" in latent else None
+            noise = ldm_patched.modules.sample.prepare_noise(latent_image, noise_seed, batch_inds)
 
         noise_mask = None
         if "noise_mask" in latent:
             noise_mask = latent["noise_mask"]
 
         x0_output = {}
-        callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output)
+        callback = ldm_patched.utils.latent_visualization.prepare_callback(model, sigmas.shape[-1] - 1, x0_output)
 
         disable_pbar = not ldm_patched.modules.utils.PROGRESS_BAR_ENABLED
         samples = ldm_patched.modules.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed)
@@ -404,207 +299,6 @@ class SamplerCustom:
             out_denoised = out
         return (out, out_denoised)
 
-class Guider_Basic(ldm_patched.modules.samplers.CFGGuider):
-    def set_conds(self, positive):
-        self.inner_set_conds({"positive": positive})
-
-class BasicGuider:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"model": ("MODEL",),
-                    "conditioning": ("CONDITIONING", ),
-                     }
-                }
-
-    RETURN_TYPES = ("GUIDER",)
-
-    FUNCTION = "get_guider"
-    CATEGORY = "sampling/custom_sampling/guiders"
-
-    def get_guider(self, model, conditioning):
-        guider = Guider_Basic(model)
-        guider.set_conds(conditioning)
-        return (guider,)
-
-class CFGGuider:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"model": ("MODEL",),
-                    "positive": ("CONDITIONING", ),
-                    "negative": ("CONDITIONING", ),
-                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
-                     }
-                }
-
-    RETURN_TYPES = ("GUIDER",)
-
-    FUNCTION = "get_guider"
-    CATEGORY = "sampling/custom_sampling/guiders"
-
-    def get_guider(self, model, positive, negative, cfg):
-        guider = ldm_patched.modules.samplers.CFGGuider(model)
-        guider.set_conds(positive, negative)
-        guider.set_cfg(cfg)
-        return (guider,)
-
-class Guider_DualCFG(ldm_patched.modules.samplers.CFGGuider):
-    def set_cfg(self, cfg1, cfg2):
-        self.cfg1 = cfg1
-        self.cfg2 = cfg2
-
-    def set_conds(self, positive, middle, negative):
-        middle = node_helpers.conditioning_set_values(middle, {"prompt_type": "negative"})
-        self.inner_set_conds({"positive": positive, "middle": middle, "negative": negative})
-
-    def predict_noise(self, x, timestep, model_options={}, seed=None):
-        negative_cond = self.conds.get("negative", None)
-        middle_cond = self.conds.get("middle", None)
-
-        out = ldm_patched.modules.samplers.calc_cond_batch(self.inner_model, [negative_cond, middle_cond, self.conds.get("positive", None)], x, timestep, model_options)
-        return ldm_patched.modules.samplers.cfg_function(self.inner_model, out[1], out[0], self.cfg2, x, timestep, model_options=model_options, cond=middle_cond, uncond=negative_cond) + (out[2] - out[1]) * self.cfg1
-
-class DualCFGGuider:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"model": ("MODEL",),
-                    "cond1": ("CONDITIONING", ),
-                    "cond2": ("CONDITIONING", ),
-                    "negative": ("CONDITIONING", ),
-                    "cfg_conds": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
-                    "cfg_cond2_negative": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
-                     }
-                }
-
-    RETURN_TYPES = ("GUIDER",)
-
-    FUNCTION = "get_guider"
-    CATEGORY = "sampling/custom_sampling/guiders"
-
-    def get_guider(self, model, cond1, cond2, negative, cfg_conds, cfg_cond2_negative):
-        guider = Guider_DualCFG(model)
-        guider.set_conds(cond1, cond2, negative)
-        guider.set_cfg(cfg_conds, cfg_cond2_negative)
-        return (guider,)
-
-class DisableNoise:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":{
-                     }
-                }
-
-    RETURN_TYPES = ("NOISE",)
-    FUNCTION = "get_noise"
-    CATEGORY = "sampling/custom_sampling/noise"
-
-    def get_noise(self):
-        return (Noise_EmptyNoise(),)
-
-
-class RandomNoise(DisableNoise):
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":{
-                    "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
-                     }
-                }
-
-    def get_noise(self, noise_seed):
-        return (Noise_RandomNoise(noise_seed),)
-
-
-class SamplerCustomAdvanced:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"noise": ("NOISE", ),
-                    "guider": ("GUIDER", ),
-                    "sampler": ("SAMPLER", ),
-                    "sigmas": ("SIGMAS", ),
-                    "latent_image": ("LATENT", ),
-                     }
-                }
-
-    RETURN_TYPES = ("LATENT","LATENT")
-    RETURN_NAMES = ("output", "denoised_output")
-
-    FUNCTION = "sample"
-
-    CATEGORY = "sampling/custom_sampling"
-
-    def sample(self, noise, guider, sampler, sigmas, latent_image):
-        latent = latent_image
-        latent_image = latent["samples"]
-
-        noise_mask = None
-        if "noise_mask" in latent:
-            noise_mask = latent["noise_mask"]
-
-        x0_output = {}
-        callback = latent_preview.prepare_callback(guider.model_patcher, sigmas.shape[-1] - 1, x0_output)
-
-        disable_pbar = not ldm_patched.modules.utils.PROGRESS_BAR_ENABLED
-        samples = guider.sample(noise.generate_noise(latent), latent_image, sampler, sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise.seed)
-        samples = samples.to(ldm_patched.modules.model_management.intermediate_device())
-
-        out = latent.copy()
-        out["samples"] = samples
-        if "x0" in x0_output:
-            out_denoised = latent.copy()
-            out_denoised["samples"] = guider.model_patcher.model.process_latent_out(x0_output["x0"].cpu())
-        else:
-            out_denoised = out
-        return (out, out_denoised)
-
-class AddNoise:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"model": ("MODEL",),
-                     "noise": ("NOISE", ),
-                     "sigmas": ("SIGMAS", ),
-                     "latent_image": ("LATENT", ),
-                     }
-                }
-
-    RETURN_TYPES = ("LATENT",)
-
-    FUNCTION = "add_noise"
-
-    CATEGORY = "_for_testing/custom_sampling/noise"
-
-    def add_noise(self, model, noise, sigmas, latent_image):
-        if len(sigmas) == 0:
-            return latent_image
-
-        latent = latent_image
-        latent_image = latent["samples"]
-
-        noisy = noise.generate_noise(latent)
-
-        model_sampling = model.get_model_object("model_sampling")
-        process_latent_out = model.get_model_object("process_latent_out")
-        process_latent_in = model.get_model_object("process_latent_in")
-
-        if len(sigmas) > 1:
-            scale = torch.abs(sigmas[0] - sigmas[-1])
-        else:
-            scale = sigmas[0]
-
-        if torch.count_nonzero(latent_image) > 0: #Don't shift the empty latent image.
-            latent_image = process_latent_in(latent_image)
-        noisy = model_sampling.noise_scaling(scale, noisy, latent_image)
-        noisy = process_latent_out(noisy)
-        noisy = torch.nan_to_num(noisy, nan=0.0, posinf=0.0, neginf=0.0)
-
-        out = latent.copy()
-        out["samples"] = noisy
-        return (out,)
-
-
 NODE_CLASS_MAPPINGS = {
     "SamplerCustom": SamplerCustom,
     "BasicScheduler": BasicScheduler,
@@ -614,21 +308,9 @@ NODE_CLASS_MAPPINGS = {
     "VPScheduler": VPScheduler,
     "SDTurboScheduler": SDTurboScheduler,
     "KSamplerSelect": KSamplerSelect,
-    "SamplerEulerAncestral": SamplerEulerAncestral,
-    "SamplerLMS": SamplerLMS,
-    "SamplerDPMPP_3M_SDE": SamplerDPMPP_3M_SDE,
     "SamplerDPMPP_2M_SDE": SamplerDPMPP_2M_SDE,
     "SamplerDPMPP_SDE": SamplerDPMPP_SDE,
-    "SamplerDPMAdaptative": SamplerDPMAdaptative,
+    "SamplerTCD": SamplerTCD,
     "SplitSigmas": SplitSigmas,
-    "SplitSigmasDenoise": SplitSigmasDenoise,
     "FlipSigmas": FlipSigmas,
-
-    "CFGGuider": CFGGuider,
-    "DualCFGGuider": DualCFGGuider,
-    "BasicGuider": BasicGuider,
-    "RandomNoise": RandomNoise,
-    "DisableNoise": DisableNoise,
-    "AddNoise": AddNoise,
-    "SamplerCustomAdvanced": SamplerCustomAdvanced,
 }

ldm_patched/contrib/external_differential_diffusion.py

@@ -1,42 +0,0 @@
-# code adapted from https://github.com/exx8/differential-diffusion
-
-import torch
-
-class DifferentialDiffusion():
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"model": ("MODEL", ),
-                            }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "apply"
-    CATEGORY = "_for_testing"
-    INIT = False
-
-    def apply(self, model):
-        model = model.clone()
-        model.set_model_denoise_mask_function(self.forward)
-        return (model,)
-
-    def forward(self, sigma: torch.Tensor, denoise_mask: torch.Tensor, extra_options: dict):
-        model = extra_options["model"]
-        step_sigmas = extra_options["sigmas"]
-        sigma_to = model.inner_model.model_sampling.sigma_min
-        if step_sigmas[-1] > sigma_to:
-            sigma_to = step_sigmas[-1]
-        sigma_from = step_sigmas[0]
-
-        ts_from = model.inner_model.model_sampling.timestep(sigma_from)
-        ts_to = model.inner_model.model_sampling.timestep(sigma_to)
-        current_ts = model.inner_model.model_sampling.timestep(sigma[0])
-
-        threshold = (current_ts - ts_to) / (ts_from - ts_to)
-
-        return (denoise_mask >= threshold).to(denoise_mask.dtype)
-
-
-NODE_CLASS_MAPPINGS = {
-    "DifferentialDiffusion": DifferentialDiffusion,
-}
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "DifferentialDiffusion": "Differential Diffusion",
-}

ldm_patched/contrib/external_freelunch.py

@@ -1,7 +1,9 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 #code originally taken from: https://github.com/ChenyangSi/FreeU (under MIT License)
 
 import torch
-import logging
+
 
 def Fourier_filter(x, threshold, scale):
     # FFT
@@ -42,14 +44,14 @@ class FreeU:
         on_cpu_devices = {}
 
         def output_block_patch(h, hsp, transformer_options):
-            scale = scale_dict.get(int(h.shape[1]), None)
+            scale = scale_dict.get(h.shape[1], None)
             if scale is not None:
                 h[:,:h.shape[1] // 2] = h[:,:h.shape[1] // 2] * scale[0]
                 if hsp.device not in on_cpu_devices:
                     try:
                         hsp = Fourier_filter(hsp, threshold=1, scale=scale[1])
                     except:
-                        logging.warning("Device {} does not support the torch.fft functions used in the FreeU node, switching to CPU.".format(hsp.device))
+                        print("Device", hsp.device, "does not support the torch.fft functions used in the FreeU node, switching to CPU.")
                         on_cpu_devices[hsp.device] = True
                         hsp = Fourier_filter(hsp.cpu(), threshold=1, scale=scale[1]).to(hsp.device)
                 else:
@@ -81,7 +83,7 @@ class FreeU_V2:
         on_cpu_devices = {}
 
         def output_block_patch(h, hsp, transformer_options):
-            scale = scale_dict.get(int(h.shape[1]), None)
+            scale = scale_dict.get(h.shape[1], None)
             if scale is not None:
                 hidden_mean = h.mean(1).unsqueeze(1)
                 B = hidden_mean.shape[0]
@@ -95,7 +97,7 @@ class FreeU_V2:
                     try:
                         hsp = Fourier_filter(hsp, threshold=1, scale=scale[1])
                     except:
-                        logging.warning("Device {} does not support the torch.fft functions used in the FreeU node, switching to CPU.".format(hsp.device))
+                        print("Device", hsp.device, "does not support the torch.fft functions used in the FreeU node, switching to CPU.")
                         on_cpu_devices[hsp.device] = True
                         hsp = Fourier_filter(hsp.cpu(), threshold=1, scale=scale[1]).to(hsp.device)
                 else:

ldm_patched/contrib/external_hypernetwork.py

@@ -1,7 +1,8 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import ldm_patched.modules.utils
-import ldm_patched.utils.path_utils as folder_paths
+import ldm_patched.utils.path_utils
 import torch
-import logging
 
 def load_hypernetwork_patch(path, strength):
     sd = ldm_patched.modules.utils.load_torch_file(path, safe_load=True)
@@ -24,7 +25,7 @@ def load_hypernetwork_patch(path, strength):
     }
 
     if activation_func not in valid_activation:
-        logging.error("Unsupported Hypernetwork format, if you report it I might implement it. {}   {} {} {} {} {}".format(path, activation_func, is_layer_norm, use_dropout, activate_output, last_layer_dropout))
+        print("Unsupported Hypernetwork format, if you report it I might implement it.", path, " ", activation_func, is_layer_norm, use_dropout, activate_output, last_layer_dropout)
         return None
 
     out = {}
@@ -98,7 +99,7 @@ class HypernetworkLoader:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": { "model": ("MODEL",),
-                              "hypernetwork_name": (folder_paths.get_filename_list("hypernetworks"), ),
+                              "hypernetwork_name": (ldm_patched.utils.path_utils.get_filename_list("hypernetworks"), ),
                               "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
                               }}
     RETURN_TYPES = ("MODEL",)
@@ -107,7 +108,7 @@ class HypernetworkLoader:
     CATEGORY = "loaders"
 
     def load_hypernetwork(self, model, hypernetwork_name, strength):
-        hypernetwork_path = folder_paths.get_full_path("hypernetworks", hypernetwork_name)
+        hypernetwork_path = ldm_patched.utils.path_utils.get_full_path("hypernetworks", hypernetwork_name)
         model_hypernetwork = model.clone()
         patch = load_hypernetwork_patch(hypernetwork_path, strength)
         if patch is not None:

ldm_patched/contrib/external_hypertile.py

@@ -1,3 +1,5 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 #Taken from: https://github.com/tfernd/HyperTile/
 
 import math

ldm_patched/contrib/external_images.py

@@ -1,5 +1,7 @@
-import external
-import ldm_patched.utils.path_utils as folder_paths
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
+import ldm_patched.contrib.external
+import ldm_patched.utils.path_utils
 from ldm_patched.modules.args_parser import args
 
 from PIL import Image
@@ -9,7 +11,7 @@ import numpy as np
 import json
 import os
 
-MAX_RESOLUTION = external.MAX_RESOLUTION
+MAX_RESOLUTION = ldm_patched.contrib.external.MAX_RESOLUTION
 
 class ImageCrop:
     @classmethod
@@ -37,7 +39,7 @@ class RepeatImageBatch:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": { "image": ("IMAGE",),
-                              "amount": ("INT", {"default": 1, "min": 1, "max": 4096}),
+                              "amount": ("INT", {"default": 1, "min": 1, "max": 64}),
                               }}
     RETURN_TYPES = ("IMAGE",)
     FUNCTION = "repeat"
@@ -48,28 +50,9 @@ class RepeatImageBatch:
         s = image.repeat((amount, 1,1,1))
         return (s,)
 
-class ImageFromBatch:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "image": ("IMAGE",),
-                              "batch_index": ("INT", {"default": 0, "min": 0, "max": 4095}),
-                              "length": ("INT", {"default": 1, "min": 1, "max": 4096}),
-                              }}
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "frombatch"
-
-    CATEGORY = "image/batch"
-
-    def frombatch(self, image, batch_index, length):
-        s_in = image
-        batch_index = min(s_in.shape[0] - 1, batch_index)
-        length = min(s_in.shape[0] - batch_index, length)
-        s = s_in[batch_index:batch_index + length].clone()
-        return (s,)
-
 class SaveAnimatedWEBP:
     def __init__(self):
-        self.output_dir = folder_paths.get_output_directory()
+        self.output_dir = ldm_patched.utils.path_utils.get_output_directory()
         self.type = "output"
         self.prefix_append = ""
 
@@ -78,7 +61,7 @@ class SaveAnimatedWEBP:
     def INPUT_TYPES(s):
         return {"required":
                     {"images": ("IMAGE", ),
-                     "filename_prefix": ("STRING", {"default": "ComfyUI"}),
+                     "filename_prefix": ("STRING", {"default": "ldm_patched"}),
                      "fps": ("FLOAT", {"default": 6.0, "min": 0.01, "max": 1000.0, "step": 0.01}),
                      "lossless": ("BOOLEAN", {"default": True}),
                      "quality": ("INT", {"default": 80, "min": 0, "max": 100}),
@@ -98,7 +81,7 @@ class SaveAnimatedWEBP:
     def save_images(self, images, fps, filename_prefix, lossless, quality, method, num_frames=0, prompt=None, extra_pnginfo=None):
         method = self.methods.get(method)
         filename_prefix += self.prefix_append
-        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
+        full_output_folder, filename, counter, subfolder, filename_prefix = ldm_patched.utils.path_utils.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
         results = list()
         pil_images = []
         for image in images:
@@ -107,7 +90,7 @@ class SaveAnimatedWEBP:
             pil_images.append(img)
 
         metadata = pil_images[0].getexif()
-        if not args.disable_metadata:
+        if not args.disable_server_info:
             if prompt is not None:
                 metadata[0x0110] = "prompt:{}".format(json.dumps(prompt))
             if extra_pnginfo is not None:
@@ -135,7 +118,7 @@ class SaveAnimatedWEBP:
 
 class SaveAnimatedPNG:
     def __init__(self):
-        self.output_dir = folder_paths.get_output_directory()
+        self.output_dir = ldm_patched.utils.path_utils.get_output_directory()
         self.type = "output"
         self.prefix_append = ""
 
@@ -143,7 +126,7 @@ class SaveAnimatedPNG:
     def INPUT_TYPES(s):
         return {"required":
                     {"images": ("IMAGE", ),
-                     "filename_prefix": ("STRING", {"default": "ComfyUI"}),
+                     "filename_prefix": ("STRING", {"default": "ldm_patched"}),
                      "fps": ("FLOAT", {"default": 6.0, "min": 0.01, "max": 1000.0, "step": 0.01}),
                      "compress_level": ("INT", {"default": 4, "min": 0, "max": 9})
                      },
@@ -157,9 +140,9 @@ class SaveAnimatedPNG:
 
     CATEGORY = "image/animation"
 
-    def save_images(self, images, fps, compress_level, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
+    def save_images(self, images, fps, compress_level, filename_prefix="ldm_patched", prompt=None, extra_pnginfo=None):
         filename_prefix += self.prefix_append
-        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
+        full_output_folder, filename, counter, subfolder, filename_prefix = ldm_patched.utils.path_utils.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
         results = list()
         pil_images = []
         for image in images:
@@ -168,13 +151,13 @@ class SaveAnimatedPNG:
             pil_images.append(img)
 
         metadata = None
-        if not args.disable_metadata:
+        if not args.disable_server_info:
             metadata = PngInfo()
             if prompt is not None:
-                metadata.add(b"comf", "prompt".encode("latin-1", "strict") + b"\0" + json.dumps(prompt).encode("latin-1", "strict"), after_idat=True)
+                metadata.add(b"ldm_patched", "prompt".encode("latin-1", "strict") + b"\0" + json.dumps(prompt).encode("latin-1", "strict"), after_idat=True)
             if extra_pnginfo is not None:
                 for x in extra_pnginfo:
-                    metadata.add(b"comf", x.encode("latin-1", "strict") + b"\0" + json.dumps(extra_pnginfo[x]).encode("latin-1", "strict"), after_idat=True)
+                    metadata.add(b"ldm_patched", x.encode("latin-1", "strict") + b"\0" + json.dumps(extra_pnginfo[x]).encode("latin-1", "strict"), after_idat=True)
 
         file = f"{filename}_{counter:05}_.png"
         pil_images[0].save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=compress_level, save_all=True, duration=int(1000.0/fps), append_images=pil_images[1:])
@@ -189,7 +172,6 @@ class SaveAnimatedPNG:
 NODE_CLASS_MAPPINGS = {
     "ImageCrop": ImageCrop,
     "RepeatImageBatch": RepeatImageBatch,
-    "ImageFromBatch": ImageFromBatch,
     "SaveAnimatedWEBP": SaveAnimatedWEBP,
     "SaveAnimatedPNG": SaveAnimatedPNG,
 }

ldm_patched/contrib/external_ip2p.py

@@ -1,45 +0,0 @@
-import torch
-
-class InstructPixToPixConditioning:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"positive": ("CONDITIONING", ),
-                             "negative": ("CONDITIONING", ),
-                             "vae": ("VAE", ),
-                             "pixels": ("IMAGE", ),
-                             }}
-
-    RETURN_TYPES = ("CONDITIONING","CONDITIONING","LATENT")
-    RETURN_NAMES = ("positive", "negative", "latent")
-    FUNCTION = "encode"
-
-    CATEGORY = "conditioning/instructpix2pix"
-
-    def encode(self, positive, negative, pixels, vae):
-        x = (pixels.shape[1] // 8) * 8
-        y = (pixels.shape[2] // 8) * 8
-
-        if pixels.shape[1] != x or pixels.shape[2] != y:
-            x_offset = (pixels.shape[1] % 8) // 2
-            y_offset = (pixels.shape[2] % 8) // 2
-            pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset,:]
-
-        concat_latent = vae.encode(pixels)
-
-        out_latent = {}
-        out_latent["samples"] = torch.zeros_like(concat_latent)
-
-        out = []
-        for conditioning in [positive, negative]:
-            c = []
-            for t in conditioning:
-                d = t[1].copy()
-                d["concat_latent_image"] = concat_latent
-                n = [t[0], d]
-                c.append(n)
-            out.append(c)
-        return (out[0], out[1], out_latent)
-
-NODE_CLASS_MAPPINGS = {
-    "InstructPixToPixConditioning": InstructPixToPixConditioning,
-}

ldm_patched/contrib/external_latent.py

@@ -1,3 +1,5 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import ldm_patched.modules.utils
 import torch
 
@@ -126,7 +128,7 @@ class LatentBatchSeedBehavior:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": { "samples": ("LATENT",),
-                              "seed_behavior": (["random", "fixed"],{"default": "fixed"}),}}
+                              "seed_behavior": (["random", "fixed"],),}}
 
     RETURN_TYPES = ("LATENT",)
     FUNCTION = "op"

ldm_patched/contrib/external_mask.py

@@ -1,9 +1,11 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import numpy as np
 import scipy.ndimage
 import torch
 import ldm_patched.modules.utils
 
-from external import MAX_RESOLUTION
+from ldm_patched.contrib.external import MAX_RESOLUTION
 
 def composite(destination, source, x, y, mask = None, multiplier = 8, resize_source = False):
     source = source.to(destination.device)
@@ -341,24 +343,6 @@ class GrowMask:
             out.append(output)
         return (torch.stack(out, dim=0),)
 
-class ThresholdMask:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-                "required": {
-                    "mask": ("MASK",),
-                    "value": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
-                }
-        }
-
-    CATEGORY = "mask"
-
-    RETURN_TYPES = ("MASK",)
-    FUNCTION = "image_to_mask"
-
-    def image_to_mask(self, mask, value):
-        mask = (mask > value).float()
-        return (mask,)
 
 
 NODE_CLASS_MAPPINGS = {
@@ -373,7 +357,6 @@ NODE_CLASS_MAPPINGS = {
     "MaskComposite": MaskComposite,
     "FeatherMask": FeatherMask,
     "GrowMask": GrowMask,
-    "ThresholdMask": ThresholdMask,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {

ldm_patched/contrib/external_model_advanced.py

@@ -1,7 +1,8 @@
-import ldm_patched.utils.path_utils as folder_paths
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
+import ldm_patched.utils.path_utils
 import ldm_patched.modules.sd
 import ldm_patched.modules.model_sampling
-import ldm_patched.modules.latent_formats
 import torch
 
 class LCM(ldm_patched.modules.model_sampling.EPS):
@@ -18,10 +19,6 @@ class LCM(ldm_patched.modules.model_sampling.EPS):
 
         return c_out * x0 + c_skip * model_input
 
-class X0(ldm_patched.modules.model_sampling.EPS):
-    def calculate_denoised(self, sigma, model_output, model_input):
-        return model_output
-
 class ModelSamplingDiscreteDistilled(ldm_patched.modules.model_sampling.ModelSamplingDiscrete):
     original_timesteps = 50
 
@@ -73,7 +70,7 @@ class ModelSamplingDiscrete:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": { "model": ("MODEL",),
-                              "sampling": (["eps", "v_prediction", "lcm", "x0"],),
+                              "sampling": (["eps", "v_prediction", "lcm", "tcd"]),
                               "zsnr": ("BOOLEAN", {"default": False}),
                               }}
 
@@ -93,8 +90,9 @@ class ModelSamplingDiscrete:
         elif sampling == "lcm":
             sampling_type = LCM
             sampling_base = ModelSamplingDiscreteDistilled
-        elif sampling == "x0":
-            sampling_type = X0
+        elif sampling == "tcd":
+            sampling_type = ldm_patched.modules.model_sampling.EPS
+            sampling_base = ModelSamplingDiscreteDistilled
 
         class ModelSamplingAdvanced(sampling_base, sampling_type):
             pass
@@ -106,32 +104,6 @@ class ModelSamplingDiscrete:
         m.add_object_patch("model_sampling", model_sampling)
         return (m, )
 
-class ModelSamplingStableCascade:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "shift": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 100.0, "step":0.01}),
-                              }}
-
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "advanced/model"
-
-    def patch(self, model, shift):
-        m = model.clone()
-
-        sampling_base = ldm_patched.modules.model_sampling.StableCascadeSampling
-        sampling_type = ldm_patched.modules.model_sampling.EPS
-
-        class ModelSamplingAdvanced(sampling_base, sampling_type):
-            pass
-
-        model_sampling = ModelSamplingAdvanced(model.model.model_config)
-        model_sampling.set_parameters(shift)
-        m.add_object_patch("model_sampling", model_sampling)
-        return (m, )
-
 class ModelSamplingContinuousEDM:
     @classmethod
     def INPUT_TYPES(s):
@@ -212,6 +184,5 @@ class RescaleCFG:
 NODE_CLASS_MAPPINGS = {
     "ModelSamplingDiscrete": ModelSamplingDiscrete,
     "ModelSamplingContinuousEDM": ModelSamplingContinuousEDM,
-    "ModelSamplingStableCascade": ModelSamplingStableCascade,
     "RescaleCFG": RescaleCFG,
 }

ldm_patched/contrib/external_model_downscale.py

@@ -1,3 +1,5 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import torch
 import ldm_patched.modules.utils
 
@@ -20,9 +22,8 @@ class PatchModelAddDownscale:
     CATEGORY = "_for_testing"
 
     def patch(self, model, block_number, downscale_factor, start_percent, end_percent, downscale_after_skip, downscale_method, upscale_method):
-        model_sampling = model.get_model_object("model_sampling")
-        sigma_start = model_sampling.percent_to_sigma(start_percent)
-        sigma_end = model_sampling.percent_to_sigma(end_percent)
+        sigma_start = model.model.model_sampling.percent_to_sigma(start_percent)
+        sigma_end = model.model.model_sampling.percent_to_sigma(end_percent)
 
         def input_block_patch(h, transformer_options):
             if transformer_options["block"][1] == block_number:

ldm_patched/contrib/external_model_merging.py

@@ -1,12 +1,11 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import ldm_patched.modules.sd
 import ldm_patched.modules.utils
 import ldm_patched.modules.model_base
 import ldm_patched.modules.model_management
-import ldm_patched.modules.model_sampling
-
-import torch
-import ldm_patched.utils.path_utils as folder_paths
 
+import ldm_patched.utils.path_utils
 import json
 import os
 
@@ -90,50 +89,6 @@ class CLIPMergeSimple:
             m.add_patches({k: kp[k]}, 1.0 - ratio, ratio)
         return (m, )
 
-
-class CLIPSubtract:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "clip1": ("CLIP",),
-                              "clip2": ("CLIP",),
-                              "multiplier": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
-                              }}
-    RETURN_TYPES = ("CLIP",)
-    FUNCTION = "merge"
-
-    CATEGORY = "advanced/model_merging"
-
-    def merge(self, clip1, clip2, multiplier):
-        m = clip1.clone()
-        kp = clip2.get_key_patches()
-        for k in kp:
-            if k.endswith(".position_ids") or k.endswith(".logit_scale"):
-                continue
-            m.add_patches({k: kp[k]}, - multiplier, multiplier)
-        return (m, )
-
-
-class CLIPAdd:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "clip1": ("CLIP",),
-                              "clip2": ("CLIP",),
-                              }}
-    RETURN_TYPES = ("CLIP",)
-    FUNCTION = "merge"
-
-    CATEGORY = "advanced/model_merging"
-
-    def merge(self, clip1, clip2):
-        m = clip1.clone()
-        kp = clip2.get_key_patches()
-        for k in kp:
-            if k.endswith(".position_ids") or k.endswith(".logit_scale"):
-                continue
-            m.add_patches({k: kp[k]}, 1.0, 1.0)
-        return (m, )
-
-
 class ModelMergeBlocks:
     @classmethod
     def INPUT_TYPES(s):
@@ -167,7 +122,7 @@ class ModelMergeBlocks:
         return (m, )
 
 def save_checkpoint(model, clip=None, vae=None, clip_vision=None, filename_prefix=None, output_dir=None, prompt=None, extra_pnginfo=None):
-    full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, output_dir)
+    full_output_folder, filename, counter, subfolder, filename_prefix = ldm_patched.utils.path_utils.get_save_image_path(filename_prefix, output_dir)
     prompt_info = ""
     if prompt is not None:
         prompt_info = json.dumps(prompt)
@@ -176,14 +131,9 @@ def save_checkpoint(model, clip=None, vae=None, clip_vision=None, filename_prefi
 
     enable_modelspec = True
     if isinstance(model.model, ldm_patched.modules.model_base.SDXL):
-        if isinstance(model.model, ldm_patched.modules.model_base.SDXL_instructpix2pix):
-            metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-edit"
-        else:
-            metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-base"
+        metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-base"
     elif isinstance(model.model, ldm_patched.modules.model_base.SDXLRefiner):
         metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-refiner"
-    elif isinstance(model.model, ldm_patched.modules.model_base.SVD_img2vid):
-        metadata["modelspec.architecture"] = "stable-video-diffusion-img2vid-v1"
     else:
         enable_modelspec = False
 
@@ -197,19 +147,12 @@ def save_checkpoint(model, clip=None, vae=None, clip_vision=None, filename_prefi
     # "stable-diffusion-v2-768-v", "stable-diffusion-v2-unclip-l", "stable-diffusion-v2-unclip-h",
     # "v2-inpainting"
 
-    extra_keys = {}
-    model_sampling = model.get_model_object("model_sampling")
-    if isinstance(model_sampling, ldm_patched.modules.model_sampling.ModelSamplingContinuousEDM):
-        if isinstance(model_sampling, ldm_patched.modules.model_sampling.V_PREDICTION):
-            extra_keys["edm_vpred.sigma_max"] = torch.tensor(model_sampling.sigma_max).float()
-            extra_keys["edm_vpred.sigma_min"] = torch.tensor(model_sampling.sigma_min).float()
-
     if model.model.model_type == ldm_patched.modules.model_base.ModelType.EPS:
         metadata["modelspec.predict_key"] = "epsilon"
     elif model.model.model_type == ldm_patched.modules.model_base.ModelType.V_PREDICTION:
         metadata["modelspec.predict_key"] = "v"
 
-    if not args.disable_metadata:
+    if not args.disable_server_info:
         metadata["prompt"] = prompt_info
         if extra_pnginfo is not None:
             for x in extra_pnginfo:
@@ -218,18 +161,18 @@ def save_checkpoint(model, clip=None, vae=None, clip_vision=None, filename_prefi
     output_checkpoint = f"{filename}_{counter:05}_.safetensors"
     output_checkpoint = os.path.join(full_output_folder, output_checkpoint)
 
-    ldm_patched.modules.sd.save_checkpoint(output_checkpoint, model, clip, vae, clip_vision, metadata=metadata, extra_keys=extra_keys)
+    ldm_patched.modules.sd.save_checkpoint(output_checkpoint, model, clip, vae, clip_vision, metadata=metadata)
 
 class CheckpointSave:
     def __init__(self):
-        self.output_dir = folder_paths.get_output_directory()
+        self.output_dir = ldm_patched.utils.path_utils.get_output_directory()
 
     @classmethod
     def INPUT_TYPES(s):
         return {"required": { "model": ("MODEL",),
                               "clip": ("CLIP",),
                               "vae": ("VAE",),
-                              "filename_prefix": ("STRING", {"default": "checkpoints/ComfyUI"}),},
+                              "filename_prefix": ("STRING", {"default": "checkpoints/ldm_patched"}),},
                 "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},}
     RETURN_TYPES = ()
     FUNCTION = "save"
@@ -243,12 +186,12 @@ class CheckpointSave:
 
 class CLIPSave:
     def __init__(self):
-        self.output_dir = folder_paths.get_output_directory()
+        self.output_dir = ldm_patched.utils.path_utils.get_output_directory()
 
     @classmethod
     def INPUT_TYPES(s):
         return {"required": { "clip": ("CLIP",),
-                              "filename_prefix": ("STRING", {"default": "clip/ComfyUI"}),},
+                              "filename_prefix": ("STRING", {"default": "clip/ldm_patched"}),},
                 "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},}
     RETURN_TYPES = ()
     FUNCTION = "save"
@@ -262,13 +205,13 @@ class CLIPSave:
             prompt_info = json.dumps(prompt)
 
         metadata = {}
-        if not args.disable_metadata:
+        if not args.disable_server_info:
             metadata["prompt"] = prompt_info
             if extra_pnginfo is not None:
                 for x in extra_pnginfo:
                     metadata[x] = json.dumps(extra_pnginfo[x])
 
-        ldm_patched.modules.model_management.load_models_gpu([clip.load_model()], force_patch_weights=True)
+        ldm_patched.modules.model_management.load_models_gpu([clip.load_model()])
         clip_sd = clip.get_sd()
 
         for prefix in ["clip_l.", "clip_g.", ""]:
@@ -287,7 +230,7 @@ class CLIPSave:
                 replace_prefix[prefix] = ""
             replace_prefix["transformer."] = ""
 
-            full_output_folder, filename, counter, subfolder, filename_prefix_ = folder_paths.get_save_image_path(filename_prefix_, self.output_dir)
+            full_output_folder, filename, counter, subfolder, filename_prefix_ = ldm_patched.utils.path_utils.get_save_image_path(filename_prefix_, self.output_dir)
 
             output_checkpoint = f"{filename}_{counter:05}_.safetensors"
             output_checkpoint = os.path.join(full_output_folder, output_checkpoint)
@@ -299,12 +242,12 @@ class CLIPSave:
 
 class VAESave:
     def __init__(self):
-        self.output_dir = folder_paths.get_output_directory()
+        self.output_dir = ldm_patched.utils.path_utils.get_output_directory()
 
     @classmethod
     def INPUT_TYPES(s):
         return {"required": { "vae": ("VAE",),
-                              "filename_prefix": ("STRING", {"default": "vae/ComfyUI_vae"}),},
+                              "filename_prefix": ("STRING", {"default": "vae/ldm_patched_vae"}),},
                 "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},}
     RETURN_TYPES = ()
     FUNCTION = "save"
@@ -313,13 +256,13 @@ class VAESave:
     CATEGORY = "advanced/model_merging"
 
     def save(self, vae, filename_prefix, prompt=None, extra_pnginfo=None):
-        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
+        full_output_folder, filename, counter, subfolder, filename_prefix = ldm_patched.utils.path_utils.get_save_image_path(filename_prefix, self.output_dir)
         prompt_info = ""
         if prompt is not None:
             prompt_info = json.dumps(prompt)
 
         metadata = {}
-        if not args.disable_metadata:
+        if not args.disable_server_info:
             metadata["prompt"] = prompt_info
             if extra_pnginfo is not None:
                 for x in extra_pnginfo:
@@ -338,8 +281,6 @@ NODE_CLASS_MAPPINGS = {
     "ModelMergeAdd": ModelAdd,
     "CheckpointSave": CheckpointSave,
     "CLIPMergeSimple": CLIPMergeSimple,
-    "CLIPMergeSubtract": CLIPSubtract,
-    "CLIPMergeAdd": CLIPAdd,
     "CLIPSave": CLIPSave,
     "VAESave": VAESave,
 }

ldm_patched/contrib/external_model_merging_model_specific.py

@@ -1,60 +0,0 @@
-import external_model_merging
-
-class ModelMergeSD1(external_model_merging.ModelMergeBlocks):
-    CATEGORY = "advanced/model_merging/model_specific"
-    @classmethod
-    def INPUT_TYPES(s):
-        arg_dict = { "model1": ("MODEL",),
-                              "model2": ("MODEL",)}
-
-        argument = ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01})
-
-        arg_dict["time_embed."] = argument
-        arg_dict["label_emb."] = argument
-
-        for i in range(12):
-            arg_dict["input_blocks.{}.".format(i)] = argument
-
-        for i in range(3):
-            arg_dict["middle_block.{}.".format(i)] = argument
-
-        for i in range(12):
-            arg_dict["output_blocks.{}.".format(i)] = argument
-
-        arg_dict["out."] = argument
-
-        return {"required": arg_dict}
-
-
-class ModelMergeSDXL(external_model_merging.ModelMergeBlocks):
-    CATEGORY = "advanced/model_merging/model_specific"
-
-    @classmethod
-    def INPUT_TYPES(s):
-        arg_dict = { "model1": ("MODEL",),
-                              "model2": ("MODEL",)}
-
-        argument = ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01})
-
-        arg_dict["time_embed."] = argument
-        arg_dict["label_emb."] = argument
-
-        for i in range(9):
-            arg_dict["input_blocks.{}".format(i)] = argument
-
-        for i in range(3):
-            arg_dict["middle_block.{}".format(i)] = argument
-
-        for i in range(9):
-            arg_dict["output_blocks.{}".format(i)] = argument
-
-        arg_dict["out."] = argument
-
-        return {"required": arg_dict}
-
-
-NODE_CLASS_MAPPINGS = {
-    "ModelMergeSD1": ModelMergeSD1,
-    "ModelMergeSD2": ModelMergeSD1, #SD1 and SD2 have the same blocks
-    "ModelMergeSDXL": ModelMergeSDXL,
-}

ldm_patched/contrib/external_morphology.py

@@ -1,49 +0,0 @@
-import torch
-import ldm_patched.modules.model_management
-
-from kornia.morphology import dilation, erosion, opening, closing, gradient, top_hat, bottom_hat
-
-
-class Morphology:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"image": ("IMAGE",),
-                                "operation": (["erode",  "dilate", "open", "close", "gradient", "bottom_hat", "top_hat"],),
-                                "kernel_size": ("INT", {"default": 3, "min": 3, "max": 999, "step": 1}),
-                                }}
-
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "process"
-
-    CATEGORY = "image/postprocessing"
-
-    def process(self, image, operation, kernel_size):
-        device = ldm_patched.modules.model_management.get_torch_device()
-        kernel = torch.ones(kernel_size, kernel_size, device=device)
-        image_k = image.to(device).movedim(-1, 1)
-        if operation == "erode":
-            output = erosion(image_k, kernel)
-        elif operation == "dilate":
-            output = dilation(image_k, kernel)
-        elif operation == "open":
-            output = opening(image_k, kernel)
-        elif operation == "close":
-            output = closing(image_k, kernel)
-        elif operation == "gradient":
-            output = gradient(image_k, kernel)
-        elif operation == "top_hat":
-            output = top_hat(image_k, kernel)
-        elif operation == "bottom_hat":
-            output = bottom_hat(image_k, kernel)
-        else:
-            raise ValueError(f"Invalid operation {operation} for morphology. Must be one of 'erode', 'dilate', 'open', 'close', 'gradient', 'tophat', 'bottomhat'")
-        img_out = output.to(ldm_patched.modules.model_management.intermediate_device()).movedim(1, -1)
-        return (img_out,)
-
-NODE_CLASS_MAPPINGS = {
-    "Morphology": Morphology,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "Morphology": "ImageMorphology",
-}

ldm_patched/contrib/external_pag.py

@@ -1,56 +0,0 @@
-#Modified/simplified version of the node from: https://github.com/pamparamm/sd-perturbed-attention
-#If you want the one with more options see the above repo.
-
-#My modified one here is more basic but has less chances of breaking with ComfyUI updates.
-
-import ldm_patched.modules.model_patcher
-import ldm_patched.modules.samplers
-
-class PerturbedAttentionGuidance:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "model": ("MODEL",),
-                "scale": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
-            }
-        }
-
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "_for_testing"
-
-    def patch(self, model, scale):
-        unet_block = "middle"
-        unet_block_id = 0
-        m = model.clone()
-
-        def perturbed_attention(q, k, v, extra_options, mask=None):
-            return v
-
-        def post_cfg_function(args):
-            model = args["model"]
-            cond_pred = args["cond_denoised"]
-            cond = args["cond"]
-            cfg_result = args["denoised"]
-            sigma = args["sigma"]
-            model_options = args["model_options"].copy()
-            x = args["input"]
-
-            if scale == 0:
-                return cfg_result
-
-            # Replace Self-attention with PAG
-            model_options = ldm_patched.modules.model_patcher.set_model_options_patch_replace(model_options, perturbed_attention, "attn1", unet_block, unet_block_id)
-            (pag,) = ldm_patched.modules.samplers.calc_cond_batch(model, [cond], x, sigma, model_options)
-
-            return cfg_result + (cond_pred - pag) * scale
-
-        m.set_model_sampler_post_cfg_function(post_cfg_function)
-
-        return (m,)
-
-NODE_CLASS_MAPPINGS = {
-    "PerturbedAttentionGuidance": PerturbedAttentionGuidance,
-}

ldm_patched/contrib/external_perpneg.py

@@ -1,26 +1,18 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import torch
 import ldm_patched.modules.model_management
-import ldm_patched.modules.sampler_helpers
+import ldm_patched.modules.sample
 import ldm_patched.modules.samplers
 import ldm_patched.modules.utils
-import ldm_patched.utils.node_helpers as node_helpers
 
-def perp_neg(x, noise_pred_pos, noise_pred_neg, noise_pred_nocond, neg_scale, cond_scale):
-    pos = noise_pred_pos - noise_pred_nocond
-    neg = noise_pred_neg - noise_pred_nocond
 
-    perp = neg - ((torch.mul(neg, pos).sum())/(torch.norm(pos)**2)) * pos
-    perp_neg = perp * neg_scale
-    cfg_result = noise_pred_nocond + cond_scale*(pos - perp_neg)
-    return cfg_result
-
-#TODO: This node should be removed, it has been replaced with PerpNegGuider
 class PerpNeg:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {"model": ("MODEL", ),
                              "empty_conditioning": ("CONDITIONING", ),
-                             "neg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
+                             "neg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0}),
                             }}
     RETURN_TYPES = ("MODEL",)
     FUNCTION = "patch"
@@ -29,7 +21,7 @@ class PerpNeg:
 
     def patch(self, model, empty_conditioning, neg_scale):
         m = model.clone()
-        nocond = ldm_patched.modules.sampler_helpers.convert_cond(empty_conditioning)
+        nocond = ldm_patched.modules.sample.convert_cond(empty_conditioning)
 
         def cfg_function(args):
             model = args["model"]
@@ -41,9 +33,14 @@ class PerpNeg:
             model_options = args["model_options"]
             nocond_processed = ldm_patched.modules.samplers.encode_model_conds(model.extra_conds, nocond, x, x.device, "negative")
 
-            (noise_pred_nocond,) = ldm_patched.modules.samplers.calc_cond_batch(model, [nocond_processed], x, sigma, model_options)
+            (noise_pred_nocond, _) = ldm_patched.modules.samplers.calc_cond_uncond_batch(model, nocond_processed, None, x, sigma, model_options)
 
-            cfg_result = x - perp_neg(x, noise_pred_pos, noise_pred_neg, noise_pred_nocond, neg_scale, cond_scale)
+            pos = noise_pred_pos - noise_pred_nocond
+            neg = noise_pred_neg - noise_pred_nocond
+            perp = ((torch.mul(pos, neg).sum())/(torch.norm(neg)**2)) * neg
+            perp_neg = perp * neg_scale
+            cfg_result = noise_pred_nocond + cond_scale*(pos - perp_neg)
+            cfg_result = x - cfg_result
             return cfg_result
 
         m.set_model_sampler_cfg_function(cfg_function)
@@ -51,78 +48,10 @@ class PerpNeg:
         return (m, )
 
 
-class Guider_PerpNeg(ldm_patched.modules.samplers.CFGGuider):
-    def set_conds(self, positive, negative, empty_negative_prompt):
-        empty_negative_prompt = node_helpers.conditioning_set_values(empty_negative_prompt, {"prompt_type": "negative"})
-        self.inner_set_conds({"positive": positive, "empty_negative_prompt": empty_negative_prompt, "negative": negative})
-
-    def set_cfg(self, cfg, neg_scale):
-        self.cfg = cfg
-        self.neg_scale = neg_scale
-
-    def predict_noise(self, x, timestep, model_options={}, seed=None):
-        # in CFGGuider.predict_noise, we call sampling_function(), which uses cfg_function() to compute pos & neg
-        # but we'd rather do a single batch of sampling pos, neg, and empty, so we call calc_cond_batch([pos,neg,empty]) directly
-        
-        positive_cond = self.conds.get("positive", None)
-        negative_cond = self.conds.get("negative", None)
-        empty_cond = self.conds.get("empty_negative_prompt", None)
-
-        (noise_pred_pos, noise_pred_neg, noise_pred_empty) = \
-            ldm_patched.modules.samplers.calc_cond_batch(self.inner_model, [positive_cond, negative_cond, empty_cond], x, timestep, model_options)
-        cfg_result = perp_neg(x, noise_pred_pos, noise_pred_neg, noise_pred_empty, self.neg_scale, self.cfg)
-
-        # normally this would be done in cfg_function, but we skipped 
-        # that for efficiency: we can compute the noise predictions in
-        # a single call to calc_cond_batch() (rather than two)
-        # so we replicate the hook here
-        for fn in model_options.get("sampler_post_cfg_function", []):
-            args = {
-                "denoised": cfg_result,
-                "cond": positive_cond,
-                "uncond": negative_cond,
-                "model": self.inner_model,
-                "uncond_denoised": noise_pred_neg,
-                "cond_denoised": noise_pred_pos,
-                "sigma": timestep,
-                "model_options": model_options,
-                "input": x,
-                # not in the original call in samplers.py:cfg_function, but made available for future hooks
-                "empty_cond": empty_cond,
-                "empty_cond_denoised": noise_pred_empty,}
-            cfg_result = fn(args)
-
-        return cfg_result
-
-class PerpNegGuider:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"model": ("MODEL",),
-                    "positive": ("CONDITIONING", ),
-                    "negative": ("CONDITIONING", ),
-                    "empty_conditioning": ("CONDITIONING", ),
-                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
-                    "neg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
-                     }
-                }
-
-    RETURN_TYPES = ("GUIDER",)
-
-    FUNCTION = "get_guider"
-    CATEGORY = "_for_testing"
-
-    def get_guider(self, model, positive, negative, empty_conditioning, cfg, neg_scale):
-        guider = Guider_PerpNeg(model)
-        guider.set_conds(positive, negative, empty_conditioning)
-        guider.set_cfg(cfg, neg_scale)
-        return (guider,)
-
 NODE_CLASS_MAPPINGS = {
     "PerpNeg": PerpNeg,
-    "PerpNegGuider": PerpNegGuider,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "PerpNeg": "Perp-Neg (DEPRECATED by PerpNegGuider)",
+    "PerpNeg": "Perp-Neg",
 }

ldm_patched/contrib/external_photomaker.py

@@ -1,6 +1,8 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import torch
 import torch.nn as nn
-import ldm_patched.utils.path_utils as folder_paths
+import ldm_patched.utils.path_utils
 import ldm_patched.modules.clip_model
 import ldm_patched.modules.clip_vision
 import ldm_patched.modules.ops
@@ -118,7 +120,7 @@ class PhotoMakerIDEncoder(ldm_patched.modules.clip_model.CLIPVisionModelProjecti
 class PhotoMakerLoader:
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": { "photomaker_model_name": (folder_paths.get_filename_list("photomaker"), )}}
+        return {"required": { "photomaker_model_name": (ldm_patched.utils.path_utils.get_filename_list("photomaker"), )}}
 
     RETURN_TYPES = ("PHOTOMAKER",)
     FUNCTION = "load_photomaker_model"
@@ -126,7 +128,7 @@ class PhotoMakerLoader:
     CATEGORY = "_for_testing/photomaker"
 
     def load_photomaker_model(self, photomaker_model_name):
-        photomaker_model_path = folder_paths.get_full_path("photomaker", photomaker_model_name)
+        photomaker_model_path = ldm_patched.utils.path_utils.get_full_path("photomaker", photomaker_model_name)
         photomaker_model = PhotoMakerIDEncoder()
         data = ldm_patched.modules.utils.load_torch_file(photomaker_model_path, safe_load=True)
         if "id_encoder" in data:
@@ -141,7 +143,7 @@ class PhotoMakerEncode:
         return {"required": { "photomaker": ("PHOTOMAKER",),
                               "image": ("IMAGE",),
                               "clip": ("CLIP", ),
-                              "text": ("STRING", {"multiline": True, "dynamicPrompts": True, "default": "photograph of photomaker"}),
+                              "text": ("STRING", {"multiline": True, "default": "photograph of photomaker"}),
                              }}
 
     RETURN_TYPES = ("CONDITIONING",)

ldm_patched/contrib/external_post_processing.py

@@ -1,3 +1,5 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -5,7 +7,6 @@ from PIL import Image
 import math
 
 import ldm_patched.modules.utils
-import ldm_patched.modules.model_management
 
 
 class Blend:
@@ -103,7 +104,6 @@ class Blur:
         if blur_radius == 0:
             return (image,)
 
-        image = image.to(ldm_patched.modules.model_management.get_torch_device())
         batch_size, height, width, channels = image.shape
 
         kernel_size = blur_radius * 2 + 1
@@ -114,7 +114,7 @@ class Blur:
         blurred = F.conv2d(padded_image, kernel, padding=kernel_size // 2, groups=channels)[:,:,blur_radius:-blur_radius, blur_radius:-blur_radius]
         blurred = blurred.permute(0, 2, 3, 1)
 
-        return (blurred.to(ldm_patched.modules.model_management.intermediate_device()),)
+        return (blurred,)
 
 class Quantize:
     def __init__(self):
@@ -206,13 +206,13 @@ class Sharpen:
                     "default": 1.0,
                     "min": 0.1,
                     "max": 10.0,
-                    "step": 0.01
+                    "step": 0.1
                 }),
                 "alpha": ("FLOAT", {
                     "default": 1.0,
                     "min": 0.0,
                     "max": 5.0,
-                    "step": 0.01
+                    "step": 0.1
                 }),
             },
         }
@@ -227,7 +227,6 @@ class Sharpen:
             return (image,)
 
         batch_size, height, width, channels = image.shape
-        image = image.to(ldm_patched.modules.model_management.get_torch_device())
 
         kernel_size = sharpen_radius * 2 + 1
         kernel = gaussian_kernel(kernel_size, sigma, device=image.device) * -(alpha*10)
@@ -242,7 +241,7 @@ class Sharpen:
 
         result = torch.clamp(sharpened, 0, 1)
 
-        return (result.to(ldm_patched.modules.model_management.intermediate_device()),)
+        return (result,)
 
 class ImageScaleToTotalPixels:
     upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]

ldm_patched/contrib/external_rebatch.py

@@ -1,3 +1,5 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import torch
 
 class LatentRebatch:

ldm_patched/contrib/external_sag.py

@@ -1,15 +1,18 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import torch
 from torch import einsum
 import torch.nn.functional as F
 import math
 
 from einops import rearrange, repeat
-from ldm_patched.ldm.modules.attention import optimized_attention
+import os
+from ldm_patched.ldm.modules.attention import optimized_attention, _ATTN_PRECISION
 import ldm_patched.modules.samplers
 
-# from comfy/ldm/modules/attention.py
+# from ldm_patched.modules/ldm/modules/attention.py
 # but modified to return attention scores as well as output
-def attention_basic_with_sim(q, k, v, heads, mask=None, attn_precision=None):
+def attention_basic_with_sim(q, k, v, heads, mask=None):
     b, _, dim_head = q.shape
     dim_head //= heads
     scale = dim_head ** -0.5
@@ -25,7 +28,7 @@ def attention_basic_with_sim(q, k, v, heads, mask=None, attn_precision=None):
     )
 
     # force cast to fp32 to avoid overflowing
-    if attn_precision == torch.float32:
+    if _ATTN_PRECISION =="fp32":
         sim = einsum('b i d, b j d -> b i j', q.float(), k.float()) * scale
     else:
         sim = einsum('b i d, b j d -> b i j', q, k) * scale
@@ -120,13 +123,13 @@ class SelfAttentionGuidance:
             if 1 in cond_or_uncond:
                 uncond_index = cond_or_uncond.index(1)
                 # do the entire attention operation, but save the attention scores to attn_scores
-                (out, sim) = attention_basic_with_sim(q, k, v, heads=heads, attn_precision=extra_options["attn_precision"])
+                (out, sim) = attention_basic_with_sim(q, k, v, heads=heads)
                 # when using a higher batch size, I BELIEVE the result batch dimension is [uc1, ... ucn, c1, ... cn]
                 n_slices = heads * b
                 attn_scores = sim[n_slices * uncond_index:n_slices * (uncond_index+1)]
                 return out
             else:
-                return optimized_attention(q, k, v, heads=heads, attn_precision=extra_options["attn_precision"])
+                return optimized_attention(q, k, v, heads=heads)
 
         def post_cfg_function(args):
             nonlocal attn_scores
@@ -149,7 +152,7 @@ class SelfAttentionGuidance:
             degraded = create_blur_map(uncond_pred, uncond_attn, sag_sigma, sag_threshold)
             degraded_noised = degraded + x - uncond_pred
             # call into the UNet
-            (sag,) = ldm_patched.modules.samplers.calc_cond_batch(model, [uncond], degraded_noised, sigma, model_options)
+            (sag, _) = ldm_patched.modules.samplers.calc_cond_uncond_batch(model, uncond, None, degraded_noised, sigma, model_options)
             return cfg_result + (degraded - sag) * sag_scale
 
         m.set_model_sampler_post_cfg_function(post_cfg_function, disable_cfg1_optimization=True)

ldm_patched/contrib/external_sdupscale.py

@@ -1,4 +1,7 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import torch
+import ldm_patched.contrib.external
 import ldm_patched.modules.utils
 
 class SD_4XUpscale_Conditioning:

ldm_patched/contrib/external_stable3d.py

@@ -1,5 +1,7 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import torch
-import external
+import ldm_patched.contrib.external
 import ldm_patched.modules.utils
 
 def camera_embeddings(elevation, azimuth):
@@ -26,11 +28,11 @@ class StableZero123_Conditioning:
         return {"required": { "clip_vision": ("CLIP_VISION",),
                               "init_image": ("IMAGE",),
                               "vae": ("VAE",),
-                              "width": ("INT", {"default": 256, "min": 16, "max": external.MAX_RESOLUTION, "step": 8}),
-                              "height": ("INT", {"default": 256, "min": 16, "max": external.MAX_RESOLUTION, "step": 8}),
+                              "width": ("INT", {"default": 256, "min": 16, "max": ldm_patched.contrib.external.MAX_RESOLUTION, "step": 8}),
+                              "height": ("INT", {"default": 256, "min": 16, "max": ldm_patched.contrib.external.MAX_RESOLUTION, "step": 8}),
                               "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
-                              "elevation": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}),
-                              "azimuth": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}),
+                              "elevation": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}),
+                              "azimuth": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}),
                              }}
     RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
     RETURN_NAMES = ("positive", "negative", "latent")
@@ -59,13 +61,13 @@ class StableZero123_Conditioning_Batched:
         return {"required": { "clip_vision": ("CLIP_VISION",),
                               "init_image": ("IMAGE",),
                               "vae": ("VAE",),
-                              "width": ("INT", {"default": 256, "min": 16, "max": external.MAX_RESOLUTION, "step": 8}),
-                              "height": ("INT", {"default": 256, "min": 16, "max": external.MAX_RESOLUTION, "step": 8}),
+                              "width": ("INT", {"default": 256, "min": 16, "max": ldm_patched.contrib.external.MAX_RESOLUTION, "step": 8}),
+                              "height": ("INT", {"default": 256, "min": 16, "max": ldm_patched.contrib.external.MAX_RESOLUTION, "step": 8}),
                               "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
-                              "elevation": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}),
-                              "azimuth": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}),
-                              "elevation_batch_increment": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}),
-                              "azimuth_batch_increment": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, "round": False}),
+                              "elevation": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}),
+                              "azimuth": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}),
+                              "elevation_batch_increment": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}),
+                              "azimuth_batch_increment": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}),
                              }}
     RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
     RETURN_NAMES = ("positive", "negative", "latent")
@@ -95,49 +97,8 @@ class StableZero123_Conditioning_Batched:
         latent = torch.zeros([batch_size, 4, height // 8, width // 8])
         return (positive, negative, {"samples":latent, "batch_index": [0] * batch_size})
 
-class SV3D_Conditioning:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "clip_vision": ("CLIP_VISION",),
-                              "init_image": ("IMAGE",),
-                              "vae": ("VAE",),
-                              "width": ("INT", {"default": 576, "min": 16, "max": external.MAX_RESOLUTION, "step": 8}),
-                              "height": ("INT", {"default": 576, "min": 16, "max": external.MAX_RESOLUTION, "step": 8}),
-                              "video_frames": ("INT", {"default": 21, "min": 1, "max": 4096}),
-                              "elevation": ("FLOAT", {"default": 0.0, "min": -90.0, "max": 90.0, "step": 0.1, "round": False}),
-                             }}
-    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
-    RETURN_NAMES = ("positive", "negative", "latent")
-
-    FUNCTION = "encode"
-
-    CATEGORY = "conditioning/3d_models"
-
-    def encode(self, clip_vision, init_image, vae, width, height, video_frames, elevation):
-        output = clip_vision.encode_image(init_image)
-        pooled = output.image_embeds.unsqueeze(0)
-        pixels = ldm_patched.modules.utils.common_upscale(init_image.movedim(-1,1), width, height, "bilinear", "center").movedim(1,-1)
-        encode_pixels = pixels[:,:,:,:3]
-        t = vae.encode(encode_pixels)
-
-        azimuth = 0
-        azimuth_increment = 360 / (max(video_frames, 2) - 1)
-
-        elevations = []
-        azimuths = []
-        for i in range(video_frames):
-            elevations.append(elevation)
-            azimuths.append(azimuth)
-            azimuth += azimuth_increment
-
-        positive = [[pooled, {"concat_latent_image": t, "elevation": elevations, "azimuth": azimuths}]]
-        negative = [[torch.zeros_like(pooled), {"concat_latent_image": torch.zeros_like(t), "elevation": elevations, "azimuth": azimuths}]]
-        latent = torch.zeros([video_frames, 4, height // 8, width // 8])
-        return (positive, negative, {"samples":latent})
-
 
 NODE_CLASS_MAPPINGS = {
     "StableZero123_Conditioning": StableZero123_Conditioning,
     "StableZero123_Conditioning_Batched": StableZero123_Conditioning_Batched,
-    "SV3D_Conditioning": SV3D_Conditioning,
 }

ldm_patched/contrib/external_stable_cascade.py

@@ -1,140 +0,0 @@
-"""
-    This file is part of ComfyUI.
-    Copyright (C) 2024 Stability AI
-
-    This program is free software: you can redistribute it and/or modify
-    it under the terms of the GNU General Public License as published by
-    the Free Software Foundation, either version 3 of the License, or
-    (at your option) any later version.
-
-    This program is distributed in the hope that it will be useful,
-    but WITHOUT ANY WARRANTY; without even the implied warranty of
-    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-    GNU General Public License for more details.
-
-    You should have received a copy of the GNU General Public License
-    along with this program.  If not, see <https://www.gnu.org/licenses/>.
-"""
-
-import torch
-import external
-import ldm_patched.modules.utils
-
-
-class StableCascade_EmptyLatentImage:
-    def __init__(self, device="cpu"):
-        self.device = device
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-            "width": ("INT", {"default": 1024, "min": 256, "max": external.MAX_RESOLUTION, "step": 8}),
-            "height": ("INT", {"default": 1024, "min": 256, "max": external.MAX_RESOLUTION, "step": 8}),
-            "compression": ("INT", {"default": 42, "min": 4, "max": 128, "step": 1}),
-            "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})
-        }}
-    RETURN_TYPES = ("LATENT", "LATENT")
-    RETURN_NAMES = ("stage_c", "stage_b")
-    FUNCTION = "generate"
-
-    CATEGORY = "latent/stable_cascade"
-
-    def generate(self, width, height, compression, batch_size=1):
-        c_latent = torch.zeros([batch_size, 16, height // compression, width // compression])
-        b_latent = torch.zeros([batch_size, 4, height // 4, width // 4])
-        return ({
-            "samples": c_latent,
-        }, {
-            "samples": b_latent,
-        })
-
-class StableCascade_StageC_VAEEncode:
-    def __init__(self, device="cpu"):
-        self.device = device
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-            "image": ("IMAGE",),
-            "vae": ("VAE", ),
-            "compression": ("INT", {"default": 42, "min": 4, "max": 128, "step": 1}),
-        }}
-    RETURN_TYPES = ("LATENT", "LATENT")
-    RETURN_NAMES = ("stage_c", "stage_b")
-    FUNCTION = "generate"
-
-    CATEGORY = "latent/stable_cascade"
-
-    def generate(self, image, vae, compression):
-        width = image.shape[-2]
-        height = image.shape[-3]
-        out_width = (width // compression) * vae.downscale_ratio
-        out_height = (height // compression) * vae.downscale_ratio
-
-        s = ldm_patched.modules.utils.common_upscale(image.movedim(-1,1), out_width, out_height, "bicubic", "center").movedim(1,-1)
-
-        c_latent = vae.encode(s[:,:,:,:3])
-        b_latent = torch.zeros([c_latent.shape[0], 4, (height // 8) * 2, (width // 8) * 2])
-        return ({
-            "samples": c_latent,
-        }, {
-            "samples": b_latent,
-        })
-
-class StableCascade_StageB_Conditioning:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "conditioning": ("CONDITIONING",),
-                              "stage_c": ("LATENT",),
-                             }}
-    RETURN_TYPES = ("CONDITIONING",)
-
-    FUNCTION = "set_prior"
-
-    CATEGORY = "conditioning/stable_cascade"
-
-    def set_prior(self, conditioning, stage_c):
-        c = []
-        for t in conditioning:
-            d = t[1].copy()
-            d['stable_cascade_prior'] = stage_c['samples']
-            n = [t[0], d]
-            c.append(n)
-        return (c, )
-
-class StableCascade_SuperResolutionControlnet:
-    def __init__(self, device="cpu"):
-        self.device = device
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-            "image": ("IMAGE",),
-            "vae": ("VAE", ),
-        }}
-    RETURN_TYPES = ("IMAGE", "LATENT", "LATENT")
-    RETURN_NAMES = ("controlnet_input", "stage_c", "stage_b")
-    FUNCTION = "generate"
-
-    CATEGORY = "_for_testing/stable_cascade"
-
-    def generate(self, image, vae):
-        width = image.shape[-2]
-        height = image.shape[-3]
-        batch_size = image.shape[0]
-        controlnet_input = vae.encode(image[:,:,:,:3]).movedim(1, -1)
-
-        c_latent = torch.zeros([batch_size, 16, height // 16, width // 16])
-        b_latent = torch.zeros([batch_size, 4, height // 2, width // 2])
-        return (controlnet_input, {
-            "samples": c_latent,
-        }, {
-            "samples": b_latent,
-        })
-
-NODE_CLASS_MAPPINGS = {
-    "StableCascade_EmptyLatentImage": StableCascade_EmptyLatentImage,
-    "StableCascade_StageB_Conditioning": StableCascade_StageB_Conditioning,
-    "StableCascade_StageC_VAEEncode": StableCascade_StageC_VAEEncode,
-    "StableCascade_SuperResolutionControlnet": StableCascade_SuperResolutionControlnet,
-}

ldm_patched/contrib/external_tomesd.py

@@ -1,3 +1,5 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 #Taken from: https://github.com/dbolya/tomesd
 
 import torch

ldm_patched/contrib/external_upscale_model.py

@@ -1,23 +1,16 @@
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
 import os
-import logging
-from spandrel import ModelLoader, ImageModelDescriptor
+from ldm_patched.pfn import model_loading
 from ldm_patched.modules import model_management
 import torch
 import ldm_patched.modules.utils
-import ldm_patched.utils.path_utils as folder_paths
-
-try:
-    from spandrel_extra_arches import EXTRA_REGISTRY
-    from spandrel import MAIN_REGISTRY
-    MAIN_REGISTRY.add(*EXTRA_REGISTRY)
-    logging.info("Successfully imported spandrel_extra_arches: support for non commercial upscale models.")
-except:
-    pass
+import ldm_patched.utils.path_utils
 
 class UpscaleModelLoader:
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": { "model_name": (folder_paths.get_filename_list("upscale_models"), ),
+        return {"required": { "model_name": (ldm_patched.utils.path_utils.get_filename_list("upscale_models"), ),
                              }}
     RETURN_TYPES = ("UPSCALE_MODEL",)
     FUNCTION = "load_model"
@@ -25,15 +18,11 @@ class UpscaleModelLoader:
     CATEGORY = "loaders"
 
     def load_model(self, model_name):
-        model_path = folder_paths.get_full_path("upscale_models", model_name)
+        model_path = ldm_patched.utils.path_utils.get_full_path("upscale_models", model_name)
         sd = ldm_patched.modules.utils.load_torch_file(model_path, safe_load=True)
         if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:
             sd = ldm_patched.modules.utils.state_dict_prefix_replace(sd, {"module.":""})
-        out = ModelLoader().load_from_state_dict(sd).eval()
-
-        if not isinstance(out, ImageModelDescriptor):
-            raise Exception("Upscale model must be a single-image model.")
-
+        out = model_loading.load_state_dict(sd).eval()
         return (out, )
 
 
@@ -50,14 +39,9 @@ class ImageUpscaleWithModel:
 
     def upscale(self, upscale_model, image):
         device = model_management.get_torch_device()
-
-        memory_required = model_management.module_size(upscale_model.model)
-        memory_required += (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0 #The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate
-        memory_required += image.nelement() * image.element_size()
-        model_management.free_memory(memory_required, device)
-
         upscale_model.to(device)
         in_img = image.movedim(-1,-3).to(device)
+        free_memory = model_management.get_free_memory(device)
 
         tile = 512
         overlap = 32
@@ -74,7 +58,7 @@ class ImageUpscaleWithModel:
                 if tile < 128:
                     raise e
 
-        upscale_model.to("cpu")
+        upscale_model.cpu()
         s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0)
         return (s,)
 

ldm_patched/contrib/external_video_model.py

@@ -1,15 +1,17 @@
-import external
+# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
+
+import ldm_patched.contrib.external
 import torch
 import ldm_patched.modules.utils
 import ldm_patched.modules.sd
-import ldm_patched.utils.path_utils as folder_paths
-import external_model_merging
+import ldm_patched.utils.path_utils
+import ldm_patched.contrib.external_model_merging
 
 
 class ImageOnlyCheckpointLoader:
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
+        return {"required": { "ckpt_name": (ldm_patched.utils.path_utils.get_filename_list("checkpoints"), ),
                              }}
     RETURN_TYPES = ("MODEL", "CLIP_VISION", "VAE")
     FUNCTION = "load_checkpoint"
@@ -17,8 +19,8 @@ class ImageOnlyCheckpointLoader:
     CATEGORY = "loaders/video_models"
 
     def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True):
-        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
-        out = ldm_patched.modules.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=False, output_clipvision=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
+        ckpt_path = ldm_patched.utils.path_utils.get_full_path("checkpoints", ckpt_name)
+        out = ldm_patched.modules.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=False, output_clipvision=True, embedding_directory=ldm_patched.utils.path_utils.get_folder_paths("embeddings"))
         return (out[0], out[3], out[2])
 
 
@@ -28,8 +30,8 @@ class SVD_img2vid_Conditioning:
         return {"required": { "clip_vision": ("CLIP_VISION",),
                               "init_image": ("IMAGE",),
                               "vae": ("VAE",),
-                              "width": ("INT", {"default": 1024, "min": 16, "max": external.MAX_RESOLUTION, "step": 8}),
-                              "height": ("INT", {"default": 576, "min": 16, "max": external.MAX_RESOLUTION, "step": 8}),
+                              "width": ("INT", {"default": 1024, "min": 16, "max": ldm_patched.contrib.external.MAX_RESOLUTION, "step": 8}),
+                              "height": ("INT", {"default": 576, "min": 16, "max": ldm_patched.contrib.external.MAX_RESOLUTION, "step": 8}),
                               "video_frames": ("INT", {"default": 14, "min": 1, "max": 4096}),
                               "motion_bucket_id": ("INT", {"default": 127, "min": 1, "max": 1023}),
                               "fps": ("INT", {"default": 6, "min": 1, "max": 1024}),
@@ -79,34 +81,7 @@ class VideoLinearCFGGuidance:
         m.set_model_sampler_cfg_function(linear_cfg)
         return (m, )
 
-class VideoTriangleCFGGuidance:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "min_cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.5, "round": 0.01}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "sampling/video_models"
-
-    def patch(self, model, min_cfg):
-        def linear_cfg(args):
-            cond = args["cond"]
-            uncond = args["uncond"]
-            cond_scale = args["cond_scale"]
-            period = 1.0
-            values = torch.linspace(0, 1, cond.shape[0], device=cond.device)
-            values = 2 * (values / period - torch.floor(values / period + 0.5)).abs()
-            scale = (values * (cond_scale - min_cfg) + min_cfg).reshape((cond.shape[0], 1, 1, 1))
-
-            return uncond + scale * (cond - uncond)
-
-        m = model.clone()
-        m.set_model_sampler_cfg_function(linear_cfg)
-        return (m, )
-
-class ImageOnlyCheckpointSave(external_model_merging.CheckpointSave):
+class ImageOnlyCheckpointSave(ldm_patched.contrib.external_model_merging.CheckpointSave):
     CATEGORY = "_for_testing"
 
     @classmethod
@@ -114,18 +89,17 @@ class ImageOnlyCheckpointSave(external_model_merging.CheckpointSave):
         return {"required": { "model": ("MODEL",),
                               "clip_vision": ("CLIP_VISION",),
                               "vae": ("VAE",),
-                              "filename_prefix": ("STRING", {"default": "checkpoints/ComfyUI"}),},
+                              "filename_prefix": ("STRING", {"default": "checkpoints/ldm_patched"}),},
                 "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},}
 
     def save(self, model, clip_vision, vae, filename_prefix, prompt=None, extra_pnginfo=None):
-        external_model_merging.save_checkpoint(model, clip_vision=clip_vision, vae=vae, filename_prefix=filename_prefix, output_dir=self.output_dir, prompt=prompt, extra_pnginfo=extra_pnginfo)
+        ldm_patched.contrib.external_model_merging.save_checkpoint(model, clip_vision=clip_vision, vae=vae, filename_prefix=filename_prefix, output_dir=self.output_dir, prompt=prompt, extra_pnginfo=extra_pnginfo)
         return {}
 
 NODE_CLASS_MAPPINGS = {
     "ImageOnlyCheckpointLoader": ImageOnlyCheckpointLoader,
     "SVD_img2vid_Conditioning": SVD_img2vid_Conditioning,
     "VideoLinearCFGGuidance": VideoLinearCFGGuidance,
-    "VideoTriangleCFGGuidance": VideoTriangleCFGGuidance,
     "ImageOnlyCheckpointSave": ImageOnlyCheckpointSave,
 }
 

ldm_patched/contrib/external_webcam.py

@@ -1,33 +0,0 @@
-import external
-import ldm_patched.utils.path_utils as folder_paths
-
-MAX_RESOLUTION = external.MAX_RESOLUTION
-
-
-class WebcamCapture(external.LoadImage):
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": ("WEBCAM", {}),
-                "width": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-                "height": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-                "capture_on_queue": ("BOOLEAN", {"default": True}),
-            }
-        }
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "load_capture"
-
-    CATEGORY = "image"
-
-    def load_capture(s, image, **kwargs):
-        return super().load_image(folder_paths.get_annotated_filepath(image))
-
-
-NODE_CLASS_MAPPINGS = {
-    "WebcamCapture": WebcamCapture,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "WebcamCapture": "Webcam Capture",
-}

ldm_patched/controlnet/cldm.py

@@ -52,7 +52,6 @@ class ControlNet(nn.Module):
         adm_in_channels=None,
         transformer_depth_middle=None,
         transformer_depth_output=None,
-        attn_precision=None,
         device=None,
         operations=ldm_patched.modules.ops.disable_weight_init,
         **kwargs,
@@ -203,7 +202,7 @@ class ControlNet(nn.Module):
                             SpatialTransformer(
                                 ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim,
                                 disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
-                                use_checkpoint=use_checkpoint, attn_precision=attn_precision, dtype=self.dtype, device=device, operations=operations
+                                use_checkpoint=use_checkpoint, dtype=self.dtype, device=device, operations=operations
                             )
                         )
                 self.input_blocks.append(TimestepEmbedSequential(*layers))
@@ -263,7 +262,7 @@ class ControlNet(nn.Module):
             mid_block += [SpatialTransformer(  # always uses a self-attn
                             ch, num_heads, dim_head, depth=transformer_depth_middle, context_dim=context_dim,
                             disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,
-                            use_checkpoint=use_checkpoint, attn_precision=attn_precision, dtype=self.dtype, device=device, operations=operations
+                            use_checkpoint=use_checkpoint, dtype=self.dtype, device=device, operations=operations
                         ),
             ResBlock(
                 ch,
@@ -310,3 +309,4 @@ class ControlNet(nn.Module):
         outs.append(self.middle_block_out(h, emb, context))
 
         return outs
+

ldm_patched/k_diffusion/sampling.py

@@ -527,9 +527,6 @@ def sample_dpmpp_2s_ancestral(model, x, sigmas, extra_args=None, callback=None,
 @torch.no_grad()
 def sample_dpmpp_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=1 / 2):
     """DPM-Solver++ (stochastic)."""
-    if len(sigmas) <= 1:
-        return x
-
     sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
     seed = extra_args.get("seed", None)
     noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
@@ -598,8 +595,6 @@ def sample_dpmpp_2m(model, x, sigmas, extra_args=None, callback=None, disable=No
 @torch.no_grad()
 def sample_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint'):
     """DPM-Solver++(2M) SDE."""
-    if len(sigmas) <= 1:
-        return x
 
     if solver_type not in {'heun', 'midpoint'}:
         raise ValueError('solver_type must be \'heun\' or \'midpoint\'')
@@ -647,9 +642,6 @@ def sample_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, disabl
 def sample_dpmpp_3m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
     """DPM-Solver++(3M) SDE."""
 
-    if len(sigmas) <= 1:
-        return x
-
     seed = extra_args.get("seed", None)
     sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
     noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
@@ -698,27 +690,18 @@ def sample_dpmpp_3m_sde(model, x, sigmas, extra_args=None, callback=None, disabl
 
 @torch.no_grad()
 def sample_dpmpp_3m_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    if len(sigmas) <= 1:
-        return x
-
     sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
     noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
     return sample_dpmpp_3m_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler)
 
 @torch.no_grad()
 def sample_dpmpp_2m_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint'):
-    if len(sigmas) <= 1:
-        return x
-
     sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
     noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
     return sample_dpmpp_2m_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, solver_type=solver_type)
 
 @torch.no_grad()
 def sample_dpmpp_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=1 / 2):
-    if len(sigmas) <= 1:
-        return x
-
     sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
     noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
     return sample_dpmpp_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, r=r)
@@ -765,11 +748,10 @@ def sample_lcm(model, x, sigmas, extra_args=None, callback=None, disable=None, n
 
         x = denoised
         if sigmas[i + 1] > 0:
-            x = model.inner_model.inner_model.model_sampling.noise_scaling(sigmas[i + 1], noise_sampler(sigmas[i], sigmas[i + 1]), x)
+            x += sigmas[i + 1] * noise_sampler(sigmas[i], sigmas[i + 1])
     return x
 
 
-
 @torch.no_grad()
 def sample_heunpp2(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
     # From MIT licensed: https://github.com/Carzit/sd-webui-samplers-scheduler/
@@ -825,3 +807,102 @@ def sample_heunpp2(model, x, sigmas, extra_args=None, callback=None, disable=Non
             d_prime = w1 * d + w2 * d_2 + w3 * d_3
             x = x + d_prime * dt
     return x
+
+
+@torch.no_grad()
+def sample_tcd(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None, eta=0.3):
+    extra_args = {} if extra_args is None else extra_args
+    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+    s_in = x.new_ones([x.shape[0]])
+
+    model_sampling = model.inner_model.inner_model.model_sampling
+    timesteps_s = torch.floor((1 - eta) * model_sampling.timestep(sigmas)).to(dtype=torch.long).detach().cpu()
+    timesteps_s[-1] = 0
+    alpha_prod_s = model_sampling.alphas_cumprod[timesteps_s]
+    beta_prod_s = 1 - alpha_prod_s
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)  # predicted_original_sample
+        eps = (x - denoised) / sigmas[i]
+        denoised = alpha_prod_s[i + 1].sqrt() * denoised + beta_prod_s[i + 1].sqrt() * eps
+
+        if callback is not None:
+            callback({"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "denoised": denoised})
+
+        x = denoised
+        if eta > 0 and sigmas[i + 1] > 0:
+            noise = noise_sampler(sigmas[i], sigmas[i + 1])
+            x = x / alpha_prod_s[i+1].sqrt() + noise * (sigmas[i+1]**2 + 1 - 1/alpha_prod_s[i+1]).sqrt()
+        else:
+            x *= torch.sqrt(1.0 + sigmas[i + 1] ** 2)
+
+    return x
+
+
+@torch.no_grad()
+def sample_restart(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1., restart_list=None):
+    """Implements restart sampling in Restart Sampling for Improving Generative Processes (2023)
+    Restart_list format: {min_sigma: [ restart_steps, restart_times, max_sigma]}
+    If restart_list is None: will choose restart_list automatically, otherwise will use the given restart_list
+    """
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+    step_id = 0
+
+    def heun_step(x, old_sigma, new_sigma, second_order=True):
+        nonlocal step_id
+        denoised = model(x, old_sigma * s_in, **extra_args)
+        d = to_d(x, old_sigma, denoised)
+        if callback is not None:
+            callback({'x': x, 'i': step_id, 'sigma': new_sigma, 'sigma_hat': old_sigma, 'denoised': denoised})
+        dt = new_sigma - old_sigma
+        if new_sigma == 0 or not second_order:
+            # Euler method
+            x = x + d * dt
+        else:
+            # Heun's method
+            x_2 = x + d * dt
+            denoised_2 = model(x_2, new_sigma * s_in, **extra_args)
+            d_2 = to_d(x_2, new_sigma, denoised_2)
+            d_prime = (d + d_2) / 2
+            x = x + d_prime * dt
+        step_id += 1
+        return x
+
+    steps = sigmas.shape[0] - 1
+    if restart_list is None:
+        if steps >= 20:
+            restart_steps = 9
+            restart_times = 1
+            if steps >= 36:
+                restart_steps = steps // 4
+                restart_times = 2
+            sigmas = get_sigmas_karras(steps - restart_steps * restart_times, sigmas[-2].item(), sigmas[0].item(), device=sigmas.device)
+            restart_list = {0.1: [restart_steps + 1, restart_times, 2]}
+        else:
+            restart_list = {}
+
+    restart_list = {int(torch.argmin(abs(sigmas - key), dim=0)): value for key, value in restart_list.items()}
+
+    step_list = []
+    for i in range(len(sigmas) - 1):
+        step_list.append((sigmas[i], sigmas[i + 1]))
+        if i + 1 in restart_list:
+            restart_steps, restart_times, restart_max = restart_list[i + 1]
+            min_idx = i + 1
+            max_idx = int(torch.argmin(abs(sigmas - restart_max), dim=0))
+            if max_idx < min_idx:
+                sigma_restart = get_sigmas_karras(restart_steps, sigmas[min_idx].item(), sigmas[max_idx].item(), device=sigmas.device)[:-1]
+                while restart_times > 0:
+                    restart_times -= 1
+                    step_list.extend(zip(sigma_restart[:-1], sigma_restart[1:]))
+
+    last_sigma = None
+    for old_sigma, new_sigma in tqdm(step_list, disable=disable):
+        if last_sigma is None:
+            last_sigma = old_sigma
+        elif last_sigma < old_sigma:
+            x = x + torch.randn_like(x) * s_noise * (old_sigma ** 2 - last_sigma ** 2) ** 0.5
+        x = heun_step(x, old_sigma, new_sigma)
+        last_sigma = new_sigma
+
+    return x

ldm_patched/ldm/cascade/common.py

@@ -1,161 +0,0 @@
-"""
-    This file is part of ComfyUI.
-    Copyright (C) 2024 Stability AI
-
-    This program is free software: you can redistribute it and/or modify
-    it under the terms of the GNU General Public License as published by
-    the Free Software Foundation, either version 3 of the License, or
-    (at your option) any later version.
-
-    This program is distributed in the hope that it will be useful,
-    but WITHOUT ANY WARRANTY; without even the implied warranty of
-    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-    GNU General Public License for more details.
-
-    You should have received a copy of the GNU General Public License
-    along with this program.  If not, see <https://www.gnu.org/licenses/>.
-"""
-
-import torch
-import torch.nn as nn
-from ldm_patched.ldm.modules.attention import optimized_attention
-
-class Linear(torch.nn.Linear):
-    def reset_parameters(self):
-        return None
-
-class Conv2d(torch.nn.Conv2d):
-    def reset_parameters(self):
-        return None
-
-class OptimizedAttention(nn.Module):
-    def __init__(self, c, nhead, dropout=0.0, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.heads = nhead
-
-        self.to_q = operations.Linear(c, c, bias=True, dtype=dtype, device=device)
-        self.to_k = operations.Linear(c, c, bias=True, dtype=dtype, device=device)
-        self.to_v = operations.Linear(c, c, bias=True, dtype=dtype, device=device)
-
-        self.out_proj = operations.Linear(c, c, bias=True, dtype=dtype, device=device)
-
-    def forward(self, q, k, v):
-        q = self.to_q(q)
-        k = self.to_k(k)
-        v = self.to_v(v)
-
-        out = optimized_attention(q, k, v, self.heads)
-
-        return self.out_proj(out)
-
-class Attention2D(nn.Module):
-    def __init__(self, c, nhead, dropout=0.0, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.attn = OptimizedAttention(c, nhead, dtype=dtype, device=device, operations=operations)
-        # self.attn = nn.MultiheadAttention(c, nhead, dropout=dropout, bias=True, batch_first=True, dtype=dtype, device=device)
-
-    def forward(self, x, kv, self_attn=False):
-        orig_shape = x.shape
-        x = x.view(x.size(0), x.size(1), -1).permute(0, 2, 1)  # Bx4xHxW -> Bx(HxW)x4
-        if self_attn:
-            kv = torch.cat([x, kv], dim=1)
-        # x = self.attn(x, kv, kv, need_weights=False)[0]
-        x = self.attn(x, kv, kv)
-        x = x.permute(0, 2, 1).view(*orig_shape)
-        return x
-
-
-def LayerNorm2d_op(operations):
-    class LayerNorm2d(operations.LayerNorm):
-        def __init__(self, *args, **kwargs):
-            super().__init__(*args, **kwargs)
-
-        def forward(self, x):
-            return super().forward(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
-    return LayerNorm2d
-
-class GlobalResponseNorm(nn.Module):
-    "from https://github.com/facebookresearch/ConvNeXt-V2/blob/3608f67cc1dae164790c5d0aead7bf2d73d9719b/models/utils.py#L105"
-    def __init__(self, dim, dtype=None, device=None):
-        super().__init__()
-        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim, dtype=dtype, device=device))
-        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim, dtype=dtype, device=device))
-
-    def forward(self, x):
-        Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
-        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
-        return self.gamma.to(device=x.device, dtype=x.dtype) * (x * Nx) + self.beta.to(device=x.device, dtype=x.dtype) + x
-
-
-class ResBlock(nn.Module):
-    def __init__(self, c, c_skip=0, kernel_size=3, dropout=0.0, dtype=None, device=None, operations=None):  # , num_heads=4, expansion=2):
-        super().__init__()
-        self.depthwise = operations.Conv2d(c, c, kernel_size=kernel_size, padding=kernel_size // 2, groups=c, dtype=dtype, device=device)
-        #         self.depthwise = SAMBlock(c, num_heads, expansion)
-        self.norm = LayerNorm2d_op(operations)(c, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-        self.channelwise = nn.Sequential(
-            operations.Linear(c + c_skip, c * 4, dtype=dtype, device=device),
-            nn.GELU(),
-            GlobalResponseNorm(c * 4, dtype=dtype, device=device),
-            nn.Dropout(dropout),
-            operations.Linear(c * 4, c, dtype=dtype, device=device)
-        )
-
-    def forward(self, x, x_skip=None):
-        x_res = x
-        x = self.norm(self.depthwise(x))
-        if x_skip is not None:
-            x = torch.cat([x, x_skip], dim=1)
-        x = self.channelwise(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
-        return x + x_res
-
-
-class AttnBlock(nn.Module):
-    def __init__(self, c, c_cond, nhead, self_attn=True, dropout=0.0, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.self_attn = self_attn
-        self.norm = LayerNorm2d_op(operations)(c, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-        self.attention = Attention2D(c, nhead, dropout, dtype=dtype, device=device, operations=operations)
-        self.kv_mapper = nn.Sequential(
-            nn.SiLU(),
-            operations.Linear(c_cond, c, dtype=dtype, device=device)
-        )
-
-    def forward(self, x, kv):
-        kv = self.kv_mapper(kv)
-        x = x + self.attention(self.norm(x), kv, self_attn=self.self_attn)
-        return x
-
-
-class FeedForwardBlock(nn.Module):
-    def __init__(self, c, dropout=0.0, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.norm = LayerNorm2d_op(operations)(c, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-        self.channelwise = nn.Sequential(
-            operations.Linear(c, c * 4, dtype=dtype, device=device),
-            nn.GELU(),
-            GlobalResponseNorm(c * 4, dtype=dtype, device=device),
-            nn.Dropout(dropout),
-            operations.Linear(c * 4, c, dtype=dtype, device=device)
-        )
-
-    def forward(self, x):
-        x = x + self.channelwise(self.norm(x).permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
-        return x
-
-
-class TimestepBlock(nn.Module):
-    def __init__(self, c, c_timestep, conds=['sca'], dtype=None, device=None, operations=None):
-        super().__init__()
-        self.mapper = operations.Linear(c_timestep, c * 2, dtype=dtype, device=device)
-        self.conds = conds
-        for cname in conds:
-            setattr(self, f"mapper_{cname}", operations.Linear(c_timestep, c * 2, dtype=dtype, device=device))
-
-    def forward(self, x, t):
-        t = t.chunk(len(self.conds) + 1, dim=1)
-        a, b = self.mapper(t[0])[:, :, None, None].chunk(2, dim=1)
-        for i, c in enumerate(self.conds):
-            ac, bc = getattr(self, f"mapper_{c}")(t[i + 1])[:, :, None, None].chunk(2, dim=1)
-            a, b = a + ac, b + bc
-        return x * (1 + a) + b

ldm_patched/ldm/cascade/controlnet.py

@@ -1,93 +0,0 @@
-"""
-    This file is part of ComfyUI.
-    Copyright (C) 2024 Stability AI
-
-    This program is free software: you can redistribute it and/or modify
-    it under the terms of the GNU General Public License as published by
-    the Free Software Foundation, either version 3 of the License, or
-    (at your option) any later version.
-
-    This program is distributed in the hope that it will be useful,
-    but WITHOUT ANY WARRANTY; without even the implied warranty of
-    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-    GNU General Public License for more details.
-
-    You should have received a copy of the GNU General Public License
-    along with this program.  If not, see <https://www.gnu.org/licenses/>.
-"""
-
-import torch
-import torchvision
-from torch import nn
-from .common import LayerNorm2d_op
-
-
-class CNetResBlock(nn.Module):
-    def __init__(self, c, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.blocks = nn.Sequential(
-            LayerNorm2d_op(operations)(c, dtype=dtype, device=device),
-            nn.GELU(),
-            operations.Conv2d(c, c, kernel_size=3, padding=1),
-            LayerNorm2d_op(operations)(c, dtype=dtype, device=device),
-            nn.GELU(),
-            operations.Conv2d(c, c, kernel_size=3, padding=1),
-        )
-
-    def forward(self, x):
-        return x + self.blocks(x)
-
-
-class ControlNet(nn.Module):
-    def __init__(self, c_in=3, c_proj=2048, proj_blocks=None, bottleneck_mode=None, dtype=None, device=None, operations=nn):
-        super().__init__()
-        if bottleneck_mode is None:
-            bottleneck_mode = 'effnet'
-        self.proj_blocks = proj_blocks
-        if bottleneck_mode == 'effnet':
-            embd_channels = 1280
-            self.backbone = torchvision.models.efficientnet_v2_s().features.eval()
-            if c_in != 3:
-                in_weights = self.backbone[0][0].weight.data
-                self.backbone[0][0] = operations.Conv2d(c_in, 24, kernel_size=3, stride=2, bias=False, dtype=dtype, device=device)
-                if c_in > 3:
-                    # nn.init.constant_(self.backbone[0][0].weight, 0)
-                    self.backbone[0][0].weight.data[:, :3] = in_weights[:, :3].clone()
-                else:
-                    self.backbone[0][0].weight.data = in_weights[:, :c_in].clone()
-        elif bottleneck_mode == 'simple':
-            embd_channels = c_in
-            self.backbone = nn.Sequential(
-                operations.Conv2d(embd_channels, embd_channels * 4, kernel_size=3, padding=1, dtype=dtype, device=device),
-                nn.LeakyReLU(0.2, inplace=True),
-                operations.Conv2d(embd_channels * 4, embd_channels, kernel_size=3, padding=1, dtype=dtype, device=device),
-            )
-        elif bottleneck_mode == 'large':
-            self.backbone = nn.Sequential(
-                operations.Conv2d(c_in, 4096 * 4, kernel_size=1, dtype=dtype, device=device),
-                nn.LeakyReLU(0.2, inplace=True),
-                operations.Conv2d(4096 * 4, 1024, kernel_size=1, dtype=dtype, device=device),
-                *[CNetResBlock(1024, dtype=dtype, device=device, operations=operations) for _ in range(8)],
-                operations.Conv2d(1024, 1280, kernel_size=1, dtype=dtype, device=device),
-            )
-            embd_channels = 1280
-        else:
-            raise ValueError(f'Unknown bottleneck mode: {bottleneck_mode}')
-        self.projections = nn.ModuleList()
-        for _ in range(len(proj_blocks)):
-            self.projections.append(nn.Sequential(
-                operations.Conv2d(embd_channels, embd_channels, kernel_size=1, bias=False, dtype=dtype, device=device),
-                nn.LeakyReLU(0.2, inplace=True),
-                operations.Conv2d(embd_channels, c_proj, kernel_size=1, bias=False, dtype=dtype, device=device),
-            ))
-            # nn.init.constant_(self.projections[-1][-1].weight, 0)  # zero output projection
-        self.xl = False
-        self.input_channels = c_in
-        self.unshuffle_amount = 8
-
-    def forward(self, x):
-        x = self.backbone(x)
-        proj_outputs = [None for _ in range(max(self.proj_blocks) + 1)]
-        for i, idx in enumerate(self.proj_blocks):
-            proj_outputs[idx] = self.projections[i](x)
-        return proj_outputs

ldm_patched/ldm/cascade/stage_a.py

@@ -1,255 +0,0 @@
-"""
-    This file is part of ComfyUI.
-    Copyright (C) 2024 Stability AI
-
-    This program is free software: you can redistribute it and/or modify
-    it under the terms of the GNU General Public License as published by
-    the Free Software Foundation, either version 3 of the License, or
-    (at your option) any later version.
-
-    This program is distributed in the hope that it will be useful,
-    but WITHOUT ANY WARRANTY; without even the implied warranty of
-    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-    GNU General Public License for more details.
-
-    You should have received a copy of the GNU General Public License
-    along with this program.  If not, see <https://www.gnu.org/licenses/>.
-"""
-
-import torch
-from torch import nn
-from torch.autograd import Function
-
-class vector_quantize(Function):
-    @staticmethod
-    def forward(ctx, x, codebook):
-        with torch.no_grad():
-            codebook_sqr = torch.sum(codebook ** 2, dim=1)
-            x_sqr = torch.sum(x ** 2, dim=1, keepdim=True)
-
-            dist = torch.addmm(codebook_sqr + x_sqr, x, codebook.t(), alpha=-2.0, beta=1.0)
-            _, indices = dist.min(dim=1)
-
-            ctx.save_for_backward(indices, codebook)
-            ctx.mark_non_differentiable(indices)
-
-            nn = torch.index_select(codebook, 0, indices)
-            return nn, indices
-
-    @staticmethod
-    def backward(ctx, grad_output, grad_indices):
-        grad_inputs, grad_codebook = None, None
-
-        if ctx.needs_input_grad[0]:
-            grad_inputs = grad_output.clone()
-        if ctx.needs_input_grad[1]:
-            # Gradient wrt. the codebook
-            indices, codebook = ctx.saved_tensors
-
-            grad_codebook = torch.zeros_like(codebook)
-            grad_codebook.index_add_(0, indices, grad_output)
-
-        return (grad_inputs, grad_codebook)
-
-
-class VectorQuantize(nn.Module):
-    def __init__(self, embedding_size, k, ema_decay=0.99, ema_loss=False):
-        """
-        Takes an input of variable size (as long as the last dimension matches the embedding size).
-        Returns one tensor containing the nearest neigbour embeddings to each of the inputs,
-        with the same size as the input, vq and commitment components for the loss as a touple
-        in the second output and the indices of the quantized vectors in the third:
-        quantized, (vq_loss, commit_loss), indices
-        """
-        super(VectorQuantize, self).__init__()
-
-        self.codebook = nn.Embedding(k, embedding_size)
-        self.codebook.weight.data.uniform_(-1./k, 1./k)
-        self.vq = vector_quantize.apply
-
-        self.ema_decay = ema_decay
-        self.ema_loss = ema_loss
-        if ema_loss:
-            self.register_buffer('ema_element_count', torch.ones(k))
-            self.register_buffer('ema_weight_sum', torch.zeros_like(self.codebook.weight))
-
-    def _laplace_smoothing(self, x, epsilon):
-        n = torch.sum(x)
-        return ((x + epsilon) / (n + x.size(0) * epsilon) * n)
-
-    def _updateEMA(self, z_e_x, indices):
-        mask = nn.functional.one_hot(indices, self.ema_element_count.size(0)).float()
-        elem_count = mask.sum(dim=0)
-        weight_sum = torch.mm(mask.t(), z_e_x)
-
-        self.ema_element_count = (self.ema_decay * self.ema_element_count) + ((1-self.ema_decay) * elem_count)
-        self.ema_element_count = self._laplace_smoothing(self.ema_element_count, 1e-5)
-        self.ema_weight_sum = (self.ema_decay * self.ema_weight_sum) + ((1-self.ema_decay) * weight_sum)
-
-        self.codebook.weight.data = self.ema_weight_sum / self.ema_element_count.unsqueeze(-1)
-
-    def idx2vq(self, idx, dim=-1):
-        q_idx = self.codebook(idx)
-        if dim != -1:
-            q_idx = q_idx.movedim(-1, dim)
-        return q_idx
-
-    def forward(self, x, get_losses=True, dim=-1):
-        if dim != -1:
-            x = x.movedim(dim, -1)
-        z_e_x = x.contiguous().view(-1, x.size(-1)) if len(x.shape) > 2 else x
-        z_q_x, indices = self.vq(z_e_x, self.codebook.weight.detach())
-        vq_loss, commit_loss = None, None
-        if self.ema_loss and self.training:
-            self._updateEMA(z_e_x.detach(), indices.detach())
-        # pick the graded embeddings after updating the codebook in order to have a more accurate commitment loss
-        z_q_x_grd = torch.index_select(self.codebook.weight, dim=0, index=indices)
-        if get_losses:
-            vq_loss = (z_q_x_grd - z_e_x.detach()).pow(2).mean()
-            commit_loss = (z_e_x - z_q_x_grd.detach()).pow(2).mean()
-
-        z_q_x = z_q_x.view(x.shape)
-        if dim != -1:
-            z_q_x = z_q_x.movedim(-1, dim)
-        return z_q_x, (vq_loss, commit_loss), indices.view(x.shape[:-1])
-
-
-class ResBlock(nn.Module):
-    def __init__(self, c, c_hidden):
-        super().__init__()
-        # depthwise/attention
-        self.norm1 = nn.LayerNorm(c, elementwise_affine=False, eps=1e-6)
-        self.depthwise = nn.Sequential(
-            nn.ReplicationPad2d(1),
-            nn.Conv2d(c, c, kernel_size=3, groups=c)
-        )
-
-        # channelwise
-        self.norm2 = nn.LayerNorm(c, elementwise_affine=False, eps=1e-6)
-        self.channelwise = nn.Sequential(
-            nn.Linear(c, c_hidden),
-            nn.GELU(),
-            nn.Linear(c_hidden, c),
-        )
-
-        self.gammas = nn.Parameter(torch.zeros(6), requires_grad=True)
-
-        # Init weights
-        def _basic_init(module):
-            if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
-                torch.nn.init.xavier_uniform_(module.weight)
-                if module.bias is not None:
-                    nn.init.constant_(module.bias, 0)
-
-        self.apply(_basic_init)
-
-    def _norm(self, x, norm):
-        return norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
-
-    def forward(self, x):
-        mods = self.gammas
-
-        x_temp = self._norm(x, self.norm1) * (1 + mods[0]) + mods[1]
-        try:
-            x = x + self.depthwise(x_temp) * mods[2]
-        except: #operation not implemented for bf16
-            x_temp = self.depthwise[0](x_temp.float()).to(x.dtype)
-            x = x + self.depthwise[1](x_temp) * mods[2]
-
-        x_temp = self._norm(x, self.norm2) * (1 + mods[3]) + mods[4]
-        x = x + self.channelwise(x_temp.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) * mods[5]
-
-        return x
-
-
-class StageA(nn.Module):
-    def __init__(self, levels=2, bottleneck_blocks=12, c_hidden=384, c_latent=4, codebook_size=8192):
-        super().__init__()
-        self.c_latent = c_latent
-        c_levels = [c_hidden // (2 ** i) for i in reversed(range(levels))]
-
-        # Encoder blocks
-        self.in_block = nn.Sequential(
-            nn.PixelUnshuffle(2),
-            nn.Conv2d(3 * 4, c_levels[0], kernel_size=1)
-        )
-        down_blocks = []
-        for i in range(levels):
-            if i > 0:
-                down_blocks.append(nn.Conv2d(c_levels[i - 1], c_levels[i], kernel_size=4, stride=2, padding=1))
-            block = ResBlock(c_levels[i], c_levels[i] * 4)
-            down_blocks.append(block)
-        down_blocks.append(nn.Sequential(
-            nn.Conv2d(c_levels[-1], c_latent, kernel_size=1, bias=False),
-            nn.BatchNorm2d(c_latent),  # then normalize them to have mean 0 and std 1
-        ))
-        self.down_blocks = nn.Sequential(*down_blocks)
-        self.down_blocks[0]
-
-        self.codebook_size = codebook_size
-        self.vquantizer = VectorQuantize(c_latent, k=codebook_size)
-
-        # Decoder blocks
-        up_blocks = [nn.Sequential(
-            nn.Conv2d(c_latent, c_levels[-1], kernel_size=1)
-        )]
-        for i in range(levels):
-            for j in range(bottleneck_blocks if i == 0 else 1):
-                block = ResBlock(c_levels[levels - 1 - i], c_levels[levels - 1 - i] * 4)
-                up_blocks.append(block)
-            if i < levels - 1:
-                up_blocks.append(
-                    nn.ConvTranspose2d(c_levels[levels - 1 - i], c_levels[levels - 2 - i], kernel_size=4, stride=2,
-                                       padding=1))
-        self.up_blocks = nn.Sequential(*up_blocks)
-        self.out_block = nn.Sequential(
-            nn.Conv2d(c_levels[0], 3 * 4, kernel_size=1),
-            nn.PixelShuffle(2),
-        )
-
-    def encode(self, x, quantize=False):
-        x = self.in_block(x)
-        x = self.down_blocks(x)
-        if quantize:
-            qe, (vq_loss, commit_loss), indices = self.vquantizer.forward(x, dim=1)
-            return qe, x, indices, vq_loss + commit_loss * 0.25
-        else:
-            return x
-
-    def decode(self, x):
-        x = self.up_blocks(x)
-        x = self.out_block(x)
-        return x
-
-    def forward(self, x, quantize=False):
-        qe, x, _, vq_loss = self.encode(x, quantize)
-        x = self.decode(qe)
-        return x, vq_loss
-
-
-class Discriminator(nn.Module):
-    def __init__(self, c_in=3, c_cond=0, c_hidden=512, depth=6):
-        super().__init__()
-        d = max(depth - 3, 3)
-        layers = [
-            nn.utils.spectral_norm(nn.Conv2d(c_in, c_hidden // (2 ** d), kernel_size=3, stride=2, padding=1)),
-            nn.LeakyReLU(0.2),
-        ]
-        for i in range(depth - 1):
-            c_in = c_hidden // (2 ** max((d - i), 0))
-            c_out = c_hidden // (2 ** max((d - 1 - i), 0))
-            layers.append(nn.utils.spectral_norm(nn.Conv2d(c_in, c_out, kernel_size=3, stride=2, padding=1)))
-            layers.append(nn.InstanceNorm2d(c_out))
-            layers.append(nn.LeakyReLU(0.2))
-        self.encoder = nn.Sequential(*layers)
-        self.shuffle = nn.Conv2d((c_hidden + c_cond) if c_cond > 0 else c_hidden, 1, kernel_size=1)
-        self.logits = nn.Sigmoid()
-
-    def forward(self, x, cond=None):
-        x = self.encoder(x)
-        if cond is not None:
-            cond = cond.view(cond.size(0), cond.size(1), 1, 1, ).expand(-1, -1, x.size(-2), x.size(-1))
-            x = torch.cat([x, cond], dim=1)
-        x = self.shuffle(x)
-        x = self.logits(x)
-        return x

ldm_patched/ldm/cascade/stage_b.py

@@ -1,256 +0,0 @@
-"""
-    This file is part of ComfyUI.
-    Copyright (C) 2024 Stability AI
-
-    This program is free software: you can redistribute it and/or modify
-    it under the terms of the GNU General Public License as published by
-    the Free Software Foundation, either version 3 of the License, or
-    (at your option) any later version.
-
-    This program is distributed in the hope that it will be useful,
-    but WITHOUT ANY WARRANTY; without even the implied warranty of
-    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-    GNU General Public License for more details.
-
-    You should have received a copy of the GNU General Public License
-    along with this program.  If not, see <https://www.gnu.org/licenses/>.
-"""
-
-import math
-import torch
-from torch import nn
-from .common import AttnBlock, LayerNorm2d_op, ResBlock, FeedForwardBlock, TimestepBlock
-
-class StageB(nn.Module):
-    def __init__(self, c_in=4, c_out=4, c_r=64, patch_size=2, c_cond=1280, c_hidden=[320, 640, 1280, 1280],
-                 nhead=[-1, -1, 20, 20], blocks=[[2, 6, 28, 6], [6, 28, 6, 2]],
-                 block_repeat=[[1, 1, 1, 1], [3, 3, 2, 2]], level_config=['CT', 'CT', 'CTA', 'CTA'], c_clip=1280,
-                 c_clip_seq=4, c_effnet=16, c_pixels=3, kernel_size=3, dropout=[0, 0, 0.0, 0.0], self_attn=True,
-                 t_conds=['sca'], stable_cascade_stage=None, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.dtype = dtype
-        self.c_r = c_r
-        self.t_conds = t_conds
-        self.c_clip_seq = c_clip_seq
-        if not isinstance(dropout, list):
-            dropout = [dropout] * len(c_hidden)
-        if not isinstance(self_attn, list):
-            self_attn = [self_attn] * len(c_hidden)
-
-        # CONDITIONING
-        self.effnet_mapper = nn.Sequential(
-            operations.Conv2d(c_effnet, c_hidden[0] * 4, kernel_size=1, dtype=dtype, device=device),
-            nn.GELU(),
-            operations.Conv2d(c_hidden[0] * 4, c_hidden[0], kernel_size=1, dtype=dtype, device=device),
-            LayerNorm2d_op(operations)(c_hidden[0], elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-        )
-        self.pixels_mapper = nn.Sequential(
-            operations.Conv2d(c_pixels, c_hidden[0] * 4, kernel_size=1, dtype=dtype, device=device),
-            nn.GELU(),
-            operations.Conv2d(c_hidden[0] * 4, c_hidden[0], kernel_size=1, dtype=dtype, device=device),
-            LayerNorm2d_op(operations)(c_hidden[0], elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-        )
-        self.clip_mapper = operations.Linear(c_clip, c_cond * c_clip_seq, dtype=dtype, device=device)
-        self.clip_norm = operations.LayerNorm(c_cond, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-
-        self.embedding = nn.Sequential(
-            nn.PixelUnshuffle(patch_size),
-            operations.Conv2d(c_in * (patch_size ** 2), c_hidden[0], kernel_size=1, dtype=dtype, device=device),
-            LayerNorm2d_op(operations)(c_hidden[0], elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-        )
-
-        def get_block(block_type, c_hidden, nhead, c_skip=0, dropout=0, self_attn=True):
-            if block_type == 'C':
-                return ResBlock(c_hidden, c_skip, kernel_size=kernel_size, dropout=dropout, dtype=dtype, device=device, operations=operations)
-            elif block_type == 'A':
-                return AttnBlock(c_hidden, c_cond, nhead, self_attn=self_attn, dropout=dropout, dtype=dtype, device=device, operations=operations)
-            elif block_type == 'F':
-                return FeedForwardBlock(c_hidden, dropout=dropout, dtype=dtype, device=device, operations=operations)
-            elif block_type == 'T':
-                return TimestepBlock(c_hidden, c_r, conds=t_conds, dtype=dtype, device=device, operations=operations)
-            else:
-                raise Exception(f'Block type {block_type} not supported')
-
-        # BLOCKS
-        # -- down blocks
-        self.down_blocks = nn.ModuleList()
-        self.down_downscalers = nn.ModuleList()
-        self.down_repeat_mappers = nn.ModuleList()
-        for i in range(len(c_hidden)):
-            if i > 0:
-                self.down_downscalers.append(nn.Sequential(
-                    LayerNorm2d_op(operations)(c_hidden[i - 1], elementwise_affine=False, eps=1e-6, dtype=dtype, device=device),
-                    operations.Conv2d(c_hidden[i - 1], c_hidden[i], kernel_size=2, stride=2, dtype=dtype, device=device),
-                ))
-            else:
-                self.down_downscalers.append(nn.Identity())
-            down_block = nn.ModuleList()
-            for _ in range(blocks[0][i]):
-                for block_type in level_config[i]:
-                    block = get_block(block_type, c_hidden[i], nhead[i], dropout=dropout[i], self_attn=self_attn[i])
-                    down_block.append(block)
-            self.down_blocks.append(down_block)
-            if block_repeat is not None:
-                block_repeat_mappers = nn.ModuleList()
-                for _ in range(block_repeat[0][i] - 1):
-                    block_repeat_mappers.append(operations.Conv2d(c_hidden[i], c_hidden[i], kernel_size=1, dtype=dtype, device=device))
-                self.down_repeat_mappers.append(block_repeat_mappers)
-
-        # -- up blocks
-        self.up_blocks = nn.ModuleList()
-        self.up_upscalers = nn.ModuleList()
-        self.up_repeat_mappers = nn.ModuleList()
-        for i in reversed(range(len(c_hidden))):
-            if i > 0:
-                self.up_upscalers.append(nn.Sequential(
-                    LayerNorm2d_op(operations)(c_hidden[i], elementwise_affine=False, eps=1e-6, dtype=dtype, device=device),
-                    operations.ConvTranspose2d(c_hidden[i], c_hidden[i - 1], kernel_size=2, stride=2, dtype=dtype, device=device),
-                ))
-            else:
-                self.up_upscalers.append(nn.Identity())
-            up_block = nn.ModuleList()
-            for j in range(blocks[1][::-1][i]):
-                for k, block_type in enumerate(level_config[i]):
-                    c_skip = c_hidden[i] if i < len(c_hidden) - 1 and j == k == 0 else 0
-                    block = get_block(block_type, c_hidden[i], nhead[i], c_skip=c_skip, dropout=dropout[i],
-                                      self_attn=self_attn[i])
-                    up_block.append(block)
-            self.up_blocks.append(up_block)
-            if block_repeat is not None:
-                block_repeat_mappers = nn.ModuleList()
-                for _ in range(block_repeat[1][::-1][i] - 1):
-                    block_repeat_mappers.append(operations.Conv2d(c_hidden[i], c_hidden[i], kernel_size=1, dtype=dtype, device=device))
-                self.up_repeat_mappers.append(block_repeat_mappers)
-
-        # OUTPUT
-        self.clf = nn.Sequential(
-            LayerNorm2d_op(operations)(c_hidden[0], elementwise_affine=False, eps=1e-6, dtype=dtype, device=device),
-            operations.Conv2d(c_hidden[0], c_out * (patch_size ** 2), kernel_size=1, dtype=dtype, device=device),
-            nn.PixelShuffle(patch_size),
-        )
-
-        # --- WEIGHT INIT ---
-    #     self.apply(self._init_weights)  # General init
-    #     nn.init.normal_(self.clip_mapper.weight, std=0.02)  # conditionings
-    #     nn.init.normal_(self.effnet_mapper[0].weight, std=0.02)  # conditionings
-    #     nn.init.normal_(self.effnet_mapper[2].weight, std=0.02)  # conditionings
-    #     nn.init.normal_(self.pixels_mapper[0].weight, std=0.02)  # conditionings
-    #     nn.init.normal_(self.pixels_mapper[2].weight, std=0.02)  # conditionings
-    #     torch.nn.init.xavier_uniform_(self.embedding[1].weight, 0.02)  # inputs
-    #     nn.init.constant_(self.clf[1].weight, 0)  # outputs
-    # 
-    #     # blocks
-    #     for level_block in self.down_blocks + self.up_blocks:
-    #         for block in level_block:
-    #             if isinstance(block, ResBlock) or isinstance(block, FeedForwardBlock):
-    #                 block.channelwise[-1].weight.data *= np.sqrt(1 / sum(blocks[0]))
-    #             elif isinstance(block, TimestepBlock):
-    #                 for layer in block.modules():
-    #                     if isinstance(layer, nn.Linear):
-    #                         nn.init.constant_(layer.weight, 0)
-    # 
-    # def _init_weights(self, m):
-    #     if isinstance(m, (nn.Conv2d, nn.Linear)):
-    #         torch.nn.init.xavier_uniform_(m.weight)
-    #         if m.bias is not None:
-    #             nn.init.constant_(m.bias, 0)
-
-    def gen_r_embedding(self, r, max_positions=10000):
-        r = r * max_positions
-        half_dim = self.c_r // 2
-        emb = math.log(max_positions) / (half_dim - 1)
-        emb = torch.arange(half_dim, device=r.device).float().mul(-emb).exp()
-        emb = r[:, None] * emb[None, :]
-        emb = torch.cat([emb.sin(), emb.cos()], dim=1)
-        if self.c_r % 2 == 1:  # zero pad
-            emb = nn.functional.pad(emb, (0, 1), mode='constant')
-        return emb
-
-    def gen_c_embeddings(self, clip):
-        if len(clip.shape) == 2:
-            clip = clip.unsqueeze(1)
-        clip = self.clip_mapper(clip).view(clip.size(0), clip.size(1) * self.c_clip_seq, -1)
-        clip = self.clip_norm(clip)
-        return clip
-
-    def _down_encode(self, x, r_embed, clip):
-        level_outputs = []
-        block_group = zip(self.down_blocks, self.down_downscalers, self.down_repeat_mappers)
-        for down_block, downscaler, repmap in block_group:
-            x = downscaler(x)
-            for i in range(len(repmap) + 1):
-                for block in down_block:
-                    if isinstance(block, ResBlock) or (
-                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
-                                                                                  ResBlock)):
-                        x = block(x)
-                    elif isinstance(block, AttnBlock) or (
-                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
-                                                                                  AttnBlock)):
-                        x = block(x, clip)
-                    elif isinstance(block, TimestepBlock) or (
-                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
-                                                                                  TimestepBlock)):
-                        x = block(x, r_embed)
-                    else:
-                        x = block(x)
-                if i < len(repmap):
-                    x = repmap[i](x)
-            level_outputs.insert(0, x)
-        return level_outputs
-
-    def _up_decode(self, level_outputs, r_embed, clip):
-        x = level_outputs[0]
-        block_group = zip(self.up_blocks, self.up_upscalers, self.up_repeat_mappers)
-        for i, (up_block, upscaler, repmap) in enumerate(block_group):
-            for j in range(len(repmap) + 1):
-                for k, block in enumerate(up_block):
-                    if isinstance(block, ResBlock) or (
-                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
-                                                                                  ResBlock)):
-                        skip = level_outputs[i] if k == 0 and i > 0 else None
-                        if skip is not None and (x.size(-1) != skip.size(-1) or x.size(-2) != skip.size(-2)):
-                            x = torch.nn.functional.interpolate(x, skip.shape[-2:], mode='bilinear',
-                                                                align_corners=True)
-                        x = block(x, skip)
-                    elif isinstance(block, AttnBlock) or (
-                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
-                                                                                  AttnBlock)):
-                        x = block(x, clip)
-                    elif isinstance(block, TimestepBlock) or (
-                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
-                                                                                  TimestepBlock)):
-                        x = block(x, r_embed)
-                    else:
-                        x = block(x)
-                if j < len(repmap):
-                    x = repmap[j](x)
-            x = upscaler(x)
-        return x
-
-    def forward(self, x, r, effnet, clip, pixels=None, **kwargs):
-        if pixels is None:
-            pixels = x.new_zeros(x.size(0), 3, 8, 8)
-
-        # Process the conditioning embeddings
-        r_embed = self.gen_r_embedding(r).to(dtype=x.dtype)
-        for c in self.t_conds:
-            t_cond = kwargs.get(c, torch.zeros_like(r))
-            r_embed = torch.cat([r_embed, self.gen_r_embedding(t_cond).to(dtype=x.dtype)], dim=1)
-        clip = self.gen_c_embeddings(clip)
-
-        # Model Blocks
-        x = self.embedding(x)
-        x = x + self.effnet_mapper(
-            nn.functional.interpolate(effnet, size=x.shape[-2:], mode='bilinear', align_corners=True))
-        x = x + nn.functional.interpolate(self.pixels_mapper(pixels), size=x.shape[-2:], mode='bilinear',
-                                          align_corners=True)
-        level_outputs = self._down_encode(x, r_embed, clip)
-        x = self._up_decode(level_outputs, r_embed, clip)
-        return self.clf(x)
-
-    def update_weights_ema(self, src_model, beta=0.999):
-        for self_params, src_params in zip(self.parameters(), src_model.parameters()):
-            self_params.data = self_params.data * beta + src_params.data.clone().to(self_params.device) * (1 - beta)
-        for self_buffers, src_buffers in zip(self.buffers(), src_model.buffers()):
-            self_buffers.data = self_buffers.data * beta + src_buffers.data.clone().to(self_buffers.device) * (1 - beta)

ldm_patched/ldm/cascade/stage_c.py

@@ -1,273 +0,0 @@
-"""
-    This file is part of ComfyUI.
-    Copyright (C) 2024 Stability AI
-
-    This program is free software: you can redistribute it and/or modify
-    it under the terms of the GNU General Public License as published by
-    the Free Software Foundation, either version 3 of the License, or
-    (at your option) any later version.
-
-    This program is distributed in the hope that it will be useful,
-    but WITHOUT ANY WARRANTY; without even the implied warranty of
-    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-    GNU General Public License for more details.
-
-    You should have received a copy of the GNU General Public License
-    along with this program.  If not, see <https://www.gnu.org/licenses/>.
-"""
-
-import torch
-from torch import nn
-import math
-from .common import AttnBlock, LayerNorm2d_op, ResBlock, FeedForwardBlock, TimestepBlock
-# from .controlnet import ControlNetDeliverer
-
-class UpDownBlock2d(nn.Module):
-    def __init__(self, c_in, c_out, mode, enabled=True, dtype=None, device=None, operations=None):
-        super().__init__()
-        assert mode in ['up', 'down']
-        interpolation = nn.Upsample(scale_factor=2 if mode == 'up' else 0.5, mode='bilinear',
-                                    align_corners=True) if enabled else nn.Identity()
-        mapping = operations.Conv2d(c_in, c_out, kernel_size=1, dtype=dtype, device=device)
-        self.blocks = nn.ModuleList([interpolation, mapping] if mode == 'up' else [mapping, interpolation])
-
-    def forward(self, x):
-        for block in self.blocks:
-            x = block(x)
-        return x
-
-
-class StageC(nn.Module):
-    def __init__(self, c_in=16, c_out=16, c_r=64, patch_size=1, c_cond=2048, c_hidden=[2048, 2048], nhead=[32, 32],
-                 blocks=[[8, 24], [24, 8]], block_repeat=[[1, 1], [1, 1]], level_config=['CTA', 'CTA'],
-                 c_clip_text=1280, c_clip_text_pooled=1280, c_clip_img=768, c_clip_seq=4, kernel_size=3,
-                 dropout=[0.0, 0.0], self_attn=True, t_conds=['sca', 'crp'], switch_level=[False], stable_cascade_stage=None,
-                 dtype=None, device=None, operations=None):
-        super().__init__()
-        self.dtype = dtype
-        self.c_r = c_r
-        self.t_conds = t_conds
-        self.c_clip_seq = c_clip_seq
-        if not isinstance(dropout, list):
-            dropout = [dropout] * len(c_hidden)
-        if not isinstance(self_attn, list):
-            self_attn = [self_attn] * len(c_hidden)
-
-        # CONDITIONING
-        self.clip_txt_mapper = operations.Linear(c_clip_text, c_cond, dtype=dtype, device=device)
-        self.clip_txt_pooled_mapper = operations.Linear(c_clip_text_pooled, c_cond * c_clip_seq, dtype=dtype, device=device)
-        self.clip_img_mapper = operations.Linear(c_clip_img, c_cond * c_clip_seq, dtype=dtype, device=device)
-        self.clip_norm = operations.LayerNorm(c_cond, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-
-        self.embedding = nn.Sequential(
-            nn.PixelUnshuffle(patch_size),
-            operations.Conv2d(c_in * (patch_size ** 2), c_hidden[0], kernel_size=1, dtype=dtype, device=device),
-            LayerNorm2d_op(operations)(c_hidden[0], elementwise_affine=False, eps=1e-6)
-        )
-
-        def get_block(block_type, c_hidden, nhead, c_skip=0, dropout=0, self_attn=True):
-            if block_type == 'C':
-                return ResBlock(c_hidden, c_skip, kernel_size=kernel_size, dropout=dropout, dtype=dtype, device=device, operations=operations)
-            elif block_type == 'A':
-                return AttnBlock(c_hidden, c_cond, nhead, self_attn=self_attn, dropout=dropout, dtype=dtype, device=device, operations=operations)
-            elif block_type == 'F':
-                return FeedForwardBlock(c_hidden, dropout=dropout, dtype=dtype, device=device, operations=operations)
-            elif block_type == 'T':
-                return TimestepBlock(c_hidden, c_r, conds=t_conds, dtype=dtype, device=device, operations=operations)
-            else:
-                raise Exception(f'Block type {block_type} not supported')
-
-        # BLOCKS
-        # -- down blocks
-        self.down_blocks = nn.ModuleList()
-        self.down_downscalers = nn.ModuleList()
-        self.down_repeat_mappers = nn.ModuleList()
-        for i in range(len(c_hidden)):
-            if i > 0:
-                self.down_downscalers.append(nn.Sequential(
-                    LayerNorm2d_op(operations)(c_hidden[i - 1], elementwise_affine=False, eps=1e-6),
-                    UpDownBlock2d(c_hidden[i - 1], c_hidden[i], mode='down', enabled=switch_level[i - 1], dtype=dtype, device=device, operations=operations)
-                ))
-            else:
-                self.down_downscalers.append(nn.Identity())
-            down_block = nn.ModuleList()
-            for _ in range(blocks[0][i]):
-                for block_type in level_config[i]:
-                    block = get_block(block_type, c_hidden[i], nhead[i], dropout=dropout[i], self_attn=self_attn[i])
-                    down_block.append(block)
-            self.down_blocks.append(down_block)
-            if block_repeat is not None:
-                block_repeat_mappers = nn.ModuleList()
-                for _ in range(block_repeat[0][i] - 1):
-                    block_repeat_mappers.append(operations.Conv2d(c_hidden[i], c_hidden[i], kernel_size=1, dtype=dtype, device=device))
-                self.down_repeat_mappers.append(block_repeat_mappers)
-
-        # -- up blocks
-        self.up_blocks = nn.ModuleList()
-        self.up_upscalers = nn.ModuleList()
-        self.up_repeat_mappers = nn.ModuleList()
-        for i in reversed(range(len(c_hidden))):
-            if i > 0:
-                self.up_upscalers.append(nn.Sequential(
-                    LayerNorm2d_op(operations)(c_hidden[i], elementwise_affine=False, eps=1e-6),
-                    UpDownBlock2d(c_hidden[i], c_hidden[i - 1], mode='up', enabled=switch_level[i - 1], dtype=dtype, device=device, operations=operations)
-                ))
-            else:
-                self.up_upscalers.append(nn.Identity())
-            up_block = nn.ModuleList()
-            for j in range(blocks[1][::-1][i]):
-                for k, block_type in enumerate(level_config[i]):
-                    c_skip = c_hidden[i] if i < len(c_hidden) - 1 and j == k == 0 else 0
-                    block = get_block(block_type, c_hidden[i], nhead[i], c_skip=c_skip, dropout=dropout[i],
-                                      self_attn=self_attn[i])
-                    up_block.append(block)
-            self.up_blocks.append(up_block)
-            if block_repeat is not None:
-                block_repeat_mappers = nn.ModuleList()
-                for _ in range(block_repeat[1][::-1][i] - 1):
-                    block_repeat_mappers.append(operations.Conv2d(c_hidden[i], c_hidden[i], kernel_size=1, dtype=dtype, device=device))
-                self.up_repeat_mappers.append(block_repeat_mappers)
-
-        # OUTPUT
-        self.clf = nn.Sequential(
-            LayerNorm2d_op(operations)(c_hidden[0], elementwise_affine=False, eps=1e-6, dtype=dtype, device=device),
-            operations.Conv2d(c_hidden[0], c_out * (patch_size ** 2), kernel_size=1, dtype=dtype, device=device),
-            nn.PixelShuffle(patch_size),
-        )
-
-        # --- WEIGHT INIT ---
-    #     self.apply(self._init_weights)  # General init
-    #     nn.init.normal_(self.clip_txt_mapper.weight, std=0.02)  # conditionings
-    #     nn.init.normal_(self.clip_txt_pooled_mapper.weight, std=0.02)  # conditionings
-    #     nn.init.normal_(self.clip_img_mapper.weight, std=0.02)  # conditionings
-    #     torch.nn.init.xavier_uniform_(self.embedding[1].weight, 0.02)  # inputs
-    #     nn.init.constant_(self.clf[1].weight, 0)  # outputs
-    # 
-    #     # blocks
-    #     for level_block in self.down_blocks + self.up_blocks:
-    #         for block in level_block:
-    #             if isinstance(block, ResBlock) or isinstance(block, FeedForwardBlock):
-    #                 block.channelwise[-1].weight.data *= np.sqrt(1 / sum(blocks[0]))
-    #             elif isinstance(block, TimestepBlock):
-    #                 for layer in block.modules():
-    #                     if isinstance(layer, nn.Linear):
-    #                         nn.init.constant_(layer.weight, 0)
-    # 
-    # def _init_weights(self, m):
-    #     if isinstance(m, (nn.Conv2d, nn.Linear)):
-    #         torch.nn.init.xavier_uniform_(m.weight)
-    #         if m.bias is not None:
-    #             nn.init.constant_(m.bias, 0)
-
-    def gen_r_embedding(self, r, max_positions=10000):
-        r = r * max_positions
-        half_dim = self.c_r // 2
-        emb = math.log(max_positions) / (half_dim - 1)
-        emb = torch.arange(half_dim, device=r.device).float().mul(-emb).exp()
-        emb = r[:, None] * emb[None, :]
-        emb = torch.cat([emb.sin(), emb.cos()], dim=1)
-        if self.c_r % 2 == 1:  # zero pad
-            emb = nn.functional.pad(emb, (0, 1), mode='constant')
-        return emb
-
-    def gen_c_embeddings(self, clip_txt, clip_txt_pooled, clip_img):
-        clip_txt = self.clip_txt_mapper(clip_txt)
-        if len(clip_txt_pooled.shape) == 2:
-            clip_txt_pooled = clip_txt_pooled.unsqueeze(1)
-        if len(clip_img.shape) == 2:
-            clip_img = clip_img.unsqueeze(1)
-        clip_txt_pool = self.clip_txt_pooled_mapper(clip_txt_pooled).view(clip_txt_pooled.size(0), clip_txt_pooled.size(1) * self.c_clip_seq, -1)
-        clip_img = self.clip_img_mapper(clip_img).view(clip_img.size(0), clip_img.size(1) * self.c_clip_seq, -1)
-        clip = torch.cat([clip_txt, clip_txt_pool, clip_img], dim=1)
-        clip = self.clip_norm(clip)
-        return clip
-
-    def _down_encode(self, x, r_embed, clip, cnet=None):
-        level_outputs = []
-        block_group = zip(self.down_blocks, self.down_downscalers, self.down_repeat_mappers)
-        for down_block, downscaler, repmap in block_group:
-            x = downscaler(x)
-            for i in range(len(repmap) + 1):
-                for block in down_block:
-                    if isinstance(block, ResBlock) or (
-                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
-                                                                                  ResBlock)):
-                        if cnet is not None:
-                            next_cnet = cnet.pop()
-                            if next_cnet is not None:
-                                x = x + nn.functional.interpolate(next_cnet, size=x.shape[-2:], mode='bilinear',
-                                                                  align_corners=True).to(x.dtype)
-                        x = block(x)
-                    elif isinstance(block, AttnBlock) or (
-                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
-                                                                                  AttnBlock)):
-                        x = block(x, clip)
-                    elif isinstance(block, TimestepBlock) or (
-                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
-                                                                                  TimestepBlock)):
-                        x = block(x, r_embed)
-                    else:
-                        x = block(x)
-                if i < len(repmap):
-                    x = repmap[i](x)
-            level_outputs.insert(0, x)
-        return level_outputs
-
-    def _up_decode(self, level_outputs, r_embed, clip, cnet=None):
-        x = level_outputs[0]
-        block_group = zip(self.up_blocks, self.up_upscalers, self.up_repeat_mappers)
-        for i, (up_block, upscaler, repmap) in enumerate(block_group):
-            for j in range(len(repmap) + 1):
-                for k, block in enumerate(up_block):
-                    if isinstance(block, ResBlock) or (
-                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
-                                                                                  ResBlock)):
-                        skip = level_outputs[i] if k == 0 and i > 0 else None
-                        if skip is not None and (x.size(-1) != skip.size(-1) or x.size(-2) != skip.size(-2)):
-                            x = torch.nn.functional.interpolate(x, skip.shape[-2:], mode='bilinear',
-                                                                align_corners=True)
-                        if cnet is not None:
-                            next_cnet = cnet.pop()
-                            if next_cnet is not None:
-                                x = x + nn.functional.interpolate(next_cnet, size=x.shape[-2:], mode='bilinear',
-                                                                  align_corners=True).to(x.dtype)
-                        x = block(x, skip)
-                    elif isinstance(block, AttnBlock) or (
-                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
-                                                                                  AttnBlock)):
-                        x = block(x, clip)
-                    elif isinstance(block, TimestepBlock) or (
-                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
-                                                                                  TimestepBlock)):
-                        x = block(x, r_embed)
-                    else:
-                        x = block(x)
-                if j < len(repmap):
-                    x = repmap[j](x)
-            x = upscaler(x)
-        return x
-
-    def forward(self, x, r, clip_text, clip_text_pooled, clip_img, control=None, **kwargs):
-        # Process the conditioning embeddings
-        r_embed = self.gen_r_embedding(r).to(dtype=x.dtype)
-        for c in self.t_conds:
-            t_cond = kwargs.get(c, torch.zeros_like(r))
-            r_embed = torch.cat([r_embed, self.gen_r_embedding(t_cond).to(dtype=x.dtype)], dim=1)
-        clip = self.gen_c_embeddings(clip_text, clip_text_pooled, clip_img)
-
-        if control is not None:
-            cnet = control.get("input")
-        else:
-            cnet = None
-
-        # Model Blocks
-        x = self.embedding(x)
-        level_outputs = self._down_encode(x, r_embed, clip, cnet)
-        x = self._up_decode(level_outputs, r_embed, clip, cnet)
-        return self.clf(x)
-
-    def update_weights_ema(self, src_model, beta=0.999):
-        for self_params, src_params in zip(self.parameters(), src_model.parameters()):
-            self_params.data = self_params.data * beta + src_params.data.clone().to(self_params.device) * (1 - beta)
-        for self_buffers, src_buffers in zip(self.buffers(), src_model.buffers()):
-            self_buffers.data = self_buffers.data * beta + src_buffers.data.clone().to(self_buffers.device) * (1 - beta)

ldm_patched/ldm/cascade/stage_c_coder.py

@@ -1,95 +0,0 @@
-"""
-    This file is part of ComfyUI.
-    Copyright (C) 2024 Stability AI
-
-    This program is free software: you can redistribute it and/or modify
-    it under the terms of the GNU General Public License as published by
-    the Free Software Foundation, either version 3 of the License, or
-    (at your option) any later version.
-
-    This program is distributed in the hope that it will be useful,
-    but WITHOUT ANY WARRANTY; without even the implied warranty of
-    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-    GNU General Public License for more details.
-
-    You should have received a copy of the GNU General Public License
-    along with this program.  If not, see <https://www.gnu.org/licenses/>.
-"""
-import torch
-import torchvision
-from torch import nn
-
-
-# EfficientNet
-class EfficientNetEncoder(nn.Module):
-    def __init__(self, c_latent=16):
-        super().__init__()
-        self.backbone = torchvision.models.efficientnet_v2_s().features.eval()
-        self.mapper = nn.Sequential(
-            nn.Conv2d(1280, c_latent, kernel_size=1, bias=False),
-            nn.BatchNorm2d(c_latent, affine=False),  # then normalize them to have mean 0 and std 1
-        )
-        self.mean = nn.Parameter(torch.tensor([0.485, 0.456, 0.406]))
-        self.std = nn.Parameter(torch.tensor([0.229, 0.224, 0.225]))
-
-    def forward(self, x):
-        x = x * 0.5 + 0.5
-        x = (x - self.mean.view([3,1,1])) / self.std.view([3,1,1])
-        o = self.mapper(self.backbone(x))
-        return o
-
-
-# Fast Decoder for Stage C latents. E.g. 16 x 24 x 24 -> 3 x 192 x 192
-class Previewer(nn.Module):
-    def __init__(self, c_in=16, c_hidden=512, c_out=3):
-        super().__init__()
-        self.blocks = nn.Sequential(
-            nn.Conv2d(c_in, c_hidden, kernel_size=1),  # 16 channels to 512 channels
-            nn.GELU(),
-            nn.BatchNorm2d(c_hidden),
-
-            nn.Conv2d(c_hidden, c_hidden, kernel_size=3, padding=1),
-            nn.GELU(),
-            nn.BatchNorm2d(c_hidden),
-
-            nn.ConvTranspose2d(c_hidden, c_hidden // 2, kernel_size=2, stride=2),  # 16 -> 32
-            nn.GELU(),
-            nn.BatchNorm2d(c_hidden // 2),
-
-            nn.Conv2d(c_hidden // 2, c_hidden // 2, kernel_size=3, padding=1),
-            nn.GELU(),
-            nn.BatchNorm2d(c_hidden // 2),
-
-            nn.ConvTranspose2d(c_hidden // 2, c_hidden // 4, kernel_size=2, stride=2),  # 32 -> 64
-            nn.GELU(),
-            nn.BatchNorm2d(c_hidden // 4),
-
-            nn.Conv2d(c_hidden // 4, c_hidden // 4, kernel_size=3, padding=1),
-            nn.GELU(),
-            nn.BatchNorm2d(c_hidden // 4),
-
-            nn.ConvTranspose2d(c_hidden // 4, c_hidden // 4, kernel_size=2, stride=2),  # 64 -> 128
-            nn.GELU(),
-            nn.BatchNorm2d(c_hidden // 4),
-
-            nn.Conv2d(c_hidden // 4, c_hidden // 4, kernel_size=3, padding=1),
-            nn.GELU(),
-            nn.BatchNorm2d(c_hidden // 4),
-
-            nn.Conv2d(c_hidden // 4, c_out, kernel_size=1),
-        )
-
-    def forward(self, x):
-        return (self.blocks(x) - 0.5) * 2.0
-
-class StageC_coder(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.previewer = Previewer()
-        self.encoder = EfficientNetEncoder()
-
-    def encode(self, x):
-        return self.encoder(x)
-
-    def decode(self, x):
-        return self.previewer(x)

ldm_patched/ldm/models/autoencoder.py

@@ -1,4 +1,6 @@
 import torch
+# import pytorch_lightning as pl
+import torch.nn.functional as F
 from contextlib import contextmanager
 from typing import Any, Dict, List, Optional, Tuple, Union
 

ldm_patched/ldm/modules/attention.py

@@ -3,10 +3,9 @@ import torch
 import torch.nn.functional as F
 from torch import nn, einsum
 from einops import rearrange, repeat
-from typing import Optional
-import logging
+from typing import Optional, Any
 
-from .diffusionmodules.util import AlphaBlender, timestep_embedding
+from .diffusionmodules.util import checkpoint, AlphaBlender, timestep_embedding
 from .sub_quadratic_attention import efficient_dot_product_attention
 
 from ldm_patched.modules import model_management
@@ -19,14 +18,13 @@ from ldm_patched.modules.args_parser import args
 import ldm_patched.modules.ops
 ops = ldm_patched.modules.ops.disable_weight_init
 
-FORCE_UPCAST_ATTENTION_DTYPE = model_management.force_upcast_attention_dtype()
+# CrossAttn precision handling
+if args.disable_attention_upcast:
+    print("disabling upcasting of attention")
+    _ATTN_PRECISION = "fp16"
+else:
+    _ATTN_PRECISION = "fp32"
 
-def get_attn_precision(attn_precision):
-    if args.dont_upcast_attention:
-        return None
-    if FORCE_UPCAST_ATTENTION_DTYPE is not None:
-        return FORCE_UPCAST_ATTENTION_DTYPE
-    return attn_precision
 
 def exists(val):
     return val is not None
@@ -86,9 +84,7 @@ class FeedForward(nn.Module):
 def Normalize(in_channels, dtype=None, device=None):
     return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, dtype=dtype, device=device)
 
-def attention_basic(q, k, v, heads, mask=None, attn_precision=None):
-    attn_precision = get_attn_precision(attn_precision)
-
+def attention_basic(q, k, v, heads, mask=None):
     b, _, dim_head = q.shape
     dim_head //= heads
     scale = dim_head ** -0.5
@@ -104,7 +100,7 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None):
     )
 
     # force cast to fp32 to avoid overflowing
-    if attn_precision == torch.float32:
+    if _ATTN_PRECISION =="fp32":
         sim = einsum('b i d, b j d -> b i j', q.float(), k.float()) * scale
     else:
         sim = einsum('b i d, b j d -> b i j', q, k) * scale
@@ -118,12 +114,7 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None):
             mask = repeat(mask, 'b j -> (b h) () j', h=h)
             sim.masked_fill_(~mask, max_neg_value)
         else:
-            if len(mask.shape) == 2:
-                bs = 1
-            else:
-                bs = mask.shape[0]
-            mask = mask.reshape(bs, -1, mask.shape[-2], mask.shape[-1]).expand(b, heads, -1, -1).reshape(-1, mask.shape[-2], mask.shape[-1])
-            sim.add_(mask)
+            sim += mask
 
     # attention, what we cannot get enough of
     sim = sim.softmax(dim=-1)
@@ -138,9 +129,7 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None):
     return out
 
 
-def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None):
-    attn_precision = get_attn_precision(attn_precision)
-
+def attention_sub_quad(query, key, value, heads, mask=None):
     b, _, dim_head = query.shape
     dim_head //= heads
 
@@ -151,7 +140,7 @@ def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None)
     key = key.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 3, 1).reshape(b * heads, dim_head, -1)
 
     dtype = query.dtype
-    upcast_attention = attn_precision == torch.float32 and query.dtype != torch.float32
+    upcast_attention = _ATTN_PRECISION =="fp32" and query.dtype != torch.float32
     if upcast_attention:
         bytes_per_token = torch.finfo(torch.float32).bits//8
     else:
@@ -176,13 +165,6 @@ def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None)
     if query_chunk_size is None:
         query_chunk_size = 512
 
-    if mask is not None:
-        if len(mask.shape) == 2:
-            bs = 1
-        else:
-            bs = mask.shape[0]
-        mask = mask.reshape(bs, -1, mask.shape[-2], mask.shape[-1]).expand(b, heads, -1, -1).reshape(-1, mask.shape[-2], mask.shape[-1])
-
     hidden_states = efficient_dot_product_attention(
         query,
         key,
@@ -200,9 +182,7 @@ def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None)
     hidden_states = hidden_states.unflatten(0, (-1, heads)).transpose(1,2).flatten(start_dim=2)
     return hidden_states
 
-def attention_split(q, k, v, heads, mask=None, attn_precision=None):
-    attn_precision = get_attn_precision(attn_precision)
-
+def attention_split(q, k, v, heads, mask=None):
     b, _, dim_head = q.shape
     dim_head //= heads
     scale = dim_head ** -0.5
@@ -221,12 +201,10 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None):
 
     mem_free_total = model_management.get_free_memory(q.device)
 
-    if attn_precision == torch.float32:
+    if _ATTN_PRECISION =="fp32":
         element_size = 4
-        upcast = True
     else:
         element_size = q.element_size()
-        upcast = False
 
     gb = 1024 ** 3
     tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * element_size
@@ -245,13 +223,6 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None):
         raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '
                             f'Need: {mem_required/64/gb:0.1f}GB free, Have:{mem_free_total/gb:0.1f}GB free')
 
-    if mask is not None:
-        if len(mask.shape) == 2:
-            bs = 1
-        else:
-            bs = mask.shape[0]
-        mask = mask.reshape(bs, -1, mask.shape[-2], mask.shape[-1]).expand(b, heads, -1, -1).reshape(-1, mask.shape[-2], mask.shape[-1])
-
     # print("steps", steps, mem_required, mem_free_total, modifier, q.element_size(), tensor_size)
     first_op_done = False
     cleared_cache = False
@@ -260,7 +231,7 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None):
             slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
             for i in range(0, q.shape[1], slice_size):
                 end = i + slice_size
-                if upcast:
+                if _ATTN_PRECISION =="fp32":
                     with torch.autocast(enabled=False, device_type = 'cuda'):
                         s1 = einsum('b i d, b j d -> b i j', q[:, i:end].float(), k.float()) * scale
                 else:
@@ -284,12 +255,12 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None):
                 model_management.soft_empty_cache(True)
                 if cleared_cache == False:
                     cleared_cache = True
-                    logging.warning("out of memory error, emptying cache and trying again")
+                    print("out of memory error, emptying cache and trying again")
                     continue
                 steps *= 2
                 if steps > 64:
                     raise e
-                logging.warning("out of memory error, increasing steps and trying again {}".format(steps))
+                print("out of memory error, increasing steps and trying again", steps)
             else:
                 raise e
 
@@ -306,30 +277,24 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None):
 BROKEN_XFORMERS = False
 try:
     x_vers = xformers.__version__
-    # XFormers bug confirmed on all versions from 0.0.21 to 0.0.26 (q with bs bigger than 65535 gives CUDA error)
-    BROKEN_XFORMERS = x_vers.startswith("0.0.2") and not x_vers.startswith("0.0.20")
+    #I think 0.0.23 is also broken (q with bs bigger than 65535 gives CUDA error)
+    BROKEN_XFORMERS = x_vers.startswith("0.0.21") or x_vers.startswith("0.0.22") or x_vers.startswith("0.0.23")
 except:
     pass
 
-def attention_xformers(q, k, v, heads, mask=None, attn_precision=None):
+def attention_xformers(q, k, v, heads, mask=None):
     b, _, dim_head = q.shape
     dim_head //= heads
-
-    disabled_xformers = False
-
     if BROKEN_XFORMERS:
         if b * heads > 65535:
-            disabled_xformers = True
-
-    if not disabled_xformers:
-        if torch.jit.is_tracing() or torch.jit.is_scripting():
-            disabled_xformers = True
-
-    if disabled_xformers:
-        return attention_pytorch(q, k, v, heads, mask)
+            return attention_pytorch(q, k, v, heads, mask)
 
     q, k, v = map(
-        lambda t: t.reshape(b, -1, heads, dim_head),
+        lambda t: t.unsqueeze(3)
+        .reshape(b, -1, heads, dim_head)
+        .permute(0, 2, 1, 3)
+        .reshape(b * heads, -1, dim_head)
+        .contiguous(),
         (q, k, v),
     )
 
@@ -342,11 +307,14 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None):
     out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask)
 
     out = (
-        out.reshape(b, -1, heads * dim_head)
+        out.unsqueeze(0)
+        .reshape(b, heads, -1, dim_head)
+        .permute(0, 2, 1, 3)
+        .reshape(b, -1, heads * dim_head)
     )
     return out
 
-def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None):
+def attention_pytorch(q, k, v, heads, mask=None):
     b, _, dim_head = q.shape
     dim_head //= heads
     q, k, v = map(
@@ -364,17 +332,17 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None):
 optimized_attention = attention_basic
 
 if model_management.xformers_enabled():
-    logging.info("Using xformers cross attention")
+    print("Using xformers cross attention")
     optimized_attention = attention_xformers
 elif model_management.pytorch_attention_enabled():
-    logging.info("Using pytorch cross attention")
+    print("Using pytorch cross attention")
     optimized_attention = attention_pytorch
 else:
-    if args.use_split_cross_attention:
-        logging.info("Using split optimization for cross attention")
+    if args.attention_split:
+        print("Using split optimization for cross attention")
         optimized_attention = attention_split
     else:
-        logging.info("Using sub quadratic optimization for cross attention, if you have memory or speed issues try using: --use-split-cross-attention")
+        print("Using sub quadratic optimization for cross attention, if you have memory or speed issues try using: --attention-split")
         optimized_attention = attention_sub_quad
 
 optimized_attention_masked = optimized_attention
@@ -396,11 +364,10 @@ def optimized_attention_for_device(device, mask=False, small_input=False):
 
 
 class CrossAttention(nn.Module):
-    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., attn_precision=None, dtype=None, device=None, operations=ops):
+    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None, operations=ops):
         super().__init__()
         inner_dim = dim_head * heads
         context_dim = default(context_dim, query_dim)
-        self.attn_precision = attn_precision
 
         self.heads = heads
         self.dim_head = dim_head
@@ -422,15 +389,15 @@ class CrossAttention(nn.Module):
             v = self.to_v(context)
 
         if mask is None:
-            out = optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision)
+            out = optimized_attention(q, k, v, self.heads)
         else:
-            out = optimized_attention_masked(q, k, v, self.heads, mask, attn_precision=self.attn_precision)
+            out = optimized_attention_masked(q, k, v, self.heads, mask)
         return self.to_out(out)
 
 
 class BasicTransformerBlock(nn.Module):
     def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True, ff_in=False, inner_dim=None,
-                 disable_self_attn=False, disable_temporal_crossattention=False, switch_temporal_ca_to_sa=False, attn_precision=None, dtype=None, device=None, operations=ops):
+                 disable_self_attn=False, disable_temporal_crossattention=False, switch_temporal_ca_to_sa=False, dtype=None, device=None, operations=ops):
         super().__init__()
 
         self.ff_in = ff_in or inner_dim is not None
@@ -438,7 +405,6 @@ class BasicTransformerBlock(nn.Module):
             inner_dim = dim
 
         self.is_res = inner_dim == dim
-        self.attn_precision = attn_precision
 
         if self.ff_in:
             self.norm_in = operations.LayerNorm(dim, dtype=dtype, device=device)
@@ -446,7 +412,7 @@ class BasicTransformerBlock(nn.Module):
 
         self.disable_self_attn = disable_self_attn
         self.attn1 = CrossAttention(query_dim=inner_dim, heads=n_heads, dim_head=d_head, dropout=dropout,
-                              context_dim=context_dim if self.disable_self_attn else None, attn_precision=self.attn_precision, dtype=dtype, device=device, operations=operations)  # is a self-attention if not self.disable_self_attn
+                              context_dim=context_dim if self.disable_self_attn else None, dtype=dtype, device=device, operations=operations)  # is a self-attention if not self.disable_self_attn
         self.ff = FeedForward(inner_dim, dim_out=dim, dropout=dropout, glu=gated_ff, dtype=dtype, device=device, operations=operations)
 
         if disable_temporal_crossattention:
@@ -460,16 +426,20 @@ class BasicTransformerBlock(nn.Module):
                 context_dim_attn2 = context_dim
 
             self.attn2 = CrossAttention(query_dim=inner_dim, context_dim=context_dim_attn2,
-                                heads=n_heads, dim_head=d_head, dropout=dropout, attn_precision=self.attn_precision, dtype=dtype, device=device, operations=operations)  # is self-attn if context is none
+                                heads=n_heads, dim_head=d_head, dropout=dropout, dtype=dtype, device=device, operations=operations)  # is self-attn if context is none
             self.norm2 = operations.LayerNorm(inner_dim, dtype=dtype, device=device)
 
         self.norm1 = operations.LayerNorm(inner_dim, dtype=dtype, device=device)
         self.norm3 = operations.LayerNorm(inner_dim, dtype=dtype, device=device)
+        self.checkpoint = checkpoint
         self.n_heads = n_heads
         self.d_head = d_head
         self.switch_temporal_ca_to_sa = switch_temporal_ca_to_sa
 
     def forward(self, x, context=None, transformer_options={}):
+        return checkpoint(self._forward, (x, context, transformer_options), self.parameters(), self.checkpoint)
+
+    def _forward(self, x, context=None, transformer_options={}):
         extra_options = {}
         block = transformer_options.get("block", None)
         block_index = transformer_options.get("block_index", 0)
@@ -486,7 +456,6 @@ class BasicTransformerBlock(nn.Module):
 
         extra_options["n_heads"] = self.n_heads
         extra_options["dim_head"] = self.d_head
-        extra_options["attn_precision"] = self.attn_precision
 
         if self.ff_in:
             x_skip = x
@@ -597,7 +566,7 @@ class SpatialTransformer(nn.Module):
     def __init__(self, in_channels, n_heads, d_head,
                  depth=1, dropout=0., context_dim=None,
                  disable_self_attn=False, use_linear=False,
-                 use_checkpoint=True, attn_precision=None, dtype=None, device=None, operations=ops):
+                 use_checkpoint=True, dtype=None, device=None, operations=ops):
         super().__init__()
         if exists(context_dim) and not isinstance(context_dim, list):
             context_dim = [context_dim] * depth
@@ -615,7 +584,7 @@ class SpatialTransformer(nn.Module):
 
         self.transformer_blocks = nn.ModuleList(
             [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim[d],
-                                   disable_self_attn=disable_self_attn, checkpoint=use_checkpoint, attn_precision=attn_precision, dtype=dtype, device=device, operations=operations)
+                                   disable_self_attn=disable_self_attn, checkpoint=use_checkpoint, dtype=dtype, device=device, operations=operations)
                 for d in range(depth)]
         )
         if not use_linear:
@@ -636,7 +605,7 @@ class SpatialTransformer(nn.Module):
         x = self.norm(x)
         if not self.use_linear:
             x = self.proj_in(x)
-        x = x.movedim(1, 3).flatten(1, 2).contiguous()
+        x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
         if self.use_linear:
             x = self.proj_in(x)
         for i, block in enumerate(self.transformer_blocks):
@@ -644,7 +613,7 @@ class SpatialTransformer(nn.Module):
             x = block(x, context=context[i], transformer_options=transformer_options)
         if self.use_linear:
             x = self.proj_out(x)
-        x = x.reshape(x.shape[0], h, w, x.shape[-1]).movedim(3, 1).contiguous()
+        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
         if not self.use_linear:
             x = self.proj_out(x)
         return x + x_in
@@ -671,7 +640,6 @@ class SpatialVideoTransformer(SpatialTransformer):
         disable_self_attn=False,
         disable_temporal_crossattention=False,
         max_time_embed_period: int = 10000,
-        attn_precision=None,
         dtype=None, device=None, operations=ops
     ):
         super().__init__(
@@ -684,7 +652,6 @@ class SpatialVideoTransformer(SpatialTransformer):
             context_dim=context_dim,
             use_linear=use_linear,
             disable_self_attn=disable_self_attn,
-            attn_precision=attn_precision,
             dtype=dtype, device=device, operations=operations
         )
         self.time_depth = time_depth
@@ -714,7 +681,6 @@ class SpatialVideoTransformer(SpatialTransformer):
                     inner_dim=time_mix_inner_dim,
                     disable_self_attn=disable_self_attn,
                     disable_temporal_crossattention=disable_temporal_crossattention,
-                    attn_precision=attn_precision,
                     dtype=dtype, device=device, operations=operations
                 )
                 for _ in range(self.depth)

ldm_patched/ldm/modules/diffusionmodules/model.py

@@ -3,8 +3,8 @@ import math
 import torch
 import torch.nn as nn
 import numpy as np
+from einops import rearrange
 from typing import Optional, Any
-import logging
 
 from ldm_patched.modules import model_management
 import ldm_patched.modules.ops
@@ -190,7 +190,7 @@ def slice_attention(q, k, v):
             steps *= 2
             if steps > 128:
                 raise e
-            logging.warning("out of memory error, increasing steps and trying again {}".format(steps))
+            print("out of memory error, increasing steps and trying again", steps)
 
     return r1
 
@@ -235,7 +235,7 @@ def pytorch_attention(q, k, v):
         out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
         out = out.transpose(2, 3).reshape(B, C, H, W)
     except model_management.OOM_EXCEPTION as e:
-        logging.warning("scaled_dot_product_attention OOMed: switched to slice attention")
+        print("scaled_dot_product_attention OOMed: switched to slice attention")
         out = slice_attention(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2), v.view(B, -1, C).transpose(1, 2)).reshape(B, C, H, W)
     return out
 
@@ -268,13 +268,13 @@ class AttnBlock(nn.Module):
                                         padding=0)
 
         if model_management.xformers_enabled_vae():
-            logging.info("Using xformers attention in VAE")
+            print("Using xformers attention in VAE")
             self.optimized_attention = xformers_attention
         elif model_management.pytorch_attention_enabled():
-            logging.info("Using pytorch attention in VAE")
+            print("Using pytorch attention in VAE")
             self.optimized_attention = pytorch_attention
         else:
-            logging.info("Using split attention in VAE")
+            print("Using split attention in VAE")
             self.optimized_attention = normal_attention
 
     def forward(self, x):
@@ -562,7 +562,7 @@ class Decoder(nn.Module):
         block_in = ch*ch_mult[self.num_resolutions-1]
         curr_res = resolution // 2**(self.num_resolutions-1)
         self.z_shape = (1,z_channels,curr_res,curr_res)
-        logging.debug("Working with z of shape {} = {} dimensions.".format(
+        print("Working with z of shape {} = {} dimensions.".format(
             self.z_shape, np.prod(self.z_shape)))
 
         # z to block_in

ldm_patched/ldm/modules/diffusionmodules/openaimodel.py

@@ -4,7 +4,6 @@ import torch as th
 import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
-import logging
 
 from .util import (
     checkpoint,
@@ -258,7 +257,7 @@ class ResBlock(TimestepBlock):
         else:
             if emb_out is not None:
                 if self.exchange_temb_dims:
-                    emb_out = emb_out.movedim(1, 2)
+                    emb_out = rearrange(emb_out, "b t c ... -> b c t ...")
                 h = h + emb_out
             h = self.out_layers(h)
         return self.skip_connection(x) + h
@@ -360,7 +359,7 @@ def apply_control(h, control, name):
             try:
                 h += ctrl
             except:
-                logging.warning("warning control could not be applied {} {}".format(h.shape, ctrl.shape))
+                print("warning control could not be applied", h.shape, ctrl.shape)
     return h
 
 class UNetModel(nn.Module):
@@ -431,7 +430,6 @@ class UNetModel(nn.Module):
         video_kernel_size=None,
         disable_temporal_crossattention=False,
         max_ddpm_temb_period=10000,
-        attn_precision=None,
         device=None,
         operations=ops,
     ):
@@ -486,6 +484,7 @@ class UNetModel(nn.Module):
         self.predict_codebook_ids = n_embed is not None
 
         self.default_num_video_frames = None
+        self.default_image_only_indicator = None
 
         time_embed_dim = model_channels * 4
         self.time_embed = nn.Sequential(
@@ -498,7 +497,7 @@ class UNetModel(nn.Module):
             if isinstance(self.num_classes, int):
                 self.label_emb = nn.Embedding(num_classes, time_embed_dim, dtype=self.dtype, device=device)
             elif self.num_classes == "continuous":
-                logging.debug("setting up linear c_adm embedding layer")
+                print("setting up linear c_adm embedding layer")
                 self.label_emb = nn.Linear(1, time_embed_dim)
             elif self.num_classes == "sequential":
                 assert adm_in_channels is not None
@@ -551,14 +550,13 @@ class UNetModel(nn.Module):
                     disable_self_attn=disable_self_attn,
                     disable_temporal_crossattention=disable_temporal_crossattention,
                     max_time_embed_period=max_ddpm_temb_period,
-                    attn_precision=attn_precision,
                     dtype=self.dtype, device=device, operations=operations
                 )
             else:
                 return SpatialTransformer(
                                 ch, num_heads, dim_head, depth=depth, context_dim=context_dim,
                                 disable_self_attn=disable_self_attn, use_linear=use_linear_in_transformer,
-                                use_checkpoint=use_checkpoint, attn_precision=attn_precision, dtype=self.dtype, device=device, operations=operations
+                                use_checkpoint=use_checkpoint, dtype=self.dtype, device=device, operations=operations
                             )
 
         def get_resblock(
@@ -710,30 +708,27 @@ class UNetModel(nn.Module):
                 device=device,
                 operations=operations
             )]
-
-        self.middle_block = None
-        if transformer_depth_middle >= -1:
-            if transformer_depth_middle >= 0:
-                mid_block += [get_attention_layer(  # always uses a self-attn
-                                ch, num_heads, dim_head, depth=transformer_depth_middle, context_dim=context_dim,
-                                disable_self_attn=disable_middle_self_attn, use_checkpoint=use_checkpoint
-                            ),
-                get_resblock(
-                    merge_factor=merge_factor,
-                    merge_strategy=merge_strategy,
-                    video_kernel_size=video_kernel_size,
-                    ch=ch,
-                    time_embed_dim=time_embed_dim,
-                    dropout=dropout,
-                    out_channels=None,
-                    dims=dims,
-                    use_checkpoint=use_checkpoint,
-                    use_scale_shift_norm=use_scale_shift_norm,
-                    dtype=self.dtype,
-                    device=device,
-                    operations=operations
-                )]
-            self.middle_block = TimestepEmbedSequential(*mid_block)
+        if transformer_depth_middle >= 0:
+            mid_block += [get_attention_layer(  # always uses a self-attn
+                            ch, num_heads, dim_head, depth=transformer_depth_middle, context_dim=context_dim,
+                            disable_self_attn=disable_middle_self_attn, use_checkpoint=use_checkpoint
+                        ),
+            get_resblock(
+                merge_factor=merge_factor,
+                merge_strategy=merge_strategy,
+                video_kernel_size=video_kernel_size,
+                ch=ch,
+                time_embed_dim=time_embed_dim,
+                dropout=dropout,
+                out_channels=None,
+                dims=dims,
+                use_checkpoint=use_checkpoint,
+                use_scale_shift_norm=use_scale_shift_norm,
+                dtype=self.dtype,
+                device=device,
+                operations=operations
+            )]
+        self.middle_block = TimestepEmbedSequential(*mid_block)
         self._feature_size += ch
 
         self.output_blocks = nn.ModuleList([])
@@ -832,7 +827,7 @@ class UNetModel(nn.Module):
         transformer_patches = transformer_options.get("patches", {})
 
         num_video_frames = kwargs.get("num_video_frames", self.default_num_video_frames)
-        image_only_indicator = kwargs.get("image_only_indicator", None)
+        image_only_indicator = kwargs.get("image_only_indicator", self.default_image_only_indicator)
         time_context = kwargs.get("time_context", None)
 
         assert (y is not None) == (
@@ -863,8 +858,7 @@ class UNetModel(nn.Module):
                     h = p(h, transformer_options)
 
         transformer_options["block"] = ("middle", 0)
-        if self.middle_block is not None:
-            h = forward_timestep_embed(self.middle_block, h, emb, context, transformer_options, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
+        h = forward_timestep_embed(self.middle_block, h, emb, context, transformer_options, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
         h = apply_control(h, control, 'middle')
 
 

ldm_patched/ldm/modules/diffusionmodules/util.py

@@ -46,25 +46,23 @@ class AlphaBlender(nn.Module):
         else:
             raise ValueError(f"unknown merge strategy {self.merge_strategy}")
 
-    def get_alpha(self, image_only_indicator: torch.Tensor, device) -> torch.Tensor:
+    def get_alpha(self, image_only_indicator: torch.Tensor) -> torch.Tensor:
         # skip_time_mix = rearrange(repeat(skip_time_mix, 'b -> (b t) () () ()', t=t), '(b t) 1 ... -> b 1 t ...', t=t)
         if self.merge_strategy == "fixed":
             # make shape compatible
             # alpha = repeat(self.mix_factor, '1 -> b () t  () ()', t=t, b=bs)
-            alpha = self.mix_factor.to(device)
+            alpha = self.mix_factor.to(image_only_indicator.device)
         elif self.merge_strategy == "learned":
-            alpha = torch.sigmoid(self.mix_factor.to(device))
+            alpha = torch.sigmoid(self.mix_factor.to(image_only_indicator.device))
             # make shape compatible
             # alpha = repeat(alpha, '1 -> s () ()', s = t * bs)
         elif self.merge_strategy == "learned_with_images":
-            if image_only_indicator is None:
-                alpha = rearrange(torch.sigmoid(self.mix_factor.to(device)), "... -> ... 1")
-            else:
-                alpha = torch.where(
-                    image_only_indicator.bool(),
-                    torch.ones(1, 1, device=image_only_indicator.device),
-                    rearrange(torch.sigmoid(self.mix_factor.to(image_only_indicator.device)), "... -> ... 1"),
-                )
+            assert image_only_indicator is not None, "need image_only_indicator ..."
+            alpha = torch.where(
+                image_only_indicator.bool(),
+                torch.ones(1, 1, device=image_only_indicator.device),
+                rearrange(torch.sigmoid(self.mix_factor.to(image_only_indicator.device)), "... -> ... 1"),
+            )
             alpha = rearrange(alpha, self.rearrange_pattern)
             # make shape compatible
             # alpha = repeat(alpha, '1 -> s () ()', s = t * bs)
@@ -78,7 +76,7 @@ class AlphaBlender(nn.Module):
         x_temporal,
         image_only_indicator=None,
     ) -> torch.Tensor:
-        alpha = self.get_alpha(image_only_indicator, x_spatial.device)
+        alpha = self.get_alpha(image_only_indicator)
         x = (
             alpha.to(x_spatial.dtype) * x_spatial
             + (1.0 - alpha).to(x_spatial.dtype) * x_temporal
@@ -100,7 +98,7 @@ def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2,
         alphas = torch.cos(alphas).pow(2)
         alphas = alphas / alphas[0]
         betas = 1 - alphas[1:] / alphas[:-1]
-        betas = torch.clamp(betas, min=0, max=0.999)
+        betas = np.clip(betas, a_min=0, a_max=0.999)
 
     elif schedule == "squaredcos_cap_v2":  # used for karlo prior
         # return early
@@ -115,7 +113,7 @@ def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2,
         betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5
     else:
         raise ValueError(f"schedule '{schedule}' unknown.")
-    return betas
+    return betas.numpy()
 
 
 def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True):

ldm_patched/ldm/modules/encoders/noise_aug_modules.py

@@ -8,7 +8,7 @@ class CLIPEmbeddingNoiseAugmentation(ImageConcatWithNoiseAugmentation):
         if clip_stats_path is None:
             clip_mean, clip_std = torch.zeros(timestep_dim), torch.ones(timestep_dim)
         else:
-            clip_mean, clip_std = torch.load(clip_stats_path, map_location="cpu")
+            clip_mean, clip_std = torch.load(clip_stats_path, map_location="cpu", weights_only=True)
         self.register_buffer("data_mean", clip_mean[None, :], persistent=False)
         self.register_buffer("data_std", clip_std[None, :], persistent=False)
         self.time_embed = Timestep(timestep_dim)

ldm_patched/ldm/modules/sub_quadratic_attention.py

@@ -14,7 +14,6 @@ import torch
 from torch import Tensor
 from torch.utils.checkpoint import checkpoint
 import math
-import logging
 
 try:
 	from typing import Optional, NamedTuple, List, Protocol
@@ -171,7 +170,7 @@ def _get_attention_scores_no_kv_chunking(
         attn_probs = attn_scores.softmax(dim=-1)
         del attn_scores
     except model_management.OOM_EXCEPTION:
-        logging.warning("ran out of memory while running softmax in  _get_attention_scores_no_kv_chunking, trying slower in place softmax instead")
+        print("ran out of memory while running softmax in  _get_attention_scores_no_kv_chunking, trying slower in place softmax instead")
         attn_scores -= attn_scores.max(dim=-1, keepdim=True).values
         torch.exp(attn_scores, out=attn_scores)
         summed = torch.sum(attn_scores, dim=-1, keepdim=True)

ldm_patched/modules/args_parser.py

@@ -33,105 +33,94 @@ class EnumAction(argparse.Action):
 
 parser = argparse.ArgumentParser()
 
-parser.add_argument("--listen", type=str, default="127.0.0.1", metavar="IP", nargs="?", const="0.0.0.0", help="Specify the IP address to listen on (default: 127.0.0.1). If --listen is provided without an argument, it defaults to 0.0.0.0. (listens on all)")
-parser.add_argument("--port", type=int, default=8188, help="Set the listen port.")
-parser.add_argument("--tls-keyfile", type=str, help="Path to TLS (SSL) key file. Enables TLS, makes app accessible at https://... requires --tls-certfile to function")
-parser.add_argument("--tls-keyfile-password", type=str, help="Password for the TLS (SSL) key file. Requires --tls-keyfile to function")
-parser.add_argument("--tls-certfile", type=str, help="Path to TLS (SSL) certificate file. Enables TLS, makes app accessible at https://... requires --tls-keyfile to function")
-parser.add_argument("--tls-verify", type=bool, default=True, help="If False, skips certificate validation which allows self-signed certificates to be used. Requires --tls-keyfile to function")
-parser.add_argument("--hf-mirror", type=str, default=None, help="Use HuggingFace mirror URL")
+parser.add_argument("--listen", type=str, default="127.0.0.1", metavar="IP", nargs="?", const="0.0.0.0")
+parser.add_argument("--port", type=int, default=8188)
+parser.add_argument("--disable-header-check", type=str, default=None, metavar="ORIGIN", nargs="?", const="*")
+parser.add_argument("--web-upload-size", type=float, default=100)
+parser.add_argument("--hf-mirror", type=str, default=None)
 
-parser.add_argument("--output-path", type=str, default=None, help="Set the output directory.")
-parser.add_argument("--temp-path", type=str, default=None, help="Set the temp directory.")
-parser.add_argument("--in-browser", action="store_true", help="Automatically launch in the default browser.")
-parser.add_argument("--disable-in-browser", action="store_true", help="Disable auto launching the browser.")
-parser.add_argument("--gpu-device-id", type=int, default=None, metavar="DEVICE_ID", help="Set the id of the cuda device this instance will use.")
+parser.add_argument("--external-working-path", type=str, default=None, metavar="PATH", nargs='+', action='append')
+parser.add_argument("--output-path", type=str, default=None)
+parser.add_argument("--temp-path", type=str, default=None)
+parser.add_argument("--cache-path", type=str, default=None)
+parser.add_argument("--in-browser", action="store_true")
+parser.add_argument("--disable-in-browser", action="store_true")
+parser.add_argument("--gpu-device-id", type=int, default=None, metavar="DEVICE_ID")
+cm_group = parser.add_mutually_exclusive_group()
+cm_group.add_argument("--async-cuda-allocation", action="store_true")
+cm_group.add_argument("--disable-async-cuda-allocation", action="store_true")
+
+parser.add_argument("--disable-attention-upcast", action="store_true")
 
 fp_group = parser.add_mutually_exclusive_group()
-fp_group.add_argument("--force-fp32", action="store_true", help="Force fp32 (If this makes your GPU work better please report it).")
-fp_group.add_argument("--force-fp16", action="store_true", help="Force fp16.")
+fp_group.add_argument("--all-in-fp32", action="store_true")
+fp_group.add_argument("--all-in-fp16", action="store_true")
 
 fpunet_group = parser.add_mutually_exclusive_group()
-fpunet_group.add_argument("--bf16-unet", action="store_true", help="Run the UNET in bf16. This should only be used for testing stuff.")
-fpunet_group.add_argument("--fp16-unet", action="store_true", help="Store unet weights in fp16.")
-fpunet_group.add_argument("--fp8_e4m3fn-unet", action="store_true", help="Store unet weights in fp8_e4m3fn.")
-fpunet_group.add_argument("--fp8_e5m2-unet", action="store_true", help="Store unet weights in fp8_e5m2.")
+fpunet_group.add_argument("--unet-in-bf16", action="store_true")
+fpunet_group.add_argument("--unet-in-fp16", action="store_true")
+fpunet_group.add_argument("--unet-in-fp8-e4m3fn", action="store_true")
+fpunet_group.add_argument("--unet-in-fp8-e5m2", action="store_true")
 
 fpvae_group = parser.add_mutually_exclusive_group()
-fpvae_group.add_argument("--fp16-vae", action="store_true", help="Run the VAE in fp16, might cause black images.")
-fpvae_group.add_argument("--fp32-vae", action="store_true", help="Run the VAE in full precision fp32.")
-fpvae_group.add_argument("--bf16-vae", action="store_true", help="Run the VAE in bf16.")
+fpvae_group.add_argument("--vae-in-fp16", action="store_true")
+fpvae_group.add_argument("--vae-in-fp32", action="store_true")
+fpvae_group.add_argument("--vae-in-bf16", action="store_true")
 
-parser.add_argument("--cpu-vae", action="store_true", help="Run the VAE on the CPU.")
+parser.add_argument("--vae-in-cpu", action="store_true")
 
 fpte_group = parser.add_mutually_exclusive_group()
-fpte_group.add_argument("--fp8_e4m3fn-text-enc", action="store_true", help="Store text encoder weights in fp8 (e4m3fn variant).")
-fpte_group.add_argument("--fp8_e5m2-text-enc", action="store_true", help="Store text encoder weights in fp8 (e5m2 variant).")
-fpte_group.add_argument("--fp16-text-enc", action="store_true", help="Store text encoder weights in fp16.")
-fpte_group.add_argument("--fp32-text-enc", action="store_true", help="Store text encoder weights in fp32.")
+fpte_group.add_argument("--clip-in-fp8-e4m3fn", action="store_true")
+fpte_group.add_argument("--clip-in-fp8-e5m2", action="store_true")
+fpte_group.add_argument("--clip-in-fp16", action="store_true")
+fpte_group.add_argument("--clip-in-fp32", action="store_true")
 
 
-parser.add_argument("--directml", type=int, nargs="?", metavar="DIRECTML_DEVICE", const=-1, help="Use torch-directml.")
+parser.add_argument("--directml", type=int, nargs="?", metavar="DIRECTML_DEVICE", const=-1)
 
-parser.add_argument("--disable-ipex-optimize", action="store_true", help="Disables ipex.optimize when loading models with Intel GPUs.")
+parser.add_argument("--disable-ipex-hijack", action="store_true")
 
 class LatentPreviewMethod(enum.Enum):
     NoPreviews = "none"
     Auto = "auto"
-    Latent2RGB = "latent2rgb"
+    Latent2RGB = "fast"
     TAESD = "taesd"
 
-parser.add_argument("--preview-method", type=LatentPreviewMethod, default=LatentPreviewMethod.NoPreviews, help="Default preview method for sampler nodes.", action=EnumAction)
+parser.add_argument("--preview-option", type=LatentPreviewMethod, default=LatentPreviewMethod.NoPreviews, action=EnumAction)
 
 attn_group = parser.add_mutually_exclusive_group()
-attn_group.add_argument("--use-split-cross-attention", action="store_true", help="Use the split cross attention optimization. Ignored when xformers is used.")
-attn_group.add_argument("--use-quad-cross-attention", action="store_true", help="Use the sub-quadratic cross attention optimization . Ignored when xformers is used.")
-attn_group.add_argument("--use-pytorch-cross-attention", action="store_true", help="Use the new pytorch 2.0 cross attention function.")
-
-parser.add_argument("--disable-xformers", action="store_true", help="Disable xformers.")
-
-upcast = parser.add_mutually_exclusive_group()
-upcast.add_argument("--force-upcast-attention", action="store_true", help="Force enable attention upcasting, please report if it fixes black images.")
-upcast.add_argument("--dont-upcast-attention", action="store_true", help="Disable all upcasting of attention. Should be unnecessary except for debugging.")
+attn_group.add_argument("--attention-split", action="store_true")
+attn_group.add_argument("--attention-quad", action="store_true")
+attn_group.add_argument("--attention-pytorch", action="store_true")
 
+parser.add_argument("--disable-xformers", action="store_true")
 
 vram_group = parser.add_mutually_exclusive_group()
-vram_group.add_argument("--gpu-only", action="store_true", help="Store and run everything (text encoders/CLIP models, etc... on the GPU).")
-vram_group.add_argument("--highvram", action="store_true", help="By default models will be unloaded to CPU memory after being used. This option keeps them in GPU memory.")
-vram_group.add_argument("--normalvram", action="store_true", help="Used to force normal vram use if lowvram gets automatically enabled.")
-vram_group.add_argument("--lowvram", action="store_true", help="Split the unet in parts to use less vram.")
-vram_group.add_argument("--novram", action="store_true", help="When lowvram isn't enough.")
-vram_group.add_argument("--cpu", action="store_true", help="To use the CPU for everything (slow).")
+vram_group.add_argument("--always-gpu", action="store_true")
+vram_group.add_argument("--always-high-vram", action="store_true")
+vram_group.add_argument("--always-normal-vram", action="store_true")
+vram_group.add_argument("--always-low-vram", action="store_true")
+vram_group.add_argument("--always-no-vram", action="store_true")
+vram_group.add_argument("--always-cpu", type=int, nargs="?", metavar="CPU_NUM_THREADS", const=-1)
 
+parser.add_argument("--always-offload-from-vram", action="store_true")
+parser.add_argument("--pytorch-deterministic", action="store_true")
 
-parser.add_argument("--disable-smart-memory", action="store_true", help="Force agressively offload to regular ram instead of keeping models in vram when it can.")
-parser.add_argument("--deterministic", action="store_true", help="Make pytorch use slower deterministic algorithms when it can. Note that this might not make images deterministic in all cases.")
+parser.add_argument("--disable-server-log", action="store_true")
+parser.add_argument("--debug-mode", action="store_true")
+parser.add_argument("--is-windows-embedded-python", action="store_true")
 
-parser.add_argument("--dont-print-server", action="store_true", help="Don't print server output.")
-parser.add_argument("--quick-test-for-ci", action="store_true", help="Quick test for CI.")
-parser.add_argument("--windows-standalone-build", action="store_true", help="Windows standalone build: Enable convenient things that most people using the standalone windows build will probably enjoy (like auto opening the page on startup).")
-
-parser.add_argument("--disable-metadata", action="store_true", help="Disable saving prompt metadata in files.")
-
-parser.add_argument("--multi-user", action="store_true", help="Enables per-user storage.")
-
-parser.add_argument("--verbose", action="store_true", help="Enables more debug prints.")
+parser.add_argument("--disable-server-info", action="store_true")
 
+parser.add_argument("--multi-user", action="store_true")
 
 if ldm_patched.modules.options.args_parsing:
-    args = parser.parse_args()
+    args = parser.parse_args([])
 else:
     args = parser.parse_args([])
 
-if args.windows_standalone_build:
+if args.is_windows_embedded_python:
     args.in_browser = True
 
 if args.disable_in_browser:
     args.in_browser = False
-
-import logging
-logging_level = logging.INFO
-if args.verbose:
-    logging_level = logging.DEBUG
-
-logging.basicConfig(format="%(message)s", level=logging_level)

ldm_patched/modules/clip_model.py

@@ -97,7 +97,7 @@ class CLIPTextModel_(torch.nn.Module):
         x = self.embeddings(input_tokens)
         mask = None
         if attention_mask is not None:
-            mask = 1.0 - attention_mask.to(x.dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1])).expand(attention_mask.shape[0], 1, attention_mask.shape[-1], attention_mask.shape[-1])
+            mask = 1.0 - attention_mask.to(x.dtype).unsqueeze(1).unsqueeze(1).expand(attention_mask.shape[0], 1, attention_mask.shape[-1], attention_mask.shape[-1])
             mask = mask.masked_fill(mask.to(torch.bool), float("-inf"))
 
         causal_mask = torch.empty(x.shape[1], x.shape[1], dtype=x.dtype, device=x.device).fill_(float("-inf")).triu_(1)
@@ -119,9 +119,6 @@ class CLIPTextModel(torch.nn.Module):
         super().__init__()
         self.num_layers = config_dict["num_hidden_layers"]
         self.text_model = CLIPTextModel_(config_dict, dtype, device, operations)
-        embed_dim = config_dict["hidden_size"]
-        self.text_projection = operations.Linear(embed_dim, embed_dim, bias=False, dtype=dtype, device=device)
-        self.text_projection.weight.copy_(torch.eye(embed_dim))
         self.dtype = dtype
 
     def get_input_embeddings(self):
@@ -131,10 +128,7 @@ class CLIPTextModel(torch.nn.Module):
         self.text_model.embeddings.token_embedding = embeddings
 
     def forward(self, *args, **kwargs):
-        x = self.text_model(*args, **kwargs)
-        out = self.text_projection(x[2])
-        return (x[0], x[1], out, x[2])
-
+        return self.text_model(*args, **kwargs)
 
 class CLIPVisionEmbeddings(torch.nn.Module):
     def __init__(self, embed_dim, num_channels=3, patch_size=14, image_size=224, dtype=None, device=None, operations=None):

ldm_patched/modules/clip_vision.py

@@ -2,7 +2,6 @@ from .utils import load_torch_file, transformers_convert, state_dict_prefix_repl
 import os
 import torch
 import json
-import logging
 
 import ldm_patched.modules.ops
 import ldm_patched.modules.model_patcher
@@ -56,7 +55,7 @@ class ClipVisionModel():
         outputs = Output()
         outputs["last_hidden_state"] = out[0].to(ldm_patched.modules.model_management.intermediate_device())
         outputs["image_embeds"] = out[2].to(ldm_patched.modules.model_management.intermediate_device())
-        outputs["penultimate_hidden_states"] = out[1].to(comfy.model_management.intermediate_device())
+        outputs["penultimate_hidden_states"] = out[1].to(ldm_patched.modules.model_management.intermediate_device())
         return outputs
 
 def convert_to_transformers(sd, prefix):
@@ -100,7 +99,7 @@ def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
     clip = ClipVisionModel(json_config)
     m, u = clip.load_sd(sd)
     if len(m) > 0:
-        logging.warning("missing clip vision: {}".format(m))
+        print("extra clip vision:", m)
     u = set(u)
     keys = list(sd.keys())
     for k in keys:

ldm_patched/modules/conds.py

@@ -3,8 +3,6 @@ import math
 import ldm_patched.modules.utils
 
 
-def lcm(a, b): #TODO: eventually replace by math.lcm (added in python3.9)
-    return abs(a*b) // math.gcd(a, b)
 
 class CONDRegular:
     def __init__(self, cond):
@@ -41,7 +39,7 @@ class CONDCrossAttn(CONDRegular):
             if s1[0] != s2[0] or s1[2] != s2[2]: #these 2 cases should not happen
                 return False
 
-            mult_min = lcm(s1[1], s2[1])
+            mult_min = math.lcm(s1[1], s2[1])
             diff = mult_min // min(s1[1], s2[1])
             if diff > 4: #arbitrary limit on the padding because it's probably going to impact performance negatively if it's too much
                 return False
@@ -52,7 +50,7 @@ class CONDCrossAttn(CONDRegular):
         crossattn_max_len = self.cond.shape[1]
         for x in others:
             c = x.cond
-            crossattn_max_len = lcm(crossattn_max_len, c.shape[1])
+            crossattn_max_len = math.lcm(crossattn_max_len, c.shape[1])
             conds.append(c)
 
         out = []

ldm_patched/modules/controlnet.py

@@ -1,7 +1,6 @@
 import torch
 import math
 import os
-import logging
 import ldm_patched.modules.utils
 import ldm_patched.modules.model_management
 import ldm_patched.modules.model_detection
@@ -10,7 +9,6 @@ import ldm_patched.modules.ops
 
 import ldm_patched.controlnet.cldm
 import ldm_patched.t2ia.adapter
-import ldm_patched.ldm.cascade.controlnet
 
 
 def broadcast_image_to(tensor, target_batch_size, batched_number):
@@ -39,8 +37,6 @@ class ControlBase:
         self.timestep_percent_range = (0.0, 1.0)
         self.global_average_pooling = False
         self.timestep_range = None
-        self.compression_ratio = 8
-        self.upscale_algorithm = 'nearest-exact'
 
         if device is None:
             device = ldm_patched.modules.model_management.get_torch_device()
@@ -81,8 +77,6 @@ class ControlBase:
         c.strength = self.strength
         c.timestep_percent_range = self.timestep_percent_range
         c.global_average_pooling = self.global_average_pooling
-        c.compression_ratio = self.compression_ratio
-        c.upscale_algorithm = self.upscale_algorithm
 
     def inference_memory_requirements(self, dtype):
         if self.previous_controlnet is not None:
@@ -138,13 +132,11 @@ class ControlBase:
         return out
 
 class ControlNet(ControlBase):
-    def __init__(self, control_model=None, global_average_pooling=False, device=None, load_device=None, manual_cast_dtype=None):
+    def __init__(self, control_model, global_average_pooling=False, device=None, load_device=None, manual_cast_dtype=None):
         super().__init__(device)
         self.control_model = control_model
         self.load_device = load_device
-        if control_model is not None:
-            self.control_model_wrapped = ldm_patched.modules.model_patcher.ModelPatcher(self.control_model, load_device=load_device, offload_device=ldm_patched.modules.model_management.unet_offload_device())
-
+        self.control_model_wrapped = ldm_patched.modules.model_patcher.ModelPatcher(self.control_model, load_device=load_device, offload_device=ldm_patched.modules.model_management.unet_offload_device())
         self.global_average_pooling = global_average_pooling
         self.model_sampling_current = None
         self.manual_cast_dtype = manual_cast_dtype
@@ -166,15 +158,15 @@ class ControlNet(ControlBase):
             dtype = self.manual_cast_dtype
 
         output_dtype = x_noisy.dtype
-        if self.cond_hint is None or x_noisy.shape[2] * self.compression_ratio != self.cond_hint.shape[2] or x_noisy.shape[3] * self.compression_ratio != self.cond_hint.shape[3]:
+        if self.cond_hint is None or x_noisy.shape[2] * 8 != self.cond_hint.shape[2] or x_noisy.shape[3] * 8 != self.cond_hint.shape[3]:
             if self.cond_hint is not None:
                 del self.cond_hint
             self.cond_hint = None
-            self.cond_hint = ldm_patched.modules.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * self.compression_ratio, x_noisy.shape[2] * self.compression_ratio, self.upscale_algorithm, "center").to(dtype).to(self.device)
+            self.cond_hint = ldm_patched.modules.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device)
         if x_noisy.shape[0] != self.cond_hint.shape[0]:
             self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
 
-        context = cond.get('crossattn_controlnet', cond['c_crossattn'])
+        context = cond['c_crossattn']
         y = cond.get('y', None)
         if y is not None:
             y = y.to(dtype)
@@ -185,9 +177,7 @@ class ControlNet(ControlBase):
         return self.control_merge(None, control, control_prev, output_dtype)
 
     def copy(self):
-        c = ControlNet(None, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype)
-        c.control_model = self.control_model
-        c.control_model_wrapped = self.control_model_wrapped
+        c = ControlNet(self.control_model, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype)
         self.copy_to(c)
         return c
 
@@ -205,7 +195,7 @@ class ControlNet(ControlBase):
         super().cleanup()
 
 class ControlLoraOps:
-    class Linear(torch.nn.Module, ldm_patched.modules.ops.CastWeightBiasOp):
+    class Linear(torch.nn.Module):
         def __init__(self, in_features: int, out_features: int, bias: bool = True,
                     device=None, dtype=None) -> None:
             factory_kwargs = {'device': device, 'dtype': dtype}
@@ -224,7 +214,7 @@ class ControlLoraOps:
             else:
                 return torch.nn.functional.linear(input, weight, bias)
 
-    class Conv2d(torch.nn.Module, ldm_patched.modules.ops.CastWeightBiasOp):
+    class Conv2d(torch.nn.Module):
         def __init__(
             self,
             in_channels,
@@ -297,13 +287,13 @@ class ControlLora(ControlNet):
         for k in sd:
             weight = sd[k]
             try:
-                ldm_patched.modules.utils.set_attr_param(self.control_model, k, weight)
+                ldm_patched.modules.utils.set_attr(self.control_model, k, weight)
             except:
                 pass
 
         for k in self.control_weights:
             if k not in {"lora_controlnet"}:
-                ldm_patched.modules.utils.set_attr_param(self.control_model, k, self.control_weights[k].to(dtype).to(ldm_patched.modules.model_management.get_torch_device()))
+                ldm_patched.modules.utils.set_attr(self.control_model, k, self.control_weights[k].to(dtype).to(ldm_patched.modules.model_management.get_torch_device()))
 
     def copy(self):
         c = ControlLora(self.control_weights, global_average_pooling=self.global_average_pooling)
@@ -328,10 +318,9 @@ def load_controlnet(ckpt_path, model=None):
         return ControlLora(controlnet_data)
 
     controlnet_config = None
-    supported_inference_dtypes = None
-
     if "controlnet_cond_embedding.conv_in.weight" in controlnet_data: #diffusers format
-        controlnet_config = ldm_patched.modules.model_detection.unet_config_from_diffusers_unet(controlnet_data)
+        unet_dtype = ldm_patched.modules.model_management.unet_dtype()
+        controlnet_config = ldm_patched.modules.model_detection.unet_config_from_diffusers_unet(controlnet_data, unet_dtype)
         diffusers_keys = ldm_patched.modules.utils.unet_to_diffusers(controlnet_config)
         diffusers_keys["controlnet_mid_block.weight"] = "middle_block_out.0.weight"
         diffusers_keys["controlnet_mid_block.bias"] = "middle_block_out.0.bias"
@@ -372,7 +361,7 @@ def load_controlnet(ckpt_path, model=None):
 
         leftover_keys = controlnet_data.keys()
         if len(leftover_keys) > 0:
-            logging.warning("leftover keys: {}".format(leftover_keys))
+            print("leftover keys:", leftover_keys)
         controlnet_data = new_sd
 
     pth_key = 'control_model.zero_convs.0.0.weight'
@@ -387,24 +376,16 @@ def load_controlnet(ckpt_path, model=None):
     else:
         net = load_t2i_adapter(controlnet_data)
         if net is None:
-            logging.error("error checkpoint does not contain controlnet or t2i adapter data {}".format(ckpt_path))
+            print("error checkpoint does not contain controlnet or t2i adapter data", ckpt_path)
         return net
 
     if controlnet_config is None:
-        model_config = ldm_patched.modules.model_detection.model_config_from_unet(controlnet_data, prefix, True)
-        supported_inference_dtypes = model_config.supported_inference_dtypes
-        controlnet_config = model_config.unet_config
-
-    load_device = ldm_patched.modules.model_management.get_torch_device()
-    if supported_inference_dtypes is None:
         unet_dtype = ldm_patched.modules.model_management.unet_dtype()
-    else:
-        unet_dtype = ldm_patched.modules.model_management.unet_dtype(supported_dtypes=supported_inference_dtypes)
-
+        controlnet_config = ldm_patched.modules.model_detection.model_config_from_unet(controlnet_data, prefix, unet_dtype, True).unet_config
+    load_device = ldm_patched.modules.model_management.get_torch_device()
     manual_cast_dtype = ldm_patched.modules.model_management.unet_manual_cast(unet_dtype, load_device)
     if manual_cast_dtype is not None:
         controlnet_config["operations"] = ldm_patched.modules.ops.manual_cast
-    controlnet_config["dtype"] = unet_dtype
     controlnet_config.pop("out_channels")
     controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1]
     control_model = ldm_patched.controlnet.cldm.ControlNet(**controlnet_config)
@@ -422,7 +403,7 @@ def load_controlnet(ckpt_path, model=None):
                             cd = controlnet_data[x]
                             cd += model_sd[sd_key].type(cd.dtype).to(cd.device)
             else:
-                logging.warning("WARNING: Loaded a diff controlnet without a model. It will very likely not work.")
+                print("WARNING: Loaded a diff controlnet without a model. It will very likely not work.")
 
         class WeightsLoader(torch.nn.Module):
             pass
@@ -431,12 +412,7 @@ def load_controlnet(ckpt_path, model=None):
         missing, unexpected = w.load_state_dict(controlnet_data, strict=False)
     else:
         missing, unexpected = control_model.load_state_dict(controlnet_data, strict=False)
-
-    if len(missing) > 0:
-        logging.warning("missing controlnet keys: {}".format(missing))
-
-    if len(unexpected) > 0:
-        logging.debug("unexpected controlnet keys: {}".format(unexpected))
+    print(missing, unexpected)
 
     global_average_pooling = False
     filename = os.path.splitext(ckpt_path)[0]
@@ -447,13 +423,11 @@ def load_controlnet(ckpt_path, model=None):
     return control
 
 class T2IAdapter(ControlBase):
-    def __init__(self, t2i_model, channels_in, compression_ratio, upscale_algorithm, device=None):
+    def __init__(self, t2i_model, channels_in, device=None):
         super().__init__(device)
         self.t2i_model = t2i_model
         self.channels_in = channels_in
         self.control_input = None
-        self.compression_ratio = compression_ratio
-        self.upscale_algorithm = upscale_algorithm
 
     def scale_image_to(self, width, height):
         unshuffle_amount = self.t2i_model.unshuffle_amount
@@ -473,13 +447,13 @@ class T2IAdapter(ControlBase):
                 else:
                     return None
 
-        if self.cond_hint is None or x_noisy.shape[2] * self.compression_ratio != self.cond_hint.shape[2] or x_noisy.shape[3] * self.compression_ratio != self.cond_hint.shape[3]:
+        if self.cond_hint is None or x_noisy.shape[2] * 8 != self.cond_hint.shape[2] or x_noisy.shape[3] * 8 != self.cond_hint.shape[3]:
             if self.cond_hint is not None:
                 del self.cond_hint
             self.control_input = None
             self.cond_hint = None
-            width, height = self.scale_image_to(x_noisy.shape[3] * self.compression_ratio, x_noisy.shape[2] * self.compression_ratio)
-            self.cond_hint = ldm_patched.modules.utils.common_upscale(self.cond_hint_original, width, height, self.upscale_algorithm, "center").float().to(self.device)
+            width, height = self.scale_image_to(x_noisy.shape[3] * 8, x_noisy.shape[2] * 8)
+            self.cond_hint = ldm_patched.modules.utils.common_upscale(self.cond_hint_original, width, height, 'nearest-exact', "center").float().to(self.device)
             if self.channels_in == 1 and self.cond_hint.shape[1] > 1:
                 self.cond_hint = torch.mean(self.cond_hint, 1, keepdim=True)
         if x_noisy.shape[0] != self.cond_hint.shape[0]:
@@ -498,14 +472,11 @@ class T2IAdapter(ControlBase):
         return self.control_merge(control_input, mid, control_prev, x_noisy.dtype)
 
     def copy(self):
-        c = T2IAdapter(self.t2i_model, self.channels_in, self.compression_ratio, self.upscale_algorithm)
+        c = T2IAdapter(self.t2i_model, self.channels_in)
         self.copy_to(c)
         return c
 
 def load_t2i_adapter(t2i_data):
-    compression_ratio = 8
-    upscale_algorithm = 'nearest-exact'
-
     if 'adapter' in t2i_data:
         t2i_data = t2i_data['adapter']
     if 'adapter.body.0.resnets.0.block1.weight' in t2i_data: #diffusers format
@@ -533,22 +504,13 @@ def load_t2i_adapter(t2i_data):
         if cin == 256 or cin == 768:
             xl = True
         model_ad = ldm_patched.t2ia.adapter.Adapter(cin=cin, channels=[channel, channel*2, channel*4, channel*4][:4], nums_rb=2, ksize=ksize, sk=True, use_conv=use_conv, xl=xl)
-    elif "backbone.0.0.weight" in keys:
-        model_ad = ldm_patched.ldm.cascade.controlnet.ControlNet(c_in=t2i_data['backbone.0.0.weight'].shape[1], proj_blocks=[0, 4, 8, 12, 51, 55, 59, 63])
-        compression_ratio = 32
-        upscale_algorithm = 'bilinear'
-    elif "backbone.10.blocks.0.weight" in keys:
-        model_ad = ldm_patched.ldm.cascade.controlnet.ControlNet(c_in=t2i_data['backbone.0.weight'].shape[1], bottleneck_mode="large", proj_blocks=[0, 4, 8, 12, 51, 55, 59, 63])
-        compression_ratio = 1
-        upscale_algorithm = 'nearest-exact'
     else:
         return None
-
     missing, unexpected = model_ad.load_state_dict(t2i_data)
     if len(missing) > 0:
-        logging.warning("t2i missing {}".format(missing))
+        print("t2i missing", missing)
 
     if len(unexpected) > 0:
-        logging.debug("t2i unexpected {}".format(unexpected))
+        print("t2i unexpected", unexpected)
 
-    return T2IAdapter(model_ad, model_ad.input_channels, compression_ratio, upscale_algorithm)
+    return T2IAdapter(model_ad, model_ad.input_channels)

ldm_patched/modules/diffusers_convert.py

@@ -1,6 +1,5 @@
 import re
 import torch
-import logging
 
 # conversion code from https://github.com/huggingface/diffusers/blob/main/scripts/convert_diffusers_to_original_stable_diffusion.py
 
@@ -178,7 +177,7 @@ def convert_vae_state_dict(vae_state_dict):
     for k, v in new_state_dict.items():
         for weight_name in weights_to_convert:
             if f"mid.attn_1.{weight_name}.weight" in k:
-                logging.debug(f"Reshaping {k} for SD format")
+                print(f"Reshaping {k} for SD format")
                 new_state_dict[k] = reshape_weight_for_sd(v)
     return new_state_dict
 
@@ -206,21 +205,6 @@ textenc_pattern = re.compile("|".join(protected.keys()))
 # Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
 code2idx = {"q": 0, "k": 1, "v": 2}
 
-# This function exists because at the time of writing torch.cat can't do fp8 with cuda
-def cat_tensors(tensors):
-    x = 0
-    for t in tensors:
-        x += t.shape[0]
-
-    shape = [x] + list(tensors[0].shape)[1:]
-    out = torch.empty(shape, device=tensors[0].device, dtype=tensors[0].dtype)
-
-    x = 0
-    for t in tensors:
-        out[x:x + t.shape[0]] = t
-        x += t.shape[0]
-
-    return out
 
 def convert_text_enc_state_dict_v20(text_enc_dict, prefix=""):
     new_state_dict = {}
@@ -253,24 +237,20 @@ def convert_text_enc_state_dict_v20(text_enc_dict, prefix=""):
             capture_qkv_bias[k_pre][code2idx[k_code]] = v
             continue
 
-        text_proj = "transformer.text_projection.weight"
-        if k.endswith(text_proj):
-            new_state_dict[k.replace(text_proj, "text_projection")] = v.transpose(0, 1).contiguous()
-        else:
-            relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k)
-            new_state_dict[relabelled_key] = v
+        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k)
+        new_state_dict[relabelled_key] = v
 
     for k_pre, tensors in capture_qkv_weight.items():
         if None in tensors:
             raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
         relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
-        new_state_dict[relabelled_key + ".in_proj_weight"] = cat_tensors(tensors)
+        new_state_dict[relabelled_key + ".in_proj_weight"] = torch.cat(tensors)
 
     for k_pre, tensors in capture_qkv_bias.items():
         if None in tensors:
             raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
         relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
-        new_state_dict[relabelled_key + ".in_proj_bias"] = cat_tensors(tensors)
+        new_state_dict[relabelled_key + ".in_proj_bias"] = torch.cat(tensors)
 
     return new_state_dict
 

ldm_patched/modules/gligen.py

@@ -1,11 +1,8 @@
-import math
-
 import torch
 from torch import nn
 from ldm_patched.ldm.modules.attention import CrossAttention
 from inspect import isfunction
-import ldm_patched.modules.ops
-ops = ldm_patched.modules.ops.manual_cast
+
 
 def exists(val):
     return val is not None
@@ -25,7 +22,7 @@ def default(val, d):
 class GEGLU(nn.Module):
     def __init__(self, dim_in, dim_out):
         super().__init__()
-        self.proj = ops.Linear(dim_in, dim_out * 2)
+        self.proj = nn.Linear(dim_in, dim_out * 2)
 
     def forward(self, x):
         x, gate = self.proj(x).chunk(2, dim=-1)
@@ -38,14 +35,14 @@ class FeedForward(nn.Module):
         inner_dim = int(dim * mult)
         dim_out = default(dim_out, dim)
         project_in = nn.Sequential(
-            ops.Linear(dim, inner_dim),
+            nn.Linear(dim, inner_dim),
             nn.GELU()
         ) if not glu else GEGLU(dim, inner_dim)
 
         self.net = nn.Sequential(
             project_in,
             nn.Dropout(dropout),
-            ops.Linear(inner_dim, dim_out)
+            nn.Linear(inner_dim, dim_out)
         )
 
     def forward(self, x):
@@ -60,12 +57,11 @@ class GatedCrossAttentionDense(nn.Module):
             query_dim=query_dim,
             context_dim=context_dim,
             heads=n_heads,
-            dim_head=d_head,
-            operations=ops)
+            dim_head=d_head)
         self.ff = FeedForward(query_dim, glu=True)
 
-        self.norm1 = ops.LayerNorm(query_dim)
-        self.norm2 = ops.LayerNorm(query_dim)
+        self.norm1 = nn.LayerNorm(query_dim)
+        self.norm2 = nn.LayerNorm(query_dim)
 
         self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)))
         self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)))
@@ -91,18 +87,17 @@ class GatedSelfAttentionDense(nn.Module):
 
         # we need a linear projection since we need cat visual feature and obj
         # feature
-        self.linear = ops.Linear(context_dim, query_dim)
+        self.linear = nn.Linear(context_dim, query_dim)
 
         self.attn = CrossAttention(
             query_dim=query_dim,
             context_dim=query_dim,
             heads=n_heads,
-            dim_head=d_head,
-            operations=ops)
+            dim_head=d_head)
         self.ff = FeedForward(query_dim, glu=True)
 
-        self.norm1 = ops.LayerNorm(query_dim)
-        self.norm2 = ops.LayerNorm(query_dim)
+        self.norm1 = nn.LayerNorm(query_dim)
+        self.norm2 = nn.LayerNorm(query_dim)
 
         self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)))
         self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)))
@@ -131,14 +126,14 @@ class GatedSelfAttentionDense2(nn.Module):
 
         # we need a linear projection since we need cat visual feature and obj
         # feature
-        self.linear = ops.Linear(context_dim, query_dim)
+        self.linear = nn.Linear(context_dim, query_dim)
 
         self.attn = CrossAttention(
-            query_dim=query_dim, context_dim=query_dim, dim_head=d_head, operations=ops)
+            query_dim=query_dim, context_dim=query_dim, dim_head=d_head)
         self.ff = FeedForward(query_dim, glu=True)
 
-        self.norm1 = ops.LayerNorm(query_dim)
-        self.norm2 = ops.LayerNorm(query_dim)
+        self.norm1 = nn.LayerNorm(query_dim)
+        self.norm2 = nn.LayerNorm(query_dim)
 
         self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)))
         self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)))
@@ -206,11 +201,11 @@ class PositionNet(nn.Module):
         self.position_dim = fourier_freqs * 2 * 4  # 2 is sin&cos, 4 is xyxy
 
         self.linears = nn.Sequential(
-            ops.Linear(self.in_dim + self.position_dim, 512),
+            nn.Linear(self.in_dim + self.position_dim, 512),
             nn.SiLU(),
-            ops.Linear(512, 512),
+            nn.Linear(512, 512),
             nn.SiLU(),
-            ops.Linear(512, out_dim),
+            nn.Linear(512, out_dim),
         )
 
         self.null_positive_feature = torch.nn.Parameter(
@@ -220,15 +215,16 @@ class PositionNet(nn.Module):
 
     def forward(self, boxes, masks, positive_embeddings):
         B, N, _ = boxes.shape
-        masks = masks.unsqueeze(-1)
-        positive_embeddings = positive_embeddings
+        dtype = self.linears[0].weight.dtype
+        masks = masks.unsqueeze(-1).to(dtype)
+        positive_embeddings = positive_embeddings.to(dtype)
 
         # embedding position (it may includes padding as placeholder)
-        xyxy_embedding = self.fourier_embedder(boxes)  # B*N*4 --> B*N*C
+        xyxy_embedding = self.fourier_embedder(boxes.to(dtype))  # B*N*4 --> B*N*C
 
         # learnable null embedding
-        positive_null = self.null_positive_feature.to(device=boxes.device, dtype=boxes.dtype).view(1, 1, -1)
-        xyxy_null = self.null_position_feature.to(device=boxes.device, dtype=boxes.dtype).view(1, 1, -1)
+        positive_null = self.null_positive_feature.view(1, 1, -1)
+        xyxy_null = self.null_position_feature.view(1, 1, -1)
 
         # replace padding with learnable null embedding
         positive_embeddings = positive_embeddings * \
@@ -255,7 +251,7 @@ class Gligen(nn.Module):
         def func(x, extra_options):
             key = extra_options["transformer_index"]
             module = self.module_list[key]
-            return module(x, objs.to(device=x.device, dtype=x.dtype))
+            return module(x, objs)
         return func
 
     def set_position(self, latent_image_shape, position_params, device):

ldm_patched/modules/latent_formats.py

@@ -101,4 +101,4 @@ class SC_B(LatentFormat):
             [-0.2093, -0.0222, -0.0195],
             [-0.3087, -0.1535,  0.0366],
             [ 0.0290, -0.1574, -0.4078]
-        ]
+        ]

ldm_patched/modules/lora.py

@@ -1,5 +1,4 @@
 import ldm_patched.modules.utils
-import logging
 
 LORA_CLIP_MAP = {
     "mlp.fc1": "mlp_fc1",
@@ -21,12 +20,6 @@ def load_lora(lora, to_load):
             alpha = lora[alpha_name].item()
             loaded_keys.add(alpha_name)
 
-        dora_scale_name = "{}.dora_scale".format(x)
-        dora_scale = None
-        if dora_scale_name in lora.keys():
-            dora_scale = lora[dora_scale_name]
-            loaded_keys.add(dora_scale_name)
-
         regular_lora = "{}.lora_up.weight".format(x)
         diffusers_lora = "{}_lora.up.weight".format(x)
         transformers_lora = "{}.lora_linear_layer.up.weight".format(x)
@@ -50,7 +43,7 @@ def load_lora(lora, to_load):
             if mid_name is not None and mid_name in lora.keys():
                 mid = lora[mid_name]
                 loaded_keys.add(mid_name)
-            patch_dict[to_load[x]] = ("lora", (lora[A_name], lora[B_name], alpha, mid, dora_scale))
+            patch_dict[to_load[x]] = ("lora", (lora[A_name], lora[B_name], alpha, mid))
             loaded_keys.add(A_name)
             loaded_keys.add(B_name)
 
@@ -71,7 +64,7 @@ def load_lora(lora, to_load):
                 loaded_keys.add(hada_t1_name)
                 loaded_keys.add(hada_t2_name)
 
-            patch_dict[to_load[x]] = ("loha", (lora[hada_w1_a_name], lora[hada_w1_b_name], alpha, lora[hada_w2_a_name], lora[hada_w2_b_name], hada_t1, hada_t2, dora_scale))
+            patch_dict[to_load[x]] = ("loha", (lora[hada_w1_a_name], lora[hada_w1_b_name], alpha, lora[hada_w2_a_name], lora[hada_w2_b_name], hada_t1, hada_t2))
             loaded_keys.add(hada_w1_a_name)
             loaded_keys.add(hada_w1_b_name)
             loaded_keys.add(hada_w2_a_name)
@@ -123,7 +116,7 @@ def load_lora(lora, to_load):
             loaded_keys.add(lokr_t2_name)
 
         if (lokr_w1 is not None) or (lokr_w2 is not None) or (lokr_w1_a is not None) or (lokr_w2_a is not None):
-            patch_dict[to_load[x]] = ("lokr", (lokr_w1, lokr_w2, alpha, lokr_w1_a, lokr_w1_b, lokr_w2_a, lokr_w2_b, lokr_t2, dora_scale))
+            patch_dict[to_load[x]] = ("lokr", (lokr_w1, lokr_w2, alpha, lokr_w1_a, lokr_w1_b, lokr_w2_a, lokr_w2_b, lokr_t2))
 
         #glora
         a1_name = "{}.a1.weight".format(x)
@@ -131,7 +124,7 @@ def load_lora(lora, to_load):
         b1_name = "{}.b1.weight".format(x)
         b2_name = "{}.b2.weight".format(x)
         if a1_name in lora:
-            patch_dict[to_load[x]] = ("glora", (lora[a1_name], lora[a2_name], lora[b1_name], lora[b2_name], alpha, dora_scale))
+            patch_dict[to_load[x]] = ("glora", (lora[a1_name], lora[a2_name], lora[b1_name], lora[b2_name], alpha))
             loaded_keys.add(a1_name)
             loaded_keys.add(a2_name)
             loaded_keys.add(b1_name)
@@ -163,7 +156,7 @@ def load_lora(lora, to_load):
 
     for x in lora.keys():
         if x not in loaded_keys:
-            logging.warning("lora key not loaded: {}".format(x))
+            print("lora key not loaded", x)
     return patch_dict
 
 def model_lora_keys_clip(model, key_map={}):
@@ -204,15 +197,6 @@ def model_lora_keys_clip(model, key_map={}):
                     key_map[lora_key] = k
                     lora_key = "text_encoder.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora
                     key_map[lora_key] = k
-                    lora_key = "lora_prior_te_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #cascade lora: TODO put lora key prefix in the model config
-                    key_map[lora_key] = k
-
-
-    k = "clip_g.transformer.text_projection.weight"
-    if k in sdk:
-        key_map["lora_prior_te_text_projection"] = k #cascade lora?
-        # key_map["text_encoder.text_projection"] = k #TODO: check if other lora have the text_projection too
-        # key_map["lora_te_text_projection"] = k
 
     return key_map
 
@@ -223,7 +207,6 @@ def model_lora_keys_unet(model, key_map={}):
         if k.startswith("diffusion_model.") and k.endswith(".weight"):
             key_lora = k[len("diffusion_model."):-len(".weight")].replace(".", "_")
             key_map["lora_unet_{}".format(key_lora)] = k
-            key_map["lora_prior_unet_{}".format(key_lora)] = k #cascade lora: TODO put lora key prefix in the model config
 
     diffusers_keys = ldm_patched.modules.utils.unet_to_diffusers(model.model_config.unet_config)
     for k in diffusers_keys:

ldm_patched/modules/model_base.py

@@ -1,8 +1,5 @@
 import torch
-import logging
 from ldm_patched.ldm.modules.diffusionmodules.openaimodel import UNetModel, Timestep
-from ldm_patched.ldm.cascade.stage_c import StageC
-from ldm_patched.ldm.cascade.stage_b import StageB
 from ldm_patched.ldm.modules.encoders.noise_aug_modules import CLIPEmbeddingNoiseAugmentation
 from ldm_patched.ldm.modules.diffusionmodules.upscaling import ImageConcatWithNoiseAugmentation
 import ldm_patched.modules.model_management
@@ -15,11 +12,9 @@ class ModelType(Enum):
     EPS = 1
     V_PREDICTION = 2
     V_PREDICTION_EDM = 3
-    STABLE_CASCADE = 4
-    EDM = 5
 
 
-from ldm_patched.modules.model_sampling import EPS, V_PREDICTION, EDM, ModelSamplingDiscrete, ModelSamplingContinuousEDM, StableCascadeSampling
+from ldm_patched.modules.model_sampling import EPS, V_PREDICTION, ModelSamplingDiscrete, ModelSamplingContinuousEDM
 
 
 def model_sampling(model_config, model_type):
@@ -32,12 +27,6 @@ def model_sampling(model_config, model_type):
     elif model_type == ModelType.V_PREDICTION_EDM:
         c = V_PREDICTION
         s = ModelSamplingContinuousEDM
-    elif model_type == ModelType.STABLE_CASCADE:
-        c = EPS
-        s = StableCascadeSampling
-    elif model_type == ModelType.EDM:
-        c = EDM
-        s = ModelSamplingContinuousEDM
 
     class ModelSampling(s, c):
         pass
@@ -46,7 +35,7 @@ def model_sampling(model_config, model_type):
 
 
 class BaseModel(torch.nn.Module):
-    def __init__(self, model_config, model_type=ModelType.EPS, device=None, unet_model=UNetModel):
+    def __init__(self, model_config, model_type=ModelType.EPS, device=None):
         super().__init__()
 
         unet_config = model_config.unet_config
@@ -59,17 +48,16 @@ class BaseModel(torch.nn.Module):
                 operations = ldm_patched.modules.ops.manual_cast
             else:
                 operations = ldm_patched.modules.ops.disable_weight_init
-            self.diffusion_model = unet_model(**unet_config, device=device, operations=operations)
+            self.diffusion_model = UNetModel(**unet_config, device=device, operations=operations)
         self.model_type = model_type
         self.model_sampling = model_sampling(model_config, model_type)
 
         self.adm_channels = unet_config.get("adm_in_channels", None)
         if self.adm_channels is None:
             self.adm_channels = 0
-
-        self.concat_keys = ()
-        logging.info("model_type {}".format(model_type.name))
-        logging.debug("adm {}".format(self.adm_channels))
+        self.inpaint_model = False
+        print("model_type", model_type.name)
+        print("UNet ADM Dimension", self.adm_channels)
 
     def apply_model(self, x, t, c_concat=None, c_crossattn=None, control=None, transformer_options={}, **kwargs):
         sigma = t
@@ -108,7 +96,8 @@ class BaseModel(torch.nn.Module):
 
     def extra_conds(self, **kwargs):
         out = {}
-        if len(self.concat_keys) > 0:
+        if self.inpaint_model:
+            concat_keys = ("mask", "masked_image")
             cond_concat = []
             denoise_mask = kwargs.get("concat_mask", kwargs.get("denoise_mask", None))
             concat_latent_image = kwargs.get("concat_latent_image", None)
@@ -125,16 +114,24 @@ class BaseModel(torch.nn.Module):
 
             concat_latent_image = utils.resize_to_batch_size(concat_latent_image, noise.shape[0])
 
-            if denoise_mask is not None:
-                if len(denoise_mask.shape) == len(noise.shape):
-                    denoise_mask = denoise_mask[:,:1]
+            if len(denoise_mask.shape) == len(noise.shape):
+                denoise_mask = denoise_mask[:,:1]
 
-                denoise_mask = denoise_mask.reshape((-1, 1, denoise_mask.shape[-2], denoise_mask.shape[-1]))
-                if denoise_mask.shape[-2:] != noise.shape[-2:]:
-                    denoise_mask = utils.common_upscale(denoise_mask, noise.shape[-1], noise.shape[-2], "bilinear", "center")
-                denoise_mask = utils.resize_to_batch_size(denoise_mask.round(), noise.shape[0])
+            denoise_mask = denoise_mask.reshape((-1, 1, denoise_mask.shape[-2], denoise_mask.shape[-1]))
+            if denoise_mask.shape[-2:] != noise.shape[-2:]:
+                denoise_mask = utils.common_upscale(denoise_mask, noise.shape[-1], noise.shape[-2], "bilinear", "center")
+            denoise_mask = utils.resize_to_batch_size(denoise_mask.round(), noise.shape[0])
 
-            for ck in self.concat_keys:
+            def blank_inpaint_image_like(latent_image):
+                blank_image = torch.ones_like(latent_image)
+                # these are the values for "zero" in pixel space translated to latent space
+                blank_image[:,0] *= 0.8223
+                blank_image[:,1] *= -0.6876
+                blank_image[:,2] *= 0.6364
+                blank_image[:,3] *= 0.1380
+                return blank_image
+
+            for ck in concat_keys:
                 if denoise_mask is not None:
                     if ck == "mask":
                         cond_concat.append(denoise_mask.to(device))
@@ -144,7 +141,7 @@ class BaseModel(torch.nn.Module):
                     if ck == "mask":
                         cond_concat.append(torch.ones_like(noise)[:,:1])
                     elif ck == "masked_image":
-                        cond_concat.append(self.blank_inpaint_image_like(noise))
+                        cond_concat.append(blank_inpaint_image_like(noise))
             data = torch.cat(cond_concat, dim=1)
             out['c_concat'] = ldm_patched.modules.conds.CONDNoiseShape(data)
 
@@ -156,14 +153,6 @@ class BaseModel(torch.nn.Module):
         if cross_attn is not None:
             out['c_crossattn'] = ldm_patched.modules.conds.CONDCrossAttn(cross_attn)
 
-        cross_attn_cnet = kwargs.get("cross_attn_controlnet", None)
-        if cross_attn_cnet is not None:
-            out['crossattn_controlnet'] = ldm_patched.modules.conds.CONDCrossAttn(cross_attn_cnet)
-
-        c_concat = kwargs.get("noise_concat", None)
-        if c_concat is not None:
-            out['c_concat'] = ldm_patched.modules.conds.CONDNoiseShape(c_concat)
-
         return out
 
     def load_model_weights(self, sd, unet_prefix=""):
@@ -176,10 +165,10 @@ class BaseModel(torch.nn.Module):
         to_load = self.model_config.process_unet_state_dict(to_load)
         m, u = self.diffusion_model.load_state_dict(to_load, strict=False)
         if len(m) > 0:
-            logging.warning("unet missing: {}".format(m))
+            print("unet missing:", m)
 
         if len(u) > 0:
-            logging.warning("unet unexpected: {}".format(u))
+            print("unet unexpected:", u)
         del to_load
         return self
 
@@ -213,16 +202,7 @@ class BaseModel(torch.nn.Module):
         return unet_state_dict
 
     def set_inpaint(self):
-        self.concat_keys = ("mask", "masked_image")
-        def blank_inpaint_image_like(latent_image):
-            blank_image = torch.ones_like(latent_image)
-            # these are the values for "zero" in pixel space translated to latent space
-            blank_image[:,0] *= 0.8223
-            blank_image[:,1] *= -0.6876
-            blank_image[:,2] *= 0.6364
-            blank_image[:,3] *= 0.1380
-            return blank_image
-        self.blank_inpaint_image_like = blank_inpaint_image_like
+        self.inpaint_model = True
 
     def memory_required(self, input_shape):
         if ldm_patched.modules.model_management.xformers_enabled() or ldm_patched.modules.model_management.pytorch_attention_flash_attention():
@@ -378,39 +358,10 @@ class SVD_img2vid(BaseModel):
         if "time_conditioning" in kwargs:
             out["time_context"] = ldm_patched.modules.conds.CONDCrossAttn(kwargs["time_conditioning"])
 
+        out['image_only_indicator'] = ldm_patched.modules.conds.CONDConstant(torch.zeros((1,), device=device))
         out['num_video_frames'] = ldm_patched.modules.conds.CONDConstant(noise.shape[0])
         return out
 
-class SV3D_u(SVD_img2vid):
-    def encode_adm(self, **kwargs):
-        augmentation = kwargs.get("augmentation_level", 0)
-
-        out = []
-        out.append(self.embedder(torch.flatten(torch.Tensor([augmentation]))))
-
-        flat = torch.flatten(torch.cat(out)).unsqueeze(dim=0)
-        return flat
-
-class SV3D_p(SVD_img2vid):
-    def __init__(self, model_config, model_type=ModelType.V_PREDICTION_EDM, device=None):
-        super().__init__(model_config, model_type, device=device)
-        self.embedder_512 = Timestep(512)
-
-    def encode_adm(self, **kwargs):
-        augmentation = kwargs.get("augmentation_level", 0)
-        elevation = kwargs.get("elevation", 0) #elevation and azimuth are in degrees here
-        azimuth = kwargs.get("azimuth", 0)
-        noise = kwargs.get("noise", None)
-
-        out = []
-        out.append(self.embedder(torch.flatten(torch.Tensor([augmentation]))))
-        out.append(self.embedder_512(torch.deg2rad(torch.fmod(torch.flatten(90 - torch.Tensor([elevation])), 360.0))))
-        out.append(self.embedder_512(torch.deg2rad(torch.fmod(torch.flatten(torch.Tensor([azimuth])), 360.0))))
-
-        out = list(map(lambda a: utils.resize_to_batch_size(a, noise.shape[0]), out))
-        return torch.cat(out, dim=1)
-
-
 class Stable_Zero123(BaseModel):
     def __init__(self, model_config, model_type=ModelType.EPS, device=None, cc_projection_weight=None, cc_projection_bias=None):
         super().__init__(model_config, model_type, device=device)
@@ -472,88 +423,3 @@ class SD_X4Upscaler(BaseModel):
         out['c_concat'] = ldm_patched.modules.conds.CONDNoiseShape(image)
         out['y'] = ldm_patched.modules.conds.CONDRegular(noise_level)
         return out
-
-class IP2P:
-    def extra_conds(self, **kwargs):
-        out = {}
-
-        image = kwargs.get("concat_latent_image", None)
-        noise = kwargs.get("noise", None)
-        device = kwargs["device"]
-
-        if image is None:
-            image = torch.zeros_like(noise)
-
-        if image.shape[1:] != noise.shape[1:]:
-            image = utils.common_upscale(image.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center")
-
-        image = utils.resize_to_batch_size(image, noise.shape[0])
-
-        out['c_concat'] = ldm_patched.modules.conds.CONDNoiseShape(self.process_ip2p_image_in(image))
-        adm = self.encode_adm(**kwargs)
-        if adm is not None:
-            out['y'] = ldm_patched.modules.conds.CONDRegular(adm)
-        return out
-
-class SD15_instructpix2pix(IP2P, BaseModel):
-    def __init__(self, model_config, model_type=ModelType.EPS, device=None):
-        super().__init__(model_config, model_type, device=device)
-        self.process_ip2p_image_in = lambda image: image
-
-class SDXL_instructpix2pix(IP2P, SDXL):
-    def __init__(self, model_config, model_type=ModelType.EPS, device=None):
-        super().__init__(model_config, model_type, device=device)
-        if model_type == ModelType.V_PREDICTION_EDM:
-            self.process_ip2p_image_in = lambda image: ldm_patched.modules.latent_formats.SDXL().process_in(image) #cosxl ip2p
-        else:
-            self.process_ip2p_image_in = lambda image: image #diffusers ip2p
-
-
-class StableCascade_C(BaseModel):
-    def __init__(self, model_config, model_type=ModelType.STABLE_CASCADE, device=None):
-        super().__init__(model_config, model_type, device=device, unet_model=StageC)
-        self.diffusion_model.eval().requires_grad_(False)
-
-    def extra_conds(self, **kwargs):
-        out = {}
-        clip_text_pooled = kwargs["pooled_output"]
-        if clip_text_pooled is not None:
-            out['clip_text_pooled'] = ldm_patched.modules.conds.CONDRegular(clip_text_pooled)
-
-        if "unclip_conditioning" in kwargs:
-            embeds = []
-            for unclip_cond in kwargs["unclip_conditioning"]:
-                weight = unclip_cond["strength"]
-                embeds.append(unclip_cond["clip_vision_output"].image_embeds.unsqueeze(0) * weight)
-            clip_img = torch.cat(embeds, dim=1)
-        else:
-            clip_img = torch.zeros((1, 1, 768))
-        out["clip_img"] = ldm_patched.modules.conds.CONDRegular(clip_img)
-        out["sca"] = ldm_patched.modules.conds.CONDRegular(torch.zeros((1,)))
-        out["crp"] = ldm_patched.modules.conds.CONDRegular(torch.zeros((1,)))
-
-        cross_attn = kwargs.get("cross_attn", None)
-        if cross_attn is not None:
-            out['clip_text'] = ldm_patched.modules.conds.CONDCrossAttn(cross_attn)
-        return out
-
-
-class StableCascade_B(BaseModel):
-    def __init__(self, model_config, model_type=ModelType.STABLE_CASCADE, device=None):
-        super().__init__(model_config, model_type, device=device, unet_model=StageB)
-        self.diffusion_model.eval().requires_grad_(False)
-
-    def extra_conds(self, **kwargs):
-        out = {}
-        noise = kwargs.get("noise", None)
-
-        clip_text_pooled = kwargs["pooled_output"]
-        if clip_text_pooled is not None:
-            out['clip'] = ldm_patched.modules.conds.CONDRegular(clip_text_pooled)
-
-        #size of prior doesn't really matter if zeros because it gets resized but I still want it to get batched
-        prior = kwargs.get("stable_cascade_prior", torch.zeros((1, 16, (noise.shape[2] * 4) // 42, (noise.shape[3] * 4) // 42), dtype=noise.dtype, layout=noise.layout, device=noise.device))
-
-        out["effnet"] = ldm_patched.modules.conds.CONDRegular(prior)
-        out["sca"] = ldm_patched.modules.conds.CONDRegular(torch.zeros((1,)))
-        return out

ldm_patched/modules/model_detection.py

@@ -1,6 +1,5 @@
 import ldm_patched.modules.supported_models
 import ldm_patched.modules.supported_models_base
-import logging
 
 def count_blocks(state_dict_keys, prefix_string):
     count = 0
@@ -29,38 +28,9 @@ def calculate_transformer_depth(prefix, state_dict_keys, state_dict):
         return last_transformer_depth, context_dim, use_linear_in_transformer, time_stack
     return None
 
-def detect_unet_config(state_dict, key_prefix):
+def detect_unet_config(state_dict, key_prefix, dtype):
     state_dict_keys = list(state_dict.keys())
 
-    if '{}clf.1.weight'.format(key_prefix) in state_dict_keys: #stable cascade
-        unet_config = {}
-        text_mapper_name = '{}clip_txt_mapper.weight'.format(key_prefix)
-        if text_mapper_name in state_dict_keys:
-            unet_config['stable_cascade_stage'] = 'c'
-            w = state_dict[text_mapper_name]
-            if w.shape[0] == 1536: #stage c lite
-                unet_config['c_cond'] = 1536
-                unet_config['c_hidden'] = [1536, 1536]
-                unet_config['nhead'] = [24, 24]
-                unet_config['blocks'] = [[4, 12], [12, 4]]
-            elif w.shape[0] == 2048: #stage c full
-                unet_config['c_cond'] = 2048
-        elif '{}clip_mapper.weight'.format(key_prefix) in state_dict_keys:
-            unet_config['stable_cascade_stage'] = 'b'
-            w = state_dict['{}down_blocks.1.0.channelwise.0.weight'.format(key_prefix)]
-            if w.shape[-1] == 640:
-                unet_config['c_hidden'] = [320, 640, 1280, 1280]
-                unet_config['nhead'] = [-1, -1, 20, 20]
-                unet_config['blocks'] = [[2, 6, 28, 6], [6, 28, 6, 2]]
-                unet_config['block_repeat'] = [[1, 1, 1, 1], [3, 3, 2, 2]]
-            elif w.shape[-1] == 576: #stage b lite
-                unet_config['c_hidden'] = [320, 576, 1152, 1152]
-                unet_config['nhead'] = [-1, 9, 18, 18]
-                unet_config['blocks'] = [[2, 4, 14, 4], [4, 14, 4, 2]]
-                unet_config['block_repeat'] = [[1, 1, 1, 1], [2, 2, 2, 2]]
-
-        return unet_config
-
     unet_config = {
         "use_checkpoint": False,
         "image_size": 32,
@@ -75,6 +45,7 @@ def detect_unet_config(state_dict, key_prefix):
     else:
         unet_config["adm_in_channels"] = None
 
+    unet_config["dtype"] = dtype
     model_channels = state_dict['{}input_blocks.0.0.weight'.format(key_prefix)].shape[0]
     in_channels = state_dict['{}input_blocks.0.0.weight'.format(key_prefix)].shape[1]
 
@@ -152,10 +123,8 @@ def detect_unet_config(state_dict, key_prefix):
     channel_mult.append(last_channel_mult)
     if "{}middle_block.1.proj_in.weight".format(key_prefix) in state_dict_keys:
         transformer_depth_middle = count_blocks(state_dict_keys, '{}middle_block.1.transformer_blocks.'.format(key_prefix) + '{}')
-    elif "{}middle_block.0.in_layers.0.weight".format(key_prefix) in state_dict_keys:
-        transformer_depth_middle = -1
     else:
-        transformer_depth_middle = -2
+        transformer_depth_middle = -1
 
     unet_config["in_channels"] = in_channels
     unet_config["out_channels"] = out_channels
@@ -182,17 +151,17 @@ def detect_unet_config(state_dict, key_prefix):
 
     return unet_config
 
-def model_config_from_unet_config(unet_config, state_dict=None):
+def model_config_from_unet_config(unet_config):
     for model_config in ldm_patched.modules.supported_models.models:
-        if model_config.matches(unet_config, state_dict):
+        if model_config.matches(unet_config):
             return model_config(unet_config)
 
-    logging.error("no match {}".format(unet_config))
+    print("no match", unet_config)
     return None
 
-def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=False):
-    unet_config = detect_unet_config(state_dict, unet_key_prefix)
-    model_config = model_config_from_unet_config(unet_config, state_dict)
+def model_config_from_unet(state_dict, unet_key_prefix, dtype, use_base_if_no_match=False):
+    unet_config = detect_unet_config(state_dict, unet_key_prefix, dtype)
+    model_config = model_config_from_unet_config(unet_config)
     if model_config is None and use_base_if_no_match:
         return ldm_patched.modules.supported_models_base.BASE(unet_config)
     else:
@@ -237,7 +206,7 @@ def convert_config(unet_config):
     return new_config
 
 
-def unet_config_from_diffusers_unet(state_dict, dtype=None):
+def unet_config_from_diffusers_unet(state_dict, dtype):
     match = {}
     transformer_depth = []
 
@@ -245,7 +214,6 @@ def unet_config_from_diffusers_unet(state_dict, dtype=None):
     down_blocks = count_blocks(state_dict, "down_blocks.{}")
     for i in range(down_blocks):
         attn_blocks = count_blocks(state_dict, "down_blocks.{}.attentions.".format(i) + '{}')
-        res_blocks = count_blocks(state_dict, "down_blocks.{}.resnets.".format(i) + '{}')
         for ab in range(attn_blocks):
             transformer_count = count_blocks(state_dict, "down_blocks.{}.attentions.{}.transformer_blocks.".format(i, ab) + '{}')
             transformer_depth.append(transformer_count)
@@ -254,8 +222,8 @@ def unet_config_from_diffusers_unet(state_dict, dtype=None):
 
         attn_res *= 2
         if attn_blocks == 0:
-            for i in range(res_blocks):
-                transformer_depth.append(0)
+            transformer_depth.append(0)
+            transformer_depth.append(0)
 
     match["transformer_depth"] = transformer_depth
 
@@ -321,12 +289,6 @@ def unet_config_from_diffusers_unet(state_dict, dtype=None):
                               'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10],
                               'use_temporal_attention': False, 'use_temporal_resblock': False}
 
-    SDXL_diffusers_ip2p = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
-                              'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 8, 'model_channels': 320,
-                              'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 10, 10], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 10,
-                              'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10],
-                              'use_temporal_attention': False, 'use_temporal_resblock': False}
-
     SSD_1B = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
               'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,
               'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 4, 4], 'transformer_depth_output': [0, 0, 0, 1, 1, 2, 10, 4, 4],
@@ -339,32 +301,7 @@ def unet_config_from_diffusers_unet(state_dict, dtype=None):
               'channel_mult': [1, 2, 4], 'transformer_depth_middle': -1, 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64,
               'use_temporal_attention': False, 'use_temporal_resblock': False}
 
-    KOALA_700M = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
-              'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,
-              'num_res_blocks': [1, 1, 1], 'transformer_depth': [0, 2, 5], 'transformer_depth_output': [0, 0, 2, 2, 5, 5],
-              'channel_mult': [1, 2, 4], 'transformer_depth_middle': -2, 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64,
-              'use_temporal_attention': False, 'use_temporal_resblock': False}
-
-    KOALA_1B = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
-              'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,
-              'num_res_blocks': [1, 1, 1], 'transformer_depth': [0, 2, 6], 'transformer_depth_output': [0, 0, 2, 2, 6, 6],
-              'channel_mult': [1, 2, 4], 'transformer_depth_middle': 6, 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64,
-              'use_temporal_attention': False, 'use_temporal_resblock': False}
-
-    SD09_XS = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
-            'adm_in_channels': None, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': [1, 1, 1],
-            'transformer_depth': [1, 1, 1], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': -2, 'use_linear_in_transformer': True,
-            'context_dim': 1024, 'num_head_channels': 64, 'transformer_depth_output': [1, 1, 1, 1, 1, 1],
-            'use_temporal_attention': False, 'use_temporal_resblock': False, 'disable_self_attentions': [True, False, False]}
-
-    SD_XS = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
-            'adm_in_channels': None, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': [1, 1, 1],
-            'transformer_depth': [0, 1, 1], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': -2, 'use_linear_in_transformer': False,
-            'context_dim': 768, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 1, 1, 1, 1],
-            'use_temporal_attention': False, 'use_temporal_resblock': False}
-
-
-    supported_models = [SDXL, SDXL_refiner, SD21, SD15, SD21_uncliph, SD21_unclipl, SDXL_mid_cnet, SDXL_small_cnet, SDXL_diffusers_inpaint, SSD_1B, Segmind_Vega, KOALA_700M, KOALA_1B, SD09_XS, SD_XS, SDXL_diffusers_ip2p]
+    supported_models = [SDXL, SDXL_refiner, SD21, SD15, SD21_uncliph, SD21_unclipl, SDXL_mid_cnet, SDXL_small_cnet, SDXL_diffusers_inpaint, SSD_1B, Segmind_Vega]
 
     for unet_config in supported_models:
         matches = True
@@ -376,8 +313,8 @@ def unet_config_from_diffusers_unet(state_dict, dtype=None):
             return convert_config(unet_config)
     return None
 
-def model_config_from_diffusers_unet(state_dict):
-    unet_config = unet_config_from_diffusers_unet(state_dict)
+def model_config_from_diffusers_unet(state_dict, dtype):
+    unet_config = unet_config_from_diffusers_unet(state_dict, dtype)
     if unet_config is not None:
         return model_config_from_unet_config(unet_config)
     return None

ldm_patched/modules/model_management.py

@@ -1,10 +1,9 @@
 import psutil
-import logging
 from enum import Enum
 from ldm_patched.modules.args_parser import args
+import ldm_patched.modules.utils
 import torch
 import sys
-import platform
 
 class VRAMState(Enum):
     DISABLED = 0    #No vram present: no need to move models to vram
@@ -29,8 +28,8 @@ total_vram = 0
 lowvram_available = True
 xpu_available = False
 
-if args.deterministic:
-    logging.info("Using deterministic algorithms for pytorch")
+if args.pytorch_deterministic:
+    print("Using deterministic algorithms for pytorch")
     torch.use_deterministic_algorithms(True, warn_only=True)
 
 directml_enabled = False
@@ -42,7 +41,7 @@ if args.directml is not None:
         directml_device = torch_directml.device()
     else:
         directml_device = torch_directml.device(device_index)
-    logging.info("Using directml with device: {}".format(torch_directml.device_name(device_index)))
+    print("Using directml with device:", torch_directml.device_name(device_index))
     # torch_directml.disable_tiled_resources(True)
     lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default.
 
@@ -60,7 +59,10 @@ try:
 except:
     pass
 
-if args.cpu:
+if args.always_cpu:
+    if args.always_cpu > 0:
+        torch.set_num_threads(args.always_cpu)
+    print(f"Running on {torch.get_num_threads()} CPU threads")
     cpu_state = CPUState.CPU
 
 def is_intel_xpu():
@@ -83,7 +85,7 @@ def get_torch_device():
         return torch.device("cpu")
     else:
         if is_intel_xpu():
-            return torch.device("xpu", torch.xpu.current_device())
+            return torch.device("xpu")
         else:
             return torch.device(torch.cuda.current_device())
 
@@ -102,8 +104,8 @@ def get_total_memory(dev=None, torch_total_too=False):
         elif is_intel_xpu():
             stats = torch.xpu.memory_stats(dev)
             mem_reserved = stats['reserved_bytes.all.current']
-            mem_total_torch = mem_reserved
             mem_total = torch.xpu.get_device_properties(dev).total_memory
+            mem_total_torch = mem_reserved
         else:
             stats = torch.cuda.memory_stats(dev)
             mem_reserved = stats['reserved_bytes.all.current']
@@ -118,12 +120,11 @@ def get_total_memory(dev=None, torch_total_too=False):
 
 total_vram = get_total_memory(get_torch_device()) / (1024 * 1024)
 total_ram = psutil.virtual_memory().total / (1024 * 1024)
-logging.info("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))
-
-try:
-    logging.info("pytorch version: {}".format(torch.version.__version__))
-except:
-    pass
+print("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))
+if not args.always_normal_vram and not args.always_cpu:
+    if lowvram_available and total_vram <= 4096:
+        print("Trying to enable lowvram mode because your GPU seems to have 4GB or less. If you don't want this use: --always-normal-vram")
+        set_vram_to = VRAMState.LOW_VRAM
 
 try:
     OOM_EXCEPTION = torch.cuda.OutOfMemoryError
@@ -145,10 +146,12 @@ else:
             pass
         try:
             XFORMERS_VERSION = xformers.version.__version__
-            logging.info("xformers version: {}".format(XFORMERS_VERSION))
+            print("xformers version:", XFORMERS_VERSION)
             if XFORMERS_VERSION.startswith("0.0.18"):
-                logging.warning("\nWARNING: This version of xformers has a major bug where you will get black images when generating high resolution images.")
-                logging.warning("Please downgrade or upgrade xformers to a different version.\n")
+                print()
+                print("WARNING: This version of xformers has a major bug where you will get black images when generating high resolution images.")
+                print("Please downgrade or upgrade xformers to a different version.")
+                print()
                 XFORMERS_ENABLED_VAE = False
         except:
             pass
@@ -163,7 +166,7 @@ def is_nvidia():
     return False
 
 ENABLE_PYTORCH_ATTENTION = False
-if args.use_pytorch_cross_attention:
+if args.attention_pytorch:
     ENABLE_PYTORCH_ATTENTION = True
     XFORMERS_IS_AVAILABLE = False
 
@@ -173,12 +176,12 @@ try:
     if is_nvidia():
         torch_version = torch.version.__version__
         if int(torch_version[0]) >= 2:
-            if ENABLE_PYTORCH_ATTENTION == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
+            if ENABLE_PYTORCH_ATTENTION == False and args.attention_split == False and args.attention_quad == False:
                 ENABLE_PYTORCH_ATTENTION = True
             if torch.cuda.is_bf16_supported() and torch.cuda.get_device_properties(torch.cuda.current_device()).major >= 8:
                 VAE_DTYPE = torch.bfloat16
     if is_intel_xpu():
-        if args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
+        if args.attention_split == False and args.attention_quad == False:
             ENABLE_PYTORCH_ATTENTION = True
 except:
     pass
@@ -186,14 +189,14 @@ except:
 if is_intel_xpu():
     VAE_DTYPE = torch.bfloat16
 
-if args.cpu_vae:
+if args.vae_in_cpu:
     VAE_DTYPE = torch.float32
 
-if args.fp16_vae:
+if args.vae_in_fp16:
     VAE_DTYPE = torch.float16
-elif args.bf16_vae:
+elif args.vae_in_bf16:
     VAE_DTYPE = torch.bfloat16
-elif args.fp32_vae:
+elif args.vae_in_fp32:
     VAE_DTYPE = torch.float32
 
 
@@ -202,22 +205,22 @@ if ENABLE_PYTORCH_ATTENTION:
     torch.backends.cuda.enable_flash_sdp(True)
     torch.backends.cuda.enable_mem_efficient_sdp(True)
 
-if args.lowvram:
+if args.always_low_vram:
     set_vram_to = VRAMState.LOW_VRAM
     lowvram_available = True
-elif args.novram:
+elif args.always_no_vram:
     set_vram_to = VRAMState.NO_VRAM
-elif args.highvram or args.gpu_only:
+elif args.always_high_vram or args.always_gpu:
     vram_state = VRAMState.HIGH_VRAM
 
 FORCE_FP32 = False
 FORCE_FP16 = False
-if args.force_fp32:
-    logging.info("Forcing FP32, if this improves things please report it.")
+if args.all_in_fp32:
+    print("Forcing FP32, if this improves things please report it.")
     FORCE_FP32 = True
 
-if args.force_fp16:
-    logging.info("Forcing FP16.")
+if args.all_in_fp16:
+    print("Forcing FP16.")
     FORCE_FP16 = True
 
 if lowvram_available:
@@ -231,12 +234,12 @@ if cpu_state != CPUState.GPU:
 if cpu_state == CPUState.MPS:
     vram_state = VRAMState.SHARED
 
-logging.info(f"Set vram state to: {vram_state.name}")
+print(f"Set vram state to: {vram_state.name}")
 
-DISABLE_SMART_MEMORY = args.disable_smart_memory
+ALWAYS_VRAM_OFFLOAD = args.always_offload_from_vram
 
-if DISABLE_SMART_MEMORY:
-    logging.info("Disabling smart memory management")
+if ALWAYS_VRAM_OFFLOAD:
+    print("Always offload VRAM")
 
 def get_torch_device_name(device):
     if hasattr(device, 'type'):
@@ -254,11 +257,11 @@ def get_torch_device_name(device):
         return "CUDA {}: {}".format(device, torch.cuda.get_device_name(device))
 
 try:
-    logging.info("Device: {}".format(get_torch_device_name(get_torch_device())))
+    print("Device:", get_torch_device_name(get_torch_device()))
 except:
-    logging.warning("Could not pick default device.")
+    print("Could not pick default device.")
 
-logging.info("VAE dtype: {}".format(VAE_DTYPE))
+print("VAE dtype:", VAE_DTYPE)
 
 current_loaded_models = []
 
@@ -273,9 +276,8 @@ def module_size(module):
 class LoadedModel:
     def __init__(self, model):
         self.model = model
+        self.model_accelerated = False
         self.device = model.load_device
-        self.weights_loaded = False
-        self.real_model = None
 
     def model_memory(self):
         return self.model.model_size()
@@ -286,40 +288,55 @@ class LoadedModel:
         else:
             return self.model_memory()
 
-    def model_load(self, lowvram_model_memory=0, force_patch_weights=False):
-        patch_model_to = self.device
+    def model_load(self, lowvram_model_memory=0):
+        patch_model_to = None
+        if lowvram_model_memory == 0:
+            patch_model_to = self.device
 
         self.model.model_patches_to(self.device)
         self.model.model_patches_to(self.model.model_dtype())
 
-        load_weights = not self.weights_loaded
-
         try:
-            if lowvram_model_memory > 0 and load_weights:
-                self.real_model = self.model.patch_model_lowvram(device_to=patch_model_to, lowvram_model_memory=lowvram_model_memory, force_patch_weights=force_patch_weights)
-            else:
-                self.real_model = self.model.patch_model(device_to=patch_model_to, patch_weights=load_weights)
+            self.real_model = self.model.patch_model(device_to=patch_model_to) #TODO: do something with loras and offloading to CPU
         except Exception as e:
             self.model.unpatch_model(self.model.offload_device)
             self.model_unload()
             raise e
 
-        if is_intel_xpu() and not args.disable_ipex_optimize:
-            self.real_model = ipex.optimize(self.real_model.eval(), graph_mode=True, concat_linear=True)
+        if lowvram_model_memory > 0:
+            print("loading in lowvram mode", lowvram_model_memory/(1024 * 1024))
+            mem_counter = 0
+            for m in self.real_model.modules():
+                if hasattr(m, "ldm_patched_cast_weights"):
+                    m.prev_ldm_patched_cast_weights = m.ldm_patched_cast_weights
+                    m.ldm_patched_cast_weights = True
+                    module_mem = module_size(m)
+                    if mem_counter + module_mem < lowvram_model_memory:
+                        m.to(self.device)
+                        mem_counter += module_mem
+                elif hasattr(m, "weight"): #only modules with ldm_patched_cast_weights can be set to lowvram mode
+                    m.to(self.device)
+                    mem_counter += module_size(m)
+                    print("lowvram: loaded module regularly", m)
+
+            self.model_accelerated = True
+
+        if is_intel_xpu() and not args.disable_ipex_hijack:
+            self.real_model = torch.xpu.optimize(self.real_model.eval(), inplace=True, auto_kernel_selection=True, graph_mode=True)
 
-        self.weights_loaded = True
         return self.real_model
 
-    def should_reload_model(self, force_patch_weights=False):
-        if force_patch_weights and self.model.lowvram_patch_counter > 0:
-            return True
-        return False
+    def model_unload(self):
+        if self.model_accelerated:
+            for m in self.real_model.modules():
+                if hasattr(m, "prev_ldm_patched_cast_weights"):
+                    m.ldm_patched_cast_weights = m.prev_ldm_patched_cast_weights
+                    del m.prev_ldm_patched_cast_weights
 
-    def model_unload(self, unpatch_weights=True):
-        self.model.unpatch_model(self.model.offload_device, unpatch_weights=unpatch_weights)
+            self.model_accelerated = False
+
+        self.model.unpatch_model(self.model.offload_device)
         self.model.model_patches_to(self.model.offload_device)
-        self.weights_loaded = self.weights_loaded and not unpatch_weights
-        self.real_model = None
 
     def __eq__(self, other):
         return self.model is other.model
@@ -327,57 +344,31 @@ class LoadedModel:
 def minimum_inference_memory():
     return (1024 * 1024 * 1024)
 
-def unload_model_clones(model, unload_weights_only=True, force_unload=True):
+def unload_model_clones(model):
     to_unload = []
     for i in range(len(current_loaded_models)):
         if model.is_clone(current_loaded_models[i].model):
             to_unload = [i] + to_unload
 
-    if len(to_unload) == 0:
-        return True
-
-    same_weights = 0
     for i in to_unload:
-        if model.clone_has_same_weights(current_loaded_models[i].model):
-            same_weights += 1
-
-    if same_weights == len(to_unload):
-        unload_weight = False
-    else:
-        unload_weight = True
-
-    if not force_unload:
-        if unload_weights_only and unload_weight == False:
-            return None
-
-    for i in to_unload:
-        logging.debug("unload clone {} {}".format(i, unload_weight))
-        current_loaded_models.pop(i).model_unload(unpatch_weights=unload_weight)
-
-    return unload_weight
+        print("unload clone", i)
+        current_loaded_models.pop(i).model_unload()
 
 def free_memory(memory_required, device, keep_loaded=[]):
-    unloaded_model = []
-    can_unload = []
-
+    unloaded_model = False
     for i in range(len(current_loaded_models) -1, -1, -1):
+        if not ALWAYS_VRAM_OFFLOAD:
+            if get_free_memory(device) > memory_required:
+                break
         shift_model = current_loaded_models[i]
         if shift_model.device == device:
             if shift_model not in keep_loaded:
-                can_unload.append((sys.getrefcount(shift_model.model), shift_model.model_memory(), i))
+                m = current_loaded_models.pop(i)
+                m.model_unload()
+                del m
+                unloaded_model = True
 
-    for x in sorted(can_unload):
-        i = x[-1]
-        if not DISABLE_SMART_MEMORY:
-            if get_free_memory(device) > memory_required:
-                break
-        current_loaded_models[i].model_unload()
-        unloaded_model.append(i)
-
-    for i in sorted(unloaded_model, reverse=True):
-        current_loaded_models.pop(i)
-
-    if len(unloaded_model) > 0:
+    if unloaded_model:
         soft_empty_cache()
     else:
         if vram_state != VRAMState.HIGH_VRAM:
@@ -385,36 +376,24 @@ def free_memory(memory_required, device, keep_loaded=[]):
             if mem_free_torch > mem_free_total * 0.25:
                 soft_empty_cache()
 
-def load_models_gpu(models, memory_required=0, force_patch_weights=False):
+def load_models_gpu(models, memory_required=0):
     global vram_state
 
     inference_memory = minimum_inference_memory()
     extra_mem = max(inference_memory, memory_required)
 
-    models = set(models)
-
     models_to_load = []
     models_already_loaded = []
     for x in models:
         loaded_model = LoadedModel(x)
-        loaded = None
 
-        try:
-            loaded_model_index = current_loaded_models.index(loaded_model)
-        except:
-            loaded_model_index = None
-
-        if loaded_model_index is not None:
-            loaded = current_loaded_models[loaded_model_index]
-            if loaded.should_reload_model(force_patch_weights=force_patch_weights): #TODO: cleanup this model reload logic
-                current_loaded_models.pop(loaded_model_index).model_unload(unpatch_weights=True)
-                loaded = None
-            else:
-                models_already_loaded.append(loaded)
-
-        if loaded is None:
+        if loaded_model in current_loaded_models:
+            index = current_loaded_models.index(loaded_model)
+            current_loaded_models.insert(0, current_loaded_models.pop(index))
+            models_already_loaded.append(loaded_model)
+        else:
             if hasattr(x, "model"):
-                logging.info(f"Requested to load {x.model.__class__.__name__}")
+                print(f"Requested to load {x.model.__class__.__name__}")
             models_to_load.append(loaded_model)
 
     if len(models_to_load) == 0:
@@ -424,22 +403,17 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False):
                 free_memory(extra_mem, d, models_already_loaded)
         return
 
-    logging.info(f"Loading {len(models_to_load)} new model{'s' if len(models_to_load) > 1 else ''}")
+    print(f"Loading {len(models_to_load)} new model{'s' if len(models_to_load) > 1 else ''}")
 
     total_memory_required = {}
     for loaded_model in models_to_load:
-        if unload_model_clones(loaded_model.model, unload_weights_only=True, force_unload=False) == True:#unload clones where the weights are different
-            total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device)
+        unload_model_clones(loaded_model.model)
+        total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device)
 
     for device in total_memory_required:
         if device != torch.device("cpu"):
             free_memory(total_memory_required[device] * 1.3 + extra_mem, device, models_already_loaded)
 
-    for loaded_model in models_to_load:
-        weights_unloaded = unload_model_clones(loaded_model.model, unload_weights_only=False, force_unload=False) #unload the rest of the clones where the weights can stay loaded
-        if weights_unloaded is not None:
-            loaded_model.weights_loaded = not weights_unloaded
-
     for loaded_model in models_to_load:
         model = loaded_model.model
         torch_dev = model.load_device
@@ -452,13 +426,15 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False):
             model_size = loaded_model.model_memory_required(torch_dev)
             current_free_mem = get_free_memory(torch_dev)
             lowvram_model_memory = int(max(64 * (1024 * 1024), (current_free_mem - 1024 * (1024 * 1024)) / 1.3 ))
-            if model_size <= (current_free_mem - inference_memory): #only switch to lowvram if really necessary
+            if model_size > (current_free_mem - inference_memory): #only switch to lowvram if really necessary
+                vram_set_state = VRAMState.LOW_VRAM
+            else:
                 lowvram_model_memory = 0
 
         if vram_set_state == VRAMState.NO_VRAM:
             lowvram_model_memory = 64 * 1024 * 1024
 
-        cur_loaded_model = loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights)
+        cur_loaded_model = loaded_model.model_load(lowvram_model_memory)
         current_loaded_models.insert(0, loaded_model)
     return
 
@@ -466,15 +442,11 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False):
 def load_model_gpu(model):
     return load_models_gpu([model])
 
-def cleanup_models(keep_clone_weights_loaded=False):
+def cleanup_models():
     to_delete = []
     for i in range(len(current_loaded_models)):
         if sys.getrefcount(current_loaded_models[i].model) <= 2:
-            if not keep_clone_weights_loaded:
-                to_delete = [i] + to_delete
-            #TODO: find a less fragile way to do this.
-            elif sys.getrefcount(current_loaded_models[i].real_model) <= 3: #references from .real_model + the .model
-                to_delete = [i] + to_delete
+            to_delete = [i] + to_delete
 
     for i in to_delete:
         x = current_loaded_models.pop(i)
@@ -506,7 +478,7 @@ def unet_inital_load_device(parameters, dtype):
         return torch_dev
 
     cpu_dev = torch.device("cpu")
-    if DISABLE_SMART_MEMORY:
+    if ALWAYS_VRAM_OFFLOAD:
         return cpu_dev
 
     model_size = dtype_size(dtype) * parameters
@@ -518,54 +490,45 @@ def unet_inital_load_device(parameters, dtype):
     else:
         return cpu_dev
 
-def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]):
-    if args.bf16_unet:
+def unet_dtype(device=None, model_params=0):
+    if args.unet_in_bf16:
         return torch.bfloat16
-    if args.fp16_unet:
+    if args.unet_in_fp16:
         return torch.float16
-    if args.fp8_e4m3fn_unet:
+    if args.unet_in_fp8_e4m3fn:
         return torch.float8_e4m3fn
-    if args.fp8_e5m2_unet:
+    if args.unet_in_fp8_e5m2:
         return torch.float8_e5m2
-    if should_use_fp16(device=device, model_params=model_params, manual_cast=True):
-        if torch.float16 in supported_dtypes:
-            return torch.float16
-    if should_use_bf16(device, model_params=model_params, manual_cast=True):
-        if torch.bfloat16 in supported_dtypes:
-            return torch.bfloat16
+    if should_use_fp16(device=device, model_params=model_params):
+        return torch.float16
     return torch.float32
 
 # None means no manual cast
-def unet_manual_cast(weight_dtype, inference_device, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]):
+def unet_manual_cast(weight_dtype, inference_device):
     if weight_dtype == torch.float32:
         return None
 
-    fp16_supported = should_use_fp16(inference_device, prioritize_performance=False)
+    fp16_supported = ldm_patched.modules.model_management.should_use_fp16(inference_device, prioritize_performance=False)
     if fp16_supported and weight_dtype == torch.float16:
         return None
 
-    bf16_supported = should_use_bf16(inference_device)
-    if bf16_supported and weight_dtype == torch.bfloat16:
-        return None
-
-    if fp16_supported and torch.float16 in supported_dtypes:
+    if fp16_supported:
         return torch.float16
-
-    elif bf16_supported and torch.bfloat16 in supported_dtypes:
-        return torch.bfloat16
     else:
         return torch.float32
 
 def text_encoder_offload_device():
-    if args.gpu_only:
+    if args.always_gpu:
         return get_torch_device()
     else:
         return torch.device("cpu")
 
 def text_encoder_device():
-    if args.gpu_only:
+    if args.always_gpu:
         return get_torch_device()
     elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM:
+        if is_intel_xpu():
+            return torch.device("cpu")
         if should_use_fp16(prioritize_performance=False):
             return get_torch_device()
         else:
@@ -574,34 +537,36 @@ def text_encoder_device():
         return torch.device("cpu")
 
 def text_encoder_dtype(device=None):
-    if args.fp8_e4m3fn_text_enc:
+    if args.clip_in_fp8_e4m3fn:
         return torch.float8_e4m3fn
-    elif args.fp8_e5m2_text_enc:
+    elif args.clip_in_fp8_e5m2:
         return torch.float8_e5m2
-    elif args.fp16_text_enc:
+    elif args.clip_in_fp16:
         return torch.float16
-    elif args.fp32_text_enc:
+    elif args.clip_in_fp32:
         return torch.float32
 
     if is_device_cpu(device):
         return torch.float16
 
-    return torch.float16
-
+    if should_use_fp16(device, prioritize_performance=False):
+        return torch.float16
+    else:
+        return torch.float32
 
 def intermediate_device():
-    if args.gpu_only:
+    if args.always_gpu:
         return get_torch_device()
     else:
         return torch.device("cpu")
 
 def vae_device():
-    if args.cpu_vae:
+    if args.vae_in_cpu:
         return torch.device("cpu")
     return get_torch_device()
 
 def vae_offload_device():
-    if args.gpu_only:
+    if args.always_gpu:
         return get_torch_device()
     else:
         return torch.device("cpu")
@@ -629,19 +594,8 @@ def supports_dtype(device, dtype): #TODO
 def device_supports_non_blocking(device):
     if is_device_mps(device):
         return False #pytorch bug? mps doesn't support non blocking
-    if args.deterministic: #TODO: figure out why deterministic breaks non blocking from gpu to cpu (previews)
-        return False
-    if directml_enabled:
-        return False
     return True
 
-def device_should_use_non_blocking(device):
-    if not device_supports_non_blocking(device):
-        return False
-    return False
-    # return True #TODO: figure out why this causes memory issues on Nvidia and possibly others
-
-
 def cast_to_device(tensor, device, dtype, copy=False):
     device_supports_cast = False
     if tensor.dtype == torch.float32 or tensor.dtype == torch.float16:
@@ -652,7 +606,7 @@ def cast_to_device(tensor, device, dtype, copy=False):
         elif is_intel_xpu():
             device_supports_cast = True
 
-    non_blocking = device_should_use_non_blocking(device)
+    non_blocking = device_supports_non_blocking(device)
 
     if device_supports_cast:
         if copy:
@@ -695,18 +649,6 @@ def pytorch_attention_flash_attention():
             return True
     return False
 
-def force_upcast_attention_dtype():
-    upcast = args.force_upcast_attention
-    try:
-        if platform.mac_ver()[0] in ['14.5']: #black image bug on OSX Sonoma 14.5
-            upcast = True
-    except:
-        pass
-    if upcast:
-        return torch.float32
-    else:
-        return None
-
 def get_free_memory(dev=None, torch_free_too=False):
     global directml_enabled
     if dev is None:
@@ -722,10 +664,10 @@ def get_free_memory(dev=None, torch_free_too=False):
         elif is_intel_xpu():
             stats = torch.xpu.memory_stats(dev)
             mem_active = stats['active_bytes.all.current']
+            mem_allocated = stats['allocated_bytes.all.current']
             mem_reserved = stats['reserved_bytes.all.current']
             mem_free_torch = mem_reserved - mem_active
-            mem_free_xpu = torch.xpu.get_device_properties(dev).total_memory - mem_reserved
-            mem_free_total = mem_free_xpu + mem_free_torch
+            mem_free_total = torch.xpu.get_device_properties(dev).total_memory - mem_allocated
         else:
             stats = torch.cuda.memory_stats(dev)
             mem_active = stats['active_bytes.all.current']
@@ -747,22 +689,19 @@ def mps_mode():
     global cpu_state
     return cpu_state == CPUState.MPS
 
-def is_device_type(device, type):
+def is_device_cpu(device):
     if hasattr(device, 'type'):
-        if (device.type == type):
+        if (device.type == 'cpu'):
             return True
     return False
 
-def is_device_cpu(device):
-    return is_device_type(device, 'cpu')
-
 def is_device_mps(device):
-    return is_device_type(device, 'mps')
+    if hasattr(device, 'type'):
+        if (device.type == 'mps'):
+            return True
+    return False
 
-def is_device_cuda(device):
-    return is_device_type(device, 'cuda')
-
-def should_use_fp16(device=None, model_params=0, prioritize_performance=True, manual_cast=False):
+def should_use_fp16(device=None, model_params=0, prioritize_performance=True):
     global directml_enabled
 
     if device is not None:
@@ -772,9 +711,9 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma
     if FORCE_FP16:
         return True
 
-    if device is not None:
+    if device is not None: #TODO
         if is_device_mps(device):
-            return True
+            return False
 
     if FORCE_FP32:
         return False
@@ -782,22 +721,16 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma
     if directml_enabled:
         return False
 
-    if mps_mode():
-        return True
-
-    if cpu_mode():
-        return False
+    if cpu_mode() or mps_mode():
+        return False #TODO ?
 
     if is_intel_xpu():
         return True
 
-    if torch.version.hip:
+    if torch.cuda.is_bf16_supported():
         return True
 
     props = torch.cuda.get_device_properties("cuda")
-    if props.major >= 8:
-        return True
-
     if props.major < 6:
         return False
 
@@ -805,12 +738,12 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma
     #FP16 is confirmed working on a 1080 (GP104) but it's a bit slower than FP32 so it should only be enabled
     #when the model doesn't actually fit on the card
     #TODO: actually test if GP106 and others have the same type of behavior
-    nvidia_10_series = ["1080", "1070", "titan x", "p3000", "p3200", "p4000", "p4200", "p5000", "p5200", "p6000", "1060", "1050", "p40", "p100", "p6", "p4"]
+    nvidia_10_series = ["1080", "1070", "titan x", "p3000", "p3200", "p4000", "p4200", "p5000", "p5200", "p6000", "1060", "1050"]
     for x in nvidia_10_series:
         if x in props.name.lower():
             fp16_works = True
 
-    if fp16_works or manual_cast:
+    if fp16_works:
         free_model_memory = (get_free_memory() * 0.9 - minimum_inference_memory())
         if (not prioritize_performance) or model_params * 4 > free_model_memory:
             return True
@@ -826,43 +759,6 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma
 
     return True
 
-def should_use_bf16(device=None, model_params=0, prioritize_performance=True, manual_cast=False):
-    if device is not None:
-        if is_device_cpu(device): #TODO ? bf16 works on CPU but is extremely slow
-            return False
-
-    if device is not None: #TODO not sure about mps bf16 support
-        if is_device_mps(device):
-            return False
-
-    if FORCE_FP32:
-        return False
-
-    if directml_enabled:
-        return False
-
-    if cpu_mode() or mps_mode():
-        return False
-
-    if is_intel_xpu():
-        return True
-
-    if device is None:
-        device = torch.device("cuda")
-
-    props = torch.cuda.get_device_properties(device)
-    if props.major >= 8:
-        return True
-
-    bf16_works = torch.cuda.is_bf16_supported()
-
-    if bf16_works or manual_cast:
-        free_model_memory = (get_free_memory() * 0.9 - minimum_inference_memory())
-        if (not prioritize_performance) or model_params * 4 > free_model_memory:
-            return True
-
-    return False
-
 def soft_empty_cache(force=False):
     global cpu_state
     if cpu_state == CPUState.MPS:
@@ -879,7 +775,6 @@ def unload_all_models():
 
 
 def resolve_lowvram_weight(weight, model, key): #TODO: remove
-    print("WARNING: The comfy.model_management.resolve_lowvram_weight function will be removed soon, please stop using it.")
     return weight
 
 #TODO: might be cleaner to put this somewhere else

ldm_patched/modules/model_patcher.py

@@ -1,55 +1,9 @@
 import torch
 import copy
 import inspect
-import logging
-import uuid
 
 import ldm_patched.modules.utils
 import ldm_patched.modules.model_management
-from ldm_patched.modules.types import UnetWrapperFunction
-
-
-def weight_decompose(dora_scale, weight, lora_diff, alpha, strength):
-    dora_scale = ldm_patched.modules.model_management.cast_to_device(dora_scale, weight.device, torch.float32)
-    lora_diff *= alpha
-    weight_calc = weight + lora_diff.type(weight.dtype)
-    weight_norm = (
-        weight_calc.transpose(0, 1)
-        .reshape(weight_calc.shape[1], -1)
-        .norm(dim=1, keepdim=True)
-        .reshape(weight_calc.shape[1], *[1] * (weight_calc.dim() - 1))
-        .transpose(0, 1)
-    )
-
-    weight_calc *= (dora_scale / weight_norm).type(weight.dtype)
-    if strength != 1.0:
-        weight_calc -= weight
-        weight += strength * (weight_calc)
-    else:
-        weight[:] = weight_calc
-    return weight
-
-
-def set_model_options_patch_replace(model_options, patch, name, block_name, number, transformer_index=None):
-    to = model_options["transformer_options"].copy()
-
-    if "patches_replace" not in to:
-        to["patches_replace"] = {}
-    else:
-        to["patches_replace"] = to["patches_replace"].copy()
-
-    if name not in to["patches_replace"]:
-        to["patches_replace"][name] = {}
-    else:
-        to["patches_replace"][name] = to["patches_replace"][name].copy()
-
-    if transformer_index is not None:
-        block = (block_name, number, transformer_index)
-    else:
-        block = (block_name, number)
-    to["patches_replace"][name][block] = patch
-    model_options["transformer_options"] = to
-    return model_options
 
 class ModelPatcher:
     def __init__(self, model, load_device, offload_device, size=0, current_device=None, weight_inplace_update=False):
@@ -69,9 +23,6 @@ class ModelPatcher:
             self.current_device = current_device
 
         self.weight_inplace_update = weight_inplace_update
-        self.model_lowvram = False
-        self.lowvram_patch_counter = 0
-        self.patches_uuid = uuid.uuid4()
 
     def model_size(self):
         if self.size > 0:
@@ -86,13 +37,10 @@ class ModelPatcher:
         n.patches = {}
         for k in self.patches:
             n.patches[k] = self.patches[k][:]
-        n.patches_uuid = self.patches_uuid
 
         n.object_patches = self.object_patches.copy()
         n.model_options = copy.deepcopy(self.model_options)
         n.model_keys = self.model_keys
-        n.backup = self.backup
-        n.object_patches_backup = self.object_patches_backup
         return n
 
     def is_clone(self, other):
@@ -100,19 +48,6 @@ class ModelPatcher:
             return True
         return False
 
-    def clone_has_same_weights(self, clone):
-        if not self.is_clone(clone):
-            return False
-
-        if len(self.patches) == 0 and len(clone.patches) == 0:
-            return True
-
-        if self.patches_uuid == clone.patches_uuid:
-            if len(self.patches) != len(clone.patches):
-                logging.warning("WARNING: something went wrong, same patch uuid but different length of patches.")
-            else:
-                return True
-
     def memory_required(self, input_shape):
         return self.model.memory_required(input_shape=input_shape)
 
@@ -129,12 +64,9 @@ class ModelPatcher:
         if disable_cfg1_optimization:
             self.model_options["disable_cfg1_optimization"] = True
 
-    def set_model_unet_function_wrapper(self, unet_wrapper_function: UnetWrapperFunction):
+    def set_model_unet_function_wrapper(self, unet_wrapper_function):
         self.model_options["model_function_wrapper"] = unet_wrapper_function
 
-    def set_model_denoise_mask_function(self, denoise_mask_function):
-        self.model_options["denoise_mask_function"] = denoise_mask_function
-
     def set_model_patch(self, patch, name):
         to = self.model_options["transformer_options"]
         if "patches" not in to:
@@ -142,7 +74,16 @@ class ModelPatcher:
         to["patches"][name] = to["patches"].get(name, []) + [patch]
 
     def set_model_patch_replace(self, patch, name, block_name, number, transformer_index=None):
-        self.model_options = set_model_options_patch_replace(self.model_options, patch, name, block_name, number, transformer_index=transformer_index)
+        to = self.model_options["transformer_options"]
+        if "patches_replace" not in to:
+            to["patches_replace"] = {}
+        if name not in to["patches_replace"]:
+            to["patches_replace"][name] = {}
+        if transformer_index is not None:
+            block = (block_name, number, transformer_index)
+        else:
+            block = (block_name, number)
+        to["patches_replace"][name][block] = patch
 
     def set_model_attn1_patch(self, patch):
         self.set_model_patch(patch, "attn1_patch")
@@ -174,15 +115,6 @@ class ModelPatcher:
     def add_object_patch(self, name, obj):
         self.object_patches[name] = obj
 
-    def get_model_object(self, name):
-        if name in self.object_patches:
-            return self.object_patches[name]
-        else:
-            if name in self.object_patches_backup:
-                return self.object_patches_backup[name]
-            else:
-                return ldm_patched.modules.utils.get_attr(self.model, name)
-
     def model_patches_to(self, device):
         to = self.model_options["transformer_options"]
         if "patches" in to:
@@ -217,7 +149,6 @@ class ModelPatcher:
                 current_patches.append((strength_patch, patches[k], strength_model))
                 self.patches[k] = current_patches
 
-        self.patches_uuid = uuid.uuid4()
         return list(p)
 
     def get_key_patches(self, filter_prefix=None):
@@ -243,41 +174,37 @@ class ModelPatcher:
                     sd.pop(k)
         return sd
 
-    def patch_weight_to_device(self, key, device_to=None):
-        if key not in self.patches:
-            return
-
-        weight = ldm_patched.modules.utils.get_attr(self.model, key)
-
-        inplace_update = self.weight_inplace_update
-
-        if key not in self.backup:
-            self.backup[key] = weight.to(device=self.offload_device, copy=inplace_update)
-
-        if device_to is not None:
-            temp_weight = ldm_patched.modules.model_management.cast_to_device(weight, device_to, torch.float32, copy=True)
-        else:
-            temp_weight = weight.to(torch.float32, copy=True)
-        out_weight = self.calculate_weight(self.patches[key], temp_weight, key).to(weight.dtype)
-        if inplace_update:
-            ldm_patched.modules.utils.copy_to_param(self.model, key, out_weight)
-        else:
-            ldm_patched.modules.utils.set_attr_param(self.model, key, out_weight)
-
     def patch_model(self, device_to=None, patch_weights=True):
         for k in self.object_patches:
-            old = ldm_patched.modules.utils.set_attr(self.model, k, self.object_patches[k])
+            old = getattr(self.model, k)
             if k not in self.object_patches_backup:
                 self.object_patches_backup[k] = old
+            setattr(self.model, k, self.object_patches[k])
 
         if patch_weights:
             model_sd = self.model_state_dict()
             for key in self.patches:
                 if key not in model_sd:
-                    logging.warning("could not patch. key doesn't exist in model: {}".format(key))
+                    print("could not patch. key doesn't exist in model:", key)
                     continue
 
-                self.patch_weight_to_device(key, device_to)
+                weight = model_sd[key]
+
+                inplace_update = self.weight_inplace_update
+
+                if key not in self.backup:
+                    self.backup[key] = weight.to(device=self.offload_device, copy=inplace_update)
+
+                if device_to is not None:
+                    temp_weight = ldm_patched.modules.model_management.cast_to_device(weight, device_to, torch.float32, copy=True)
+                else:
+                    temp_weight = weight.to(torch.float32, copy=True)
+                out_weight = self.calculate_weight(self.patches[key], temp_weight, key).to(weight.dtype)
+                if inplace_update:
+                    ldm_patched.modules.utils.copy_to_param(self.model, key, out_weight)
+                else:
+                    ldm_patched.modules.utils.set_attr(self.model, key, out_weight)
+                del temp_weight
 
             if device_to is not None:
                 self.model.to(device_to)
@@ -285,60 +212,9 @@ class ModelPatcher:
 
         return self.model
 
-    def patch_model_lowvram(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False):
-        self.patch_model(device_to, patch_weights=False)
-
-        logging.info("loading in lowvram mode {}".format(lowvram_model_memory/(1024 * 1024)))
-        class LowVramPatch:
-            def __init__(self, key, model_patcher):
-                self.key = key
-                self.model_patcher = model_patcher
-            def __call__(self, weight):
-                return self.model_patcher.calculate_weight(self.model_patcher.patches[self.key], weight, self.key)
-
-        mem_counter = 0
-        patch_counter = 0
-        for n, m in self.model.named_modules():
-            lowvram_weight = False
-            if hasattr(m, "comfy_cast_weights"):
-                module_mem = ldm_patched.modules.model_management.module_size(m)
-                if mem_counter + module_mem >= lowvram_model_memory:
-                    lowvram_weight = True
-
-            weight_key = "{}.weight".format(n)
-            bias_key = "{}.bias".format(n)
-
-            if lowvram_weight:
-                if weight_key in self.patches:
-                    if force_patch_weights:
-                        self.patch_weight_to_device(weight_key)
-                    else:
-                        m.weight_function = LowVramPatch(weight_key, self)
-                        patch_counter += 1
-                if bias_key in self.patches:
-                    if force_patch_weights:
-                        self.patch_weight_to_device(bias_key)
-                    else:
-                        m.bias_function = LowVramPatch(bias_key, self)
-                        patch_counter += 1
-
-                m.prev_comfy_cast_weights = m.comfy_cast_weights
-                m.comfy_cast_weights = True
-            else:
-                if hasattr(m, "weight"):
-                    self.patch_weight_to_device(weight_key, device_to)
-                    self.patch_weight_to_device(bias_key, device_to)
-                    m.to(device_to)
-                    mem_counter += ldm_patched.modules.model_management.module_size(m)
-                    logging.debug("lowvram: loaded module regularly {}".format(m))
-
-        self.model_lowvram = True
-        self.lowvram_patch_counter = patch_counter
-        return self.model
-
     def calculate_weight(self, patches, weight, key):
         for p in patches:
-            strength = p[0]
+            alpha = p[0]
             v = p[1]
             strength_model = p[2]
 
@@ -356,33 +232,25 @@ class ModelPatcher:
 
             if patch_type == "diff":
                 w1 = v[0]
-                if strength != 0.0:
+                if alpha != 0.0:
                     if w1.shape != weight.shape:
-                        logging.warning("WARNING SHAPE MISMATCH {} WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
+                        print("WARNING SHAPE MISMATCH {} WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
                     else:
-                        weight += strength * ldm_patched.modules.model_management.cast_to_device(w1, weight.device, weight.dtype)
+                        weight += alpha * ldm_patched.modules.model_management.cast_to_device(w1, weight.device, weight.dtype)
             elif patch_type == "lora": #lora/locon
                 mat1 = ldm_patched.modules.model_management.cast_to_device(v[0], weight.device, torch.float32)
                 mat2 = ldm_patched.modules.model_management.cast_to_device(v[1], weight.device, torch.float32)
-                dora_scale = v[4]
                 if v[2] is not None:
-                    alpha = v[2] / mat2.shape[0]
-                else:
-                    alpha = 1.0
-
+                    alpha *= v[2] / mat2.shape[0]
                 if v[3] is not None:
                     #locon mid weights, hopefully the math is fine because I didn't properly test it
                     mat3 = ldm_patched.modules.model_management.cast_to_device(v[3], weight.device, torch.float32)
                     final_shape = [mat2.shape[1], mat2.shape[0], mat3.shape[2], mat3.shape[3]]
                     mat2 = torch.mm(mat2.transpose(0, 1).flatten(start_dim=1), mat3.transpose(0, 1).flatten(start_dim=1)).reshape(final_shape).transpose(0, 1)
                 try:
-                    lora_diff = torch.mm(mat1.flatten(start_dim=1), mat2.flatten(start_dim=1)).reshape(weight.shape)
-                    if dora_scale is not None:
-                        weight = weight_decompose(dora_scale, weight, lora_diff, alpha, strength)
-                    else:
-                        weight += ((strength * alpha) * lora_diff).type(weight.dtype)
+                    weight += (alpha * torch.mm(mat1.flatten(start_dim=1), mat2.flatten(start_dim=1))).reshape(weight.shape).type(weight.dtype)
                 except Exception as e:
-                    logging.error("ERROR {} {} {}".format(patch_type, key, e))
+                    print("ERROR", key, e)
             elif patch_type == "lokr":
                 w1 = v[0]
                 w2 = v[1]
@@ -391,7 +259,6 @@ class ModelPatcher:
                 w2_a = v[5]
                 w2_b = v[6]
                 t2 = v[7]
-                dora_scale = v[8]
                 dim = None
 
                 if w1 is None:
@@ -417,29 +284,19 @@ class ModelPatcher:
                 if len(w2.shape) == 4:
                     w1 = w1.unsqueeze(2).unsqueeze(2)
                 if v[2] is not None and dim is not None:
-                    alpha = v[2] / dim
-                else:
-                    alpha = 1.0
+                    alpha *= v[2] / dim
 
                 try:
-                    lora_diff = torch.kron(w1, w2).reshape(weight.shape)
-                    if dora_scale is not None:
-                        weight = weight_decompose(dora_scale, weight, lora_diff, alpha, strength)
-                    else:
-                        weight += ((strength * alpha) * lora_diff).type(weight.dtype)
+                    weight += alpha * torch.kron(w1, w2).reshape(weight.shape).type(weight.dtype)
                 except Exception as e:
-                    logging.error("ERROR {} {} {}".format(patch_type, key, e))
+                    print("ERROR", key, e)
             elif patch_type == "loha":
                 w1a = v[0]
                 w1b = v[1]
                 if v[2] is not None:
-                    alpha = v[2] / w1b.shape[0]
-                else:
-                    alpha = 1.0
-
+                    alpha *= v[2] / w1b.shape[0]
                 w2a = v[3]
                 w2b = v[4]
-                dora_scale = v[7]
                 if v[5] is not None: #cp decomposition
                     t1 = v[5]
                     t2 = v[6]
@@ -459,69 +316,42 @@ class ModelPatcher:
                                   ldm_patched.modules.model_management.cast_to_device(w2b, weight.device, torch.float32))
 
                 try:
-                    lora_diff = (m1 * m2).reshape(weight.shape)
-                    if dora_scale is not None:
-                        weight = weight_decompose(dora_scale, weight, lora_diff, alpha, strength)
-                    else:
-                        weight += ((strength * alpha) * lora_diff).type(weight.dtype)
+                    weight += (alpha * m1 * m2).reshape(weight.shape).type(weight.dtype)
                 except Exception as e:
-                    logging.error("ERROR {} {} {}".format(patch_type, key, e))
+                    print("ERROR", key, e)
             elif patch_type == "glora":
                 if v[4] is not None:
-                    alpha = v[4] / v[0].shape[0]
-                else:
-                    alpha = 1.0
-
-                dora_scale = v[5]
+                    alpha *= v[4] / v[0].shape[0]
 
                 a1 = ldm_patched.modules.model_management.cast_to_device(v[0].flatten(start_dim=1), weight.device, torch.float32)
                 a2 = ldm_patched.modules.model_management.cast_to_device(v[1].flatten(start_dim=1), weight.device, torch.float32)
                 b1 = ldm_patched.modules.model_management.cast_to_device(v[2].flatten(start_dim=1), weight.device, torch.float32)
                 b2 = ldm_patched.modules.model_management.cast_to_device(v[3].flatten(start_dim=1), weight.device, torch.float32)
 
-                try:
-                    lora_diff = (torch.mm(b2, b1) + torch.mm(torch.mm(weight.flatten(start_dim=1), a2), a1)).reshape(weight.shape)
-                    if dora_scale is not None:
-                        weight = weight_decompose(dora_scale, weight, lora_diff, alpha, strength)
-                    else:
-                        weight += ((strength * alpha) * lora_diff).type(weight.dtype)
-                except Exception as e:
-                    logging.error("ERROR {} {} {}".format(patch_type, key, e))
+                weight += ((torch.mm(b2, b1) + torch.mm(torch.mm(weight.flatten(start_dim=1), a2), a1)) * alpha).reshape(weight.shape).type(weight.dtype)
             else:
-                logging.warning("patch type not recognized {} {}".format(patch_type, key))
+                print("patch type not recognized", patch_type, key)
 
         return weight
 
-    def unpatch_model(self, device_to=None, unpatch_weights=True):
-        if unpatch_weights:
-            if self.model_lowvram:
-                for m in self.model.modules():
-                    if hasattr(m, "prev_comfy_cast_weights"):
-                        m.comfy_cast_weights = m.prev_comfy_cast_weights
-                        del m.prev_comfy_cast_weights
-                    m.weight_function = None
-                    m.bias_function = None
+    def unpatch_model(self, device_to=None):
+        keys = list(self.backup.keys())
 
-                self.model_lowvram = False
-                self.lowvram_patch_counter = 0
+        if self.weight_inplace_update:
+            for k in keys:
+                ldm_patched.modules.utils.copy_to_param(self.model, k, self.backup[k])
+        else:
+            for k in keys:
+                ldm_patched.modules.utils.set_attr(self.model, k, self.backup[k])
 
-            keys = list(self.backup.keys())
+        self.backup = {}
 
-            if self.weight_inplace_update:
-                for k in keys:
-                    ldm_patched.modules.utils.copy_to_param(self.model, k, self.backup[k])
-            else:
-                for k in keys:
-                    ldm_patched.modules.utils.set_attr_param(self.model, k, self.backup[k])
-
-            self.backup.clear()
-
-            if device_to is not None:
-                self.model.to(device_to)
-                self.current_device = device_to
+        if device_to is not None:
+            self.model.to(device_to)
+            self.current_device = device_to
 
         keys = list(self.object_patches_backup.keys())
         for k in keys:
-            ldm_patched.modules.utils.set_attr(self.model, k, self.object_patches_backup[k])
+            setattr(self.model, k, self.object_patches_backup[k])
 
-        self.object_patches_backup.clear()
+        self.object_patches_backup = {}

ldm_patched/modules/model_sampling.py

@@ -1,6 +1,7 @@
 import torch
 from ldm_patched.ldm.modules.diffusionmodules.util import make_beta_schedule
 import math
+import numpy as np
 
 class EPS:
     def calculate_input(self, sigma, noise):
@@ -69,12 +70,17 @@ class ModelSamplingDiscrete(torch.nn.Module):
         # self.register_buffer('alphas_cumprod_prev', torch.tensor(alphas_cumprod_prev, dtype=torch.float32))
 
         sigmas = ((1 - alphas_cumprod) / alphas_cumprod) ** 0.5
+        alphas_cumprod = torch.tensor(np.cumprod(alphas, axis=0), dtype=torch.float32)
         self.set_sigmas(sigmas)
+        self.set_alphas_cumprod(alphas_cumprod.float())
 
     def set_sigmas(self, sigmas):
         self.register_buffer('sigmas', sigmas.float())
         self.register_buffer('log_sigmas', sigmas.log().float())
 
+    def set_alphas_cumprod(self, alphas_cumprod):
+        self.register_buffer("alphas_cumprod", alphas_cumprod.float())
+
     @property
     def sigma_min(self):
         return self.sigmas[0]
@@ -200,4 +206,4 @@ class StableCascadeSampling(ModelSamplingDiscrete):
             return 0.0
 
         percent = 1.0 - percent
-        return self.sigma(torch.tensor(percent))
+        return self.sigma(torch.tensor(percent))

ldm_patched/modules/ops.py

@@ -1,135 +1,89 @@
-"""
-    This file is part of ComfyUI.
-    Copyright (C) 2024 Stability AI
-
-    This program is free software: you can redistribute it and/or modify
-    it under the terms of the GNU General Public License as published by
-    the Free Software Foundation, either version 3 of the License, or
-    (at your option) any later version.
-
-    This program is distributed in the hope that it will be useful,
-    but WITHOUT ANY WARRANTY; without even the implied warranty of
-    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-    GNU General Public License for more details.
-
-    You should have received a copy of the GNU General Public License
-    along with this program.  If not, see <https://www.gnu.org/licenses/>.
-"""
-
 import torch
 import ldm_patched.modules.model_management
 
 def cast_bias_weight(s, input):
     bias = None
-    non_blocking = ldm_patched.modules.model_management.device_should_use_non_blocking(input.device)
+    non_blocking = ldm_patched.modules.model_management.device_supports_non_blocking(input.device)
     if s.bias is not None:
         bias = s.bias.to(device=input.device, dtype=input.dtype, non_blocking=non_blocking)
-        if s.bias_function is not None:
-            bias = s.bias_function(bias)
     weight = s.weight.to(device=input.device, dtype=input.dtype, non_blocking=non_blocking)
-    if s.weight_function is not None:
-        weight = s.weight_function(weight)
     return weight, bias
 
-class CastWeightBiasOp:
-    comfy_cast_weights = False
-    weight_function = None
-    bias_function = None
 
 class disable_weight_init:
-    class Linear(torch.nn.Linear, CastWeightBiasOp):
+    class Linear(torch.nn.Linear):
+        ldm_patched_cast_weights = False
         def reset_parameters(self):
             return None
 
-        def forward_comfy_cast_weights(self, input):
+        def forward_ldm_patched_cast_weights(self, input):
             weight, bias = cast_bias_weight(self, input)
             return torch.nn.functional.linear(input, weight, bias)
 
         def forward(self, *args, **kwargs):
-            if self.comfy_cast_weights:
-                return self.forward_comfy_cast_weights(*args, **kwargs)
+            if self.ldm_patched_cast_weights:
+                return self.forward_ldm_patched_cast_weights(*args, **kwargs)
             else:
                 return super().forward(*args, **kwargs)
 
-    class Conv2d(torch.nn.Conv2d, CastWeightBiasOp):
+    class Conv2d(torch.nn.Conv2d):
+        ldm_patched_cast_weights = False
         def reset_parameters(self):
             return None
 
-        def forward_comfy_cast_weights(self, input):
+        def forward_ldm_patched_cast_weights(self, input):
             weight, bias = cast_bias_weight(self, input)
             return self._conv_forward(input, weight, bias)
 
         def forward(self, *args, **kwargs):
-            if self.comfy_cast_weights:
-                return self.forward_comfy_cast_weights(*args, **kwargs)
+            if self.ldm_patched_cast_weights:
+                return self.forward_ldm_patched_cast_weights(*args, **kwargs)
             else:
                 return super().forward(*args, **kwargs)
 
-    class Conv3d(torch.nn.Conv3d, CastWeightBiasOp):
+    class Conv3d(torch.nn.Conv3d):
+        ldm_patched_cast_weights = False
         def reset_parameters(self):
             return None
 
-        def forward_comfy_cast_weights(self, input):
+        def forward_ldm_patched_cast_weights(self, input):
             weight, bias = cast_bias_weight(self, input)
             return self._conv_forward(input, weight, bias)
 
         def forward(self, *args, **kwargs):
-            if self.comfy_cast_weights:
-                return self.forward_comfy_cast_weights(*args, **kwargs)
+            if self.ldm_patched_cast_weights:
+                return self.forward_ldm_patched_cast_weights(*args, **kwargs)
             else:
                 return super().forward(*args, **kwargs)
 
-    class GroupNorm(torch.nn.GroupNorm, CastWeightBiasOp):
+    class GroupNorm(torch.nn.GroupNorm):
+        ldm_patched_cast_weights = False
         def reset_parameters(self):
             return None
 
-        def forward_comfy_cast_weights(self, input):
+        def forward_ldm_patched_cast_weights(self, input):
             weight, bias = cast_bias_weight(self, input)
             return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
 
         def forward(self, *args, **kwargs):
-            if self.comfy_cast_weights:
-                return self.forward_comfy_cast_weights(*args, **kwargs)
+            if self.ldm_patched_cast_weights:
+                return self.forward_ldm_patched_cast_weights(*args, **kwargs)
             else:
                 return super().forward(*args, **kwargs)
 
 
-    class LayerNorm(torch.nn.LayerNorm, CastWeightBiasOp):
+    class LayerNorm(torch.nn.LayerNorm):
+        ldm_patched_cast_weights = False
         def reset_parameters(self):
             return None
 
-        def forward_comfy_cast_weights(self, input):
-            if self.weight is not None:
-                weight, bias = cast_bias_weight(self, input)
-            else:
-                weight = None
-                bias = None
+        def forward_ldm_patched_cast_weights(self, input):
+            weight, bias = cast_bias_weight(self, input)
             return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)
 
         def forward(self, *args, **kwargs):
-            if self.comfy_cast_weights:
-                return self.forward_comfy_cast_weights(*args, **kwargs)
-            else:
-                return super().forward(*args, **kwargs)
-
-    class ConvTranspose2d(torch.nn.ConvTranspose2d, CastWeightBiasOp):
-        def reset_parameters(self):
-            return None
-
-        def forward_comfy_cast_weights(self, input, output_size=None):
-            num_spatial_dims = 2
-            output_padding = self._output_padding(
-                input, output_size, self.stride, self.padding, self.kernel_size,
-                num_spatial_dims, self.dilation)
-
-            weight, bias = cast_bias_weight(self, input)
-            return torch.nn.functional.conv_transpose2d(
-                input, weight, bias, self.stride, self.padding,
-                output_padding, self.groups, self.dilation)
-
-        def forward(self, *args, **kwargs):
-            if self.comfy_cast_weights:
-                return self.forward_comfy_cast_weights(*args, **kwargs)
+            if self.ldm_patched_cast_weights:
+                return self.forward_ldm_patched_cast_weights(*args, **kwargs)
             else:
                 return super().forward(*args, **kwargs)
 
@@ -145,19 +99,16 @@ class disable_weight_init:
 
 class manual_cast(disable_weight_init):
     class Linear(disable_weight_init.Linear):
-        comfy_cast_weights = True
+        ldm_patched_cast_weights = True
 
     class Conv2d(disable_weight_init.Conv2d):
-        comfy_cast_weights = True
+        ldm_patched_cast_weights = True
 
     class Conv3d(disable_weight_init.Conv3d):
-        comfy_cast_weights = True
+        ldm_patched_cast_weights = True
 
     class GroupNorm(disable_weight_init.GroupNorm):
-        comfy_cast_weights = True
+        ldm_patched_cast_weights = True
 
     class LayerNorm(disable_weight_init.LayerNorm):
-        comfy_cast_weights = True
-
-    class ConvTranspose2d(disable_weight_init.ConvTranspose2d):
-        comfy_cast_weights = True
+        ldm_patched_cast_weights = True

ldm_patched/modules/sample.py

@@ -1,9 +1,10 @@
 import torch
 import ldm_patched.modules.model_management
 import ldm_patched.modules.samplers
+import ldm_patched.modules.conds
 import ldm_patched.modules.utils
+import math
 import numpy as np
-import logging
 
 def prepare_noise(latent_image, seed, noise_inds=None):
     """
@@ -24,21 +25,94 @@ def prepare_noise(latent_image, seed, noise_inds=None):
     noises = torch.cat(noises, axis=0)
     return noises
 
-def prepare_sampling(model, noise_shape, positive, negative, noise_mask):
-    logging.warning("Warning: ldm_patched.modules.sample.prepare_sampling isn't used anymore and can be removed")
-    return model, positive, negative, noise_mask, []
+def prepare_mask(noise_mask, shape, device):
+    """ensures noise mask is of proper dimensions"""
+    noise_mask = torch.nn.functional.interpolate(noise_mask.reshape((-1, 1, noise_mask.shape[-2], noise_mask.shape[-1])), size=(shape[2], shape[3]), mode="bilinear")
+    noise_mask = torch.cat([noise_mask] * shape[1], dim=1)
+    noise_mask = ldm_patched.modules.utils.repeat_to_batch_size(noise_mask, shape[0])
+    noise_mask = noise_mask.to(device)
+    return noise_mask
+
+def get_models_from_cond(cond, model_type):
+    models = []
+    for c in cond:
+        if model_type in c:
+            models += [c[model_type]]
+    return models
+
+def convert_cond(cond):
+    out = []
+    for c in cond:
+        temp = c[1].copy()
+        model_conds = temp.get("model_conds", {})
+        if c[0] is not None:
+            model_conds["c_crossattn"] = ldm_patched.modules.conds.CONDCrossAttn(c[0]) #TODO: remove
+            temp["cross_attn"] = c[0]
+        temp["model_conds"] = model_conds
+        out.append(temp)
+    return out
+
+def get_additional_models(positive, negative, dtype):
+    """loads additional models in positive and negative conditioning"""
+    control_nets = set(get_models_from_cond(positive, "control") + get_models_from_cond(negative, "control"))
+
+    inference_memory = 0
+    control_models = []
+    for m in control_nets:
+        control_models += m.get_models()
+        inference_memory += m.inference_memory_requirements(dtype)
+
+    gligen = get_models_from_cond(positive, "gligen") + get_models_from_cond(negative, "gligen")
+    gligen = [x[1] for x in gligen]
+    models = control_models + gligen
+    return models, inference_memory
 
 def cleanup_additional_models(models):
-    logging.warning("Warning: ldm_patched.modules.sample.cleanup_additional_models isn't used anymore and can be removed")
+    """cleanup additional models that were loaded"""
+    for m in models:
+        if hasattr(m, 'cleanup'):
+            m.cleanup()
+
+def prepare_sampling(model, noise_shape, positive, negative, noise_mask):
+    device = model.load_device
+    positive = convert_cond(positive)
+    negative = convert_cond(negative)
+
+    if noise_mask is not None:
+        noise_mask = prepare_mask(noise_mask, noise_shape, device)
+
+    real_model = None
+    models, inference_memory = get_additional_models(positive, negative, model.model_dtype())
+    ldm_patched.modules.model_management.load_models_gpu([model] + models, model.memory_required([noise_shape[0] * 2] + list(noise_shape[1:])) + inference_memory)
+    real_model = model.model
+
+    return real_model, positive, negative, noise_mask, models
+
 
 def sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False, noise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None):
-    sampler = ldm_patched.modules.samplers.KSampler(model, steps=steps, device=model.load_device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options)
+    real_model, positive_copy, negative_copy, noise_mask, models = prepare_sampling(model, noise.shape, positive, negative, noise_mask)
 
-    samples = sampler.sample(noise, positive, negative, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed)
+    noise = noise.to(model.load_device)
+    latent_image = latent_image.to(model.load_device)
+
+    sampler = ldm_patched.modules.samplers.KSampler(real_model, steps=steps, device=model.load_device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options)
+
+    samples = sampler.sample(noise, positive_copy, negative_copy, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed)
     samples = samples.to(ldm_patched.modules.model_management.intermediate_device())
+
+    cleanup_additional_models(models)
+    cleanup_additional_models(set(get_models_from_cond(positive_copy, "control") + get_models_from_cond(negative_copy, "control")))
     return samples
 
 def sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=None, callback=None, disable_pbar=False, seed=None):
-    samples = ldm_patched.modules.samplers.sample(model, noise, positive, negative, cfg, model.load_device, sampler, sigmas, model_options=model.model_options, latent_image=latent_image, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
+    real_model, positive_copy, negative_copy, noise_mask, models = prepare_sampling(model, noise.shape, positive, negative, noise_mask)
+    noise = noise.to(model.load_device)
+    latent_image = latent_image.to(model.load_device)
+    sigmas = sigmas.to(model.load_device)
+
+    samples = ldm_patched.modules.samplers.sample(real_model, noise, positive_copy, negative_copy, cfg, model.load_device, sampler, sigmas, model_options=model.model_options, latent_image=latent_image, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
     samples = samples.to(ldm_patched.modules.model_management.intermediate_device())
+    cleanup_additional_models(models)
+    cleanup_additional_models(set(get_models_from_cond(positive_copy, "control") + get_models_from_cond(negative_copy, "control")))
     return samples
+

ldm_patched/modules/sampler_helpers.py

@@ -1,76 +0,0 @@
-import torch
-import ldm_patched.modules.model_management
-import ldm_patched.modules.conds
-
-def prepare_mask(noise_mask, shape, device):
-    """ensures noise mask is of proper dimensions"""
-    noise_mask = torch.nn.functional.interpolate(noise_mask.reshape((-1, 1, noise_mask.shape[-2], noise_mask.shape[-1])), size=(shape[2], shape[3]), mode="bilinear")
-    noise_mask = torch.cat([noise_mask] * shape[1], dim=1)
-    noise_mask = ldm_patched.modules.utils.repeat_to_batch_size(noise_mask, shape[0])
-    noise_mask = noise_mask.to(device)
-    return noise_mask
-
-def get_models_from_cond(cond, model_type):
-    models = []
-    for c in cond:
-        if model_type in c:
-            models += [c[model_type]]
-    return models
-
-def convert_cond(cond):
-    out = []
-    for c in cond:
-        temp = c[1].copy()
-        model_conds = temp.get("model_conds", {})
-        if c[0] is not None:
-            model_conds["c_crossattn"] = ldm_patched.modules.conds.CONDCrossAttn(c[0]) #TODO: remove
-            temp["cross_attn"] = c[0]
-        temp["model_conds"] = model_conds
-        out.append(temp)
-    return out
-
-def get_additional_models(conds, dtype):
-    """loads additional models in conditioning"""
-    cnets = []
-    gligen = []
-
-    for k in conds:
-        cnets += get_models_from_cond(conds[k], "control")
-        gligen += get_models_from_cond(conds[k], "gligen")
-
-    control_nets = set(cnets)
-
-    inference_memory = 0
-    control_models = []
-    for m in control_nets:
-        control_models += m.get_models()
-        inference_memory += m.inference_memory_requirements(dtype)
-
-    gligen = [x[1] for x in gligen]
-    models = control_models + gligen
-    return models, inference_memory
-
-def cleanup_additional_models(models):
-    """cleanup additional models that were loaded"""
-    for m in models:
-        if hasattr(m, 'cleanup'):
-            m.cleanup()
-
-
-def prepare_sampling(model, noise_shape, conds):
-    device = model.load_device
-    real_model = None
-    models, inference_memory = get_additional_models(conds, model.model_dtype())
-    ldm_patched.modules.model_management.load_models_gpu([model] + models, model.memory_required([noise_shape[0] * 2] + list(noise_shape[1:])) + inference_memory)
-    real_model = model.model
-
-    return real_model, conds, models
-
-def cleanup_models(conds, models):
-    cleanup_additional_models(models)
-
-    control_cleanup = []
-    for k in conds:
-        control_cleanup += get_models_from_cond(conds[k], "control")
-
-    cleanup_additional_models(set(control_cleanup))

ldm_patched/modules/samplers.py

@@ -4,8 +4,6 @@ import torch
 import collections
 from ldm_patched.modules import model_management
 import math
-import logging
-import ldm_patched.modules.sampler_helpers
 
 def get_area_and_mult(conds, x_in, timestep_in):
     area = (x_in.shape[2], x_in.shape[3], 0, 0)
@@ -34,7 +32,7 @@ def get_area_and_mult(conds, x_in, timestep_in):
         mask = conds['mask']
         assert(mask.shape[1] == x_in.shape[2])
         assert(mask.shape[2] == x_in.shape[3])
-        mask = mask[:input_x.shape[0],area[2]:area[0] + area[2],area[3]:area[1] + area[3]] * mask_strength
+        mask = mask[:,area[2]:area[0] + area[2],area[3]:area[1] + area[3]] * mask_strength
         mask = mask.unsqueeze(1).repeat(input_x.shape[0] // mask.shape[0], input_x.shape[1], 1, 1)
     else:
         mask = torch.ones_like(input_x)
@@ -128,23 +126,30 @@ def cond_cat(c_list):
 
     return out
 
-def calc_cond_batch(model, conds, x_in, timestep, model_options):
-    out_conds = []
-    out_counts = []
+def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options):
+    out_cond = torch.zeros_like(x_in)
+    out_count = torch.ones_like(x_in) * 1e-37
+
+    out_uncond = torch.zeros_like(x_in)
+    out_uncond_count = torch.ones_like(x_in) * 1e-37
+
+    COND = 0
+    UNCOND = 1
+
     to_run = []
+    for x in cond:
+        p = get_area_and_mult(x, x_in, timestep)
+        if p is None:
+            continue
 
-    for i in range(len(conds)):
-        out_conds.append(torch.zeros_like(x_in))
-        out_counts.append(torch.ones_like(x_in) * 1e-37)
+        to_run += [(p, COND)]
+    if uncond is not None:
+        for x in uncond:
+            p = get_area_and_mult(x, x_in, timestep)
+            if p is None:
+                continue
 
-        cond = conds[i]
-        if cond is not None:
-            for x in cond:
-                p = get_area_and_mult(x, x_in, timestep)
-                if p is None:
-                    continue
-
-                to_run += [(p, i)]
+            to_run += [(p, UNCOND)]
 
     while len(to_run) > 0:
         first = to_run[0]
@@ -203,7 +208,6 @@ def calc_cond_batch(model, conds, x_in, timestep, model_options):
                         cur_patches[p] = cur_patches[p] + patches[p]
                     else:
                         cur_patches[p] = patches[p]
-                transformer_options["patches"] = cur_patches
             else:
                 transformer_options["patches"] = patches
 
@@ -216,66 +220,71 @@ def calc_cond_batch(model, conds, x_in, timestep, model_options):
             output = model_options['model_function_wrapper'](model.apply_model, {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": cond_or_uncond}).chunk(batch_chunks)
         else:
             output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks)
+        del input_x
 
         for o in range(batch_chunks):
-            cond_index = cond_or_uncond[o]
-            out_conds[cond_index][:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += output[o] * mult[o]
-            out_counts[cond_index][:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += mult[o]
+            if cond_or_uncond[o] == COND:
+                out_cond[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += output[o] * mult[o]
+                out_count[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += mult[o]
+            else:
+                out_uncond[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += output[o] * mult[o]
+                out_uncond_count[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += mult[o]
+        del mult
 
-    for i in range(len(out_conds)):
-        out_conds[i] /= out_counts[i]
-
-    return out_conds
-
-def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options): #TODO: remove
-    logging.warning("WARNING: The comfy.samplers.calc_cond_uncond_batch function is deprecated please use the calc_cond_batch one instead.")
-    return tuple(calc_cond_batch(model, [cond, uncond], x_in, timestep, model_options))
-
-def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options={}, cond=None, uncond=None):
-    if "sampler_cfg_function" in model_options:
-        args = {"cond": x - cond_pred, "uncond": x - uncond_pred, "cond_scale": cond_scale, "timestep": timestep, "input": x, "sigma": timestep,
-                "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options}
-        cfg_result = x - model_options["sampler_cfg_function"](args)
-    else:
-        cfg_result = uncond_pred + (cond_pred - uncond_pred) * cond_scale
-
-    for fn in model_options.get("sampler_post_cfg_function", []):
-        args = {"denoised": cfg_result, "cond": cond, "uncond": uncond, "model": model, "uncond_denoised": uncond_pred, "cond_denoised": cond_pred,
-                "sigma": timestep, "model_options": model_options, "input": x}
-        cfg_result = fn(args)
-
-    return cfg_result
+    out_cond /= out_count
+    del out_count
+    out_uncond /= out_uncond_count
+    del out_uncond_count
+    return out_cond, out_uncond
 
 #The main sampling function shared by all the samplers
 #Returns denoised
 def sampling_function(model, x, timestep, uncond, cond, cond_scale, model_options={}, seed=None):
-    if math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False:
-        uncond_ = None
-    else:
-        uncond_ = uncond
+        if math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False:
+            uncond_ = None
+        else:
+            uncond_ = uncond
 
-    conds = [cond, uncond_]
-    out = calc_cond_batch(model, conds, x, timestep, model_options)
-    return cfg_function(model, out[0], out[1], cond_scale, x, timestep, model_options=model_options, cond=cond, uncond=uncond_)
+        cond_pred, uncond_pred = calc_cond_uncond_batch(model, cond, uncond_, x, timestep, model_options)
+        if "sampler_cfg_function" in model_options:
+            args = {"cond": x - cond_pred, "uncond": x - uncond_pred, "cond_scale": cond_scale, "timestep": timestep, "input": x, "sigma": timestep,
+                    "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options}
+            cfg_result = x - model_options["sampler_cfg_function"](args)
+        else:
+            cfg_result = uncond_pred + (cond_pred - uncond_pred) * cond_scale
 
+        for fn in model_options.get("sampler_post_cfg_function", []):
+            args = {"denoised": cfg_result, "cond": cond, "uncond": uncond, "model": model, "uncond_denoised": uncond_pred, "cond_denoised": cond_pred,
+                    "sigma": timestep, "model_options": model_options, "input": x}
+            cfg_result = fn(args)
 
-class KSamplerX0Inpaint:
-    def __init__(self, model, sigmas):
+        return cfg_result
+
+class CFGNoisePredictor(torch.nn.Module):
+    def __init__(self, model):
+        super().__init__()
         self.inner_model = model
-        self.sigmas = sigmas
-    def __call__(self, x, sigma, denoise_mask, model_options={}, seed=None):
+    def apply_model(self, x, timestep, cond, uncond, cond_scale, model_options={}, seed=None):
+        out = sampling_function(self.inner_model, x, timestep, uncond, cond, cond_scale, model_options=model_options, seed=seed)
+        return out
+    def forward(self, *args, **kwargs):
+        return self.apply_model(*args, **kwargs)
+
+class KSamplerX0Inpaint(torch.nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.inner_model = model
+    def forward(self, x, sigma, uncond, cond, cond_scale, denoise_mask, model_options={}, seed=None):
         if denoise_mask is not None:
-            if "denoise_mask_function" in model_options:
-                denoise_mask = model_options["denoise_mask_function"](sigma, denoise_mask, extra_options={"model": self.inner_model, "sigmas": self.sigmas})
             latent_mask = 1. - denoise_mask
-            x = x * denoise_mask + self.inner_model.inner_model.model_sampling.noise_scaling(sigma.reshape([sigma.shape[0]] + [1] * (len(self.noise.shape) - 1)), self.noise, self.latent_image) * latent_mask
-        out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
+            x = x * denoise_mask + (self.latent_image + self.noise * sigma.reshape([sigma.shape[0]] + [1] * (len(self.noise.shape) - 1))) * latent_mask
+        out = self.inner_model(x, sigma, cond=cond, uncond=uncond, cond_scale=cond_scale, model_options=model_options, seed=seed)
         if denoise_mask is not None:
             out = out * denoise_mask + self.latent_image * latent_mask
         return out
 
-def simple_scheduler(model_sampling, steps):
-    s = model_sampling
+def simple_scheduler(model, steps):
+    s = model.model_sampling
     sigs = []
     ss = len(s.sigmas) / steps
     for x in range(steps):
@@ -283,10 +292,10 @@ def simple_scheduler(model_sampling, steps):
     sigs += [0.0]
     return torch.FloatTensor(sigs)
 
-def ddim_scheduler(model_sampling, steps):
-    s = model_sampling
+def ddim_scheduler(model, steps):
+    s = model.model_sampling
     sigs = []
-    ss = max(len(s.sigmas) // steps, 1)
+    ss = len(s.sigmas) // steps
     x = 1
     while x < len(s.sigmas):
         sigs += [float(s.sigmas[x])]
@@ -295,8 +304,8 @@ def ddim_scheduler(model_sampling, steps):
     sigs += [0.0]
     return torch.FloatTensor(sigs)
 
-def normal_scheduler(model_sampling, steps, sgm=False, floor=False):
-    s = model_sampling
+def normal_scheduler(model, steps, sgm=False, floor=False):
+    s = model.model_sampling
     start = s.timestep(s.sigma_max)
     end = s.timestep(s.sigma_min)
 
@@ -504,9 +513,17 @@ class Sampler:
         sigma = float(sigmas[0])
         return math.isclose(max_sigma, sigma, rel_tol=1e-05) or sigma > max_sigma
 
+class UNIPC(Sampler):
+    def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
+        return uni_pc.sample_unipc(model_wrap, noise, latent_image, sigmas, max_denoise=self.max_denoise(model_wrap, sigmas), extra_args=extra_args, noise_mask=denoise_mask, callback=callback, disable=disable_pbar)
+
+class UNIPCBH2(Sampler):
+    def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
+        return uni_pc.sample_unipc(model_wrap, noise, latent_image, sigmas, max_denoise=self.max_denoise(model_wrap, sigmas), extra_args=extra_args, noise_mask=denoise_mask, callback=callback, variant='bh2', disable=disable_pbar)
+
 KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
                   "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
-                  "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm"]
+                  "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm", "tcd", "edm_playground_v2.5", "restart"]
 
 class KSAMPLER(Sampler):
     def __init__(self, sampler_function, extra_options={}, inpaint_options={}):
@@ -516,7 +533,7 @@ class KSAMPLER(Sampler):
 
     def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
         extra_args["denoise_mask"] = denoise_mask
-        model_k = KSamplerX0Inpaint(model_wrap, sigmas)
+        model_k = KSamplerX0Inpaint(model_wrap)
         model_k.latent_image = latent_image
         if self.inpaint_options.get("random", False): #TODO: Should this be the default?
             generator = torch.manual_seed(extra_args.get("seed", 41) + 1)
@@ -524,24 +541,26 @@ class KSAMPLER(Sampler):
         else:
             model_k.noise = noise
 
-        noise = model_wrap.inner_model.model_sampling.noise_scaling(sigmas[0], noise, latent_image, self.max_denoise(model_wrap, sigmas))
+        if self.max_denoise(model_wrap, sigmas):
+            noise = noise * torch.sqrt(1.0 + sigmas[0] ** 2.0)
+        else:
+            noise = noise * sigmas[0]
 
         k_callback = None
         total_steps = len(sigmas) - 1
         if callback is not None:
             k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps)
 
+        if latent_image is not None:
+            noise += latent_image
+
         samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
-        samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples)
         return samples
 
 
 def ksampler(sampler_name, extra_options={}, inpaint_options={}):
     if sampler_name == "dpm_fast":
         def dpm_fast_function(model, noise, sigmas, extra_args, callback, disable):
-            if len(sigmas) <= 1:
-                return noise
-
             sigma_min = sigmas[-1]
             if sigma_min == 0:
                 sigma_min = sigmas[-2]
@@ -549,145 +568,81 @@ def ksampler(sampler_name, extra_options={}, inpaint_options={}):
             return k_diffusion_sampling.sample_dpm_fast(model, noise, sigma_min, sigmas[0], total_steps, extra_args=extra_args, callback=callback, disable=disable)
         sampler_function = dpm_fast_function
     elif sampler_name == "dpm_adaptive":
-        def dpm_adaptive_function(model, noise, sigmas, extra_args, callback, disable, **extra_options):
-            if len(sigmas) <= 1:
-                return noise
-
+        def dpm_adaptive_function(model, noise, sigmas, extra_args, callback, disable):
             sigma_min = sigmas[-1]
             if sigma_min == 0:
                 sigma_min = sigmas[-2]
-            return k_diffusion_sampling.sample_dpm_adaptive(model, noise, sigma_min, sigmas[0], extra_args=extra_args, callback=callback, disable=disable, **extra_options)
+            return k_diffusion_sampling.sample_dpm_adaptive(model, noise, sigma_min, sigmas[0], extra_args=extra_args, callback=callback, disable=disable)
         sampler_function = dpm_adaptive_function
     else:
         sampler_function = getattr(k_diffusion_sampling, "sample_{}".format(sampler_name))
 
     return KSAMPLER(sampler_function, extra_options, inpaint_options)
 
-
-def process_conds(model, noise, conds, device, latent_image=None, denoise_mask=None, seed=None):
-    for k in conds:
-        conds[k] = conds[k][:]
-        resolve_areas_and_cond_masks(conds[k], noise.shape[2], noise.shape[3], device)
-
-    for k in conds:
-        calculate_start_end_timesteps(model, conds[k])
-
-    if hasattr(model, 'extra_conds'):
-        for k in conds:
-            conds[k] = encode_model_conds(model.extra_conds, conds[k], noise, device, k, latent_image=latent_image, denoise_mask=denoise_mask, seed=seed)
-
-    #make sure each cond area has an opposite one with the same area
-    for k in conds:
-        for c in conds[k]:
-            for kk in conds:
-                if k != kk:
-                    create_cond_with_same_area_if_none(conds[kk], c)
-
-    for k in conds:
-        pre_run_control(model, conds[k])
-
-    if "positive" in conds:
-        positive = conds["positive"]
-        for k in conds:
-            if k != "positive":
-                apply_empty_x_to_equal_area(list(filter(lambda c: c.get('control_apply_to_uncond', False) == True, positive)), conds[k], 'control', lambda cond_cnets, x: cond_cnets[x])
-                apply_empty_x_to_equal_area(positive, conds[k], 'gligen', lambda cond_cnets, x: cond_cnets[x])
-
-    return conds
-
-class CFGGuider:
-    def __init__(self, model_patcher):
-        self.model_patcher = model_patcher
-        self.model_options = model_patcher.model_options
-        self.original_conds = {}
-        self.cfg = 1.0
-
-    def set_conds(self, positive, negative):
-        self.inner_set_conds({"positive": positive, "negative": negative})
-
-    def set_cfg(self, cfg):
-        self.cfg = cfg
-
-    def inner_set_conds(self, conds):
-        for k in conds:
-            self.original_conds[k] = ldm_patched.modules.sampler_helpers.convert_cond(conds[k])
-
-    def __call__(self, *args, **kwargs):
-        return self.predict_noise(*args, **kwargs)
-
-    def predict_noise(self, x, timestep, model_options={}, seed=None):
-        return sampling_function(self.inner_model, x, timestep, self.conds.get("negative", None), self.conds.get("positive", None), self.cfg, model_options=model_options, seed=seed)
-
-    def inner_sample(self, noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed):
-        if latent_image is not None and torch.count_nonzero(latent_image) > 0: #Don't shift the empty latent image.
-            latent_image = self.inner_model.process_latent_in(latent_image)
-
-        self.conds = process_conds(self.inner_model, noise, self.conds, device, latent_image, denoise_mask, seed)
-
-        extra_args = {"model_options": self.model_options, "seed":seed}
-
-        samples = sampler.sample(self, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
-        return self.inner_model.process_latent_out(samples.to(torch.float32))
-
-    def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
-        if sigmas.shape[-1] == 0:
-            return latent_image
-
-        self.conds = {}
-        for k in self.original_conds:
-            self.conds[k] = list(map(lambda a: a.copy(), self.original_conds[k]))
-
-        self.inner_model, self.conds, self.loaded_models = ldm_patched.modules.sampler_helpers.prepare_sampling(self.model_patcher, noise.shape, self.conds)
-        device = self.model_patcher.load_device
-
-        if denoise_mask is not None:
-            denoise_mask = ldm_patched.modules.sampler_helpers.prepare_mask(denoise_mask, noise.shape, device)
-
-        noise = noise.to(device)
-        latent_image = latent_image.to(device)
-        sigmas = sigmas.to(device)
-
-        output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
-
-        ldm_patched.modules.sampler_helpers.cleanup_models(self.conds, self.loaded_models)
-        del self.inner_model
-        del self.conds
-        del self.loaded_models
-        return output
-
+def wrap_model(model):
+    model_denoise = CFGNoisePredictor(model)
+    return model_denoise
 
 def sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
-    cfg_guider = CFGGuider(model)
-    cfg_guider.set_conds(positive, negative)
-    cfg_guider.set_cfg(cfg)
-    return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
+    positive = positive[:]
+    negative = negative[:]
 
+    resolve_areas_and_cond_masks(positive, noise.shape[2], noise.shape[3], device)
+    resolve_areas_and_cond_masks(negative, noise.shape[2], noise.shape[3], device)
+
+    model_wrap = wrap_model(model)
+
+    calculate_start_end_timesteps(model, negative)
+    calculate_start_end_timesteps(model, positive)
+
+    if latent_image is not None:
+        latent_image = model.process_latent_in(latent_image)
+
+    if hasattr(model, 'extra_conds'):
+        positive = encode_model_conds(model.extra_conds, positive, noise, device, "positive", latent_image=latent_image, denoise_mask=denoise_mask, seed=seed)
+        negative = encode_model_conds(model.extra_conds, negative, noise, device, "negative", latent_image=latent_image, denoise_mask=denoise_mask, seed=seed)
+
+    #make sure each cond area has an opposite one with the same area
+    for c in positive:
+        create_cond_with_same_area_if_none(negative, c)
+    for c in negative:
+        create_cond_with_same_area_if_none(positive, c)
+
+    pre_run_control(model, negative + positive)
+
+    apply_empty_x_to_equal_area(list(filter(lambda c: c.get('control_apply_to_uncond', False) == True, positive)), negative, 'control', lambda cond_cnets, x: cond_cnets[x])
+    apply_empty_x_to_equal_area(positive, negative, 'gligen', lambda cond_cnets, x: cond_cnets[x])
+
+    extra_args = {"cond":positive, "uncond":negative, "cond_scale": cfg, "model_options": model_options, "seed":seed}
+
+    samples = sampler.sample(model_wrap, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
+    return model.process_latent_out(samples.to(torch.float32))
 
 SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"]
 SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"]
 
-def calculate_sigmas(model_sampling, scheduler_name, steps):
+def calculate_sigmas_scheduler(model, scheduler_name, steps):
     if scheduler_name == "karras":
-        sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=float(model_sampling.sigma_min), sigma_max=float(model_sampling.sigma_max))
+        sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=float(model.model_sampling.sigma_min), sigma_max=float(model.model_sampling.sigma_max))
     elif scheduler_name == "exponential":
-        sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=float(model_sampling.sigma_min), sigma_max=float(model_sampling.sigma_max))
+        sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=float(model.model_sampling.sigma_min), sigma_max=float(model.model_sampling.sigma_max))
     elif scheduler_name == "normal":
-        sigmas = normal_scheduler(model_sampling, steps)
+        sigmas = normal_scheduler(model, steps)
     elif scheduler_name == "simple":
-        sigmas = simple_scheduler(model_sampling, steps)
+        sigmas = simple_scheduler(model, steps)
     elif scheduler_name == "ddim_uniform":
-        sigmas = ddim_scheduler(model_sampling, steps)
+        sigmas = ddim_scheduler(model, steps)
     elif scheduler_name == "sgm_uniform":
-        sigmas = normal_scheduler(model_sampling, steps, sgm=True)
+        sigmas = normal_scheduler(model, steps, sgm=True)
     else:
-        logging.error("error invalid scheduler {}".format(scheduler_name))
+        print("error invalid scheduler", scheduler_name)
     return sigmas
 
 def sampler_object(name):
     if name == "uni_pc":
-        sampler = KSAMPLER(uni_pc.sample_unipc)
+        sampler = UNIPC()
     elif name == "uni_pc_bh2":
-        sampler = KSAMPLER(uni_pc.sample_unipc_bh2)
+        sampler = UNIPCBH2()
     elif name == "ddim":
         sampler = ksampler("euler", inpaint_options={"random": True})
     else:
@@ -697,7 +652,6 @@ def sampler_object(name):
 class KSampler:
     SCHEDULERS = SCHEDULER_NAMES
     SAMPLERS = SAMPLER_NAMES
-    DISCARD_PENULTIMATE_SIGMA_SAMPLERS = set(('dpm_2', 'dpm_2_ancestral', 'uni_pc', 'uni_pc_bh2'))
 
     def __init__(self, model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
         self.model = model
@@ -716,11 +670,11 @@ class KSampler:
         sigmas = None
 
         discard_penultimate_sigma = False
-        if self.sampler in self.DISCARD_PENULTIMATE_SIGMA_SAMPLERS:
+        if self.sampler in ['dpm_2', 'dpm_2_ancestral', 'uni_pc', 'uni_pc_bh2']:
             steps += 1
             discard_penultimate_sigma = True
 
-        sigmas = calculate_sigmas(self.model, self.scheduler, steps)
+        sigmas = calculate_sigmas_scheduler(self.model, self.scheduler, steps)
 
         if discard_penultimate_sigma:
             sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
@@ -731,12 +685,9 @@ class KSampler:
         if denoise is None or denoise > 0.9999:
             self.sigmas = self.calculate_sigmas(steps).to(self.device)
         else:
-            if denoise <= 0.0:
-                self.sigmas = torch.FloatTensor([])
-            else:
-                new_steps = int(steps/denoise)
-                sigmas = self.calculate_sigmas(new_steps).to(self.device)
-                self.sigmas = sigmas[-(steps + 1):]
+            new_steps = int(steps/denoise)
+            sigmas = self.calculate_sigmas(new_steps).to(self.device)
+            self.sigmas = sigmas[-(steps + 1):]
 
     def sample(self, noise, positive, negative, cfg, latent_image=None, start_step=None, last_step=None, force_full_denoise=False, denoise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None):
         if sigmas is None:

ldm_patched/modules/sd.py

@@ -1,12 +1,7 @@
 import torch
-from enum import Enum
-import logging
 
 from ldm_patched.modules import model_management
 from ldm_patched.ldm.models.autoencoder import AutoencoderKL, AutoencodingEngine
-from ldm_patched.ldm.cascade.stage_a import StageA
-from ldm_patched.ldm.cascade.stage_c_coder import StageC_coder
-
 import yaml
 
 import ldm_patched.modules.utils
@@ -14,6 +9,7 @@ import ldm_patched.modules.utils
 from . import clip_vision
 from . import gligen
 from . import diffusers_convert
+from . import model_base
 from . import model_detection
 
 from . import sd1_clip
@@ -37,7 +33,7 @@ def load_model_weights(model, sd):
             w = sd.pop(x)
             del w
     if len(m) > 0:
-        logging.warning("missing {}".format(m))
+        print("extra", m)
     return model
 
 def load_clip_weights(model, sd):
@@ -52,7 +48,7 @@ def load_clip_weights(model, sd):
         if ids.dtype == torch.float32:
             sd['cond_stage_model.transformer.text_model.embeddings.position_ids'] = ids.round()
 
-    sd = ldm_patched.modules.utils.clip_text_transformers_convert(sd, "cond_stage_model.model.", "cond_stage_model.transformer.")
+    sd = ldm_patched.modules.utils.transformers_convert(sd, "cond_stage_model.model.", "cond_stage_model.transformer.text_model.", 24)
     return load_model_weights(model, sd)
 
 
@@ -81,7 +77,7 @@ def load_lora_for_models(model, clip, lora, strength_model, strength_clip):
     k1 = set(k1)
     for x in loaded:
         if (x not in k) and (x not in k1):
-            logging.warning("NOT LOADED {}".format(x))
+            print("NOT LOADED", x)
 
     return (new_modelpatcher, new_clip)
 
@@ -123,13 +119,10 @@ class CLIP:
         return self.tokenizer.tokenize_with_weights(text, return_word_ids)
 
     def encode_from_tokens(self, tokens, return_pooled=False):
-        self.cond_stage_model.reset_clip_options()
-
         if self.layer_idx is not None:
-            self.cond_stage_model.set_clip_options({"layer": self.layer_idx})
-
-        if return_pooled == "unprojected":
-            self.cond_stage_model.set_clip_options({"projected_pooled": False})
+            self.cond_stage_model.clip_layer(self.layer_idx)
+        else:
+            self.cond_stage_model.reset_clip_layer()
 
         self.load_model()
         cond, pooled = self.cond_stage_model.encode_token_weights(tokens)
@@ -141,11 +134,8 @@ class CLIP:
         tokens = self.tokenize(text)
         return self.encode_from_tokens(tokens)
 
-    def load_sd(self, sd, full_model=False):
-        if full_model:
-            return self.cond_stage_model.load_state_dict(sd, strict=False)
-        else:
-            return self.cond_stage_model.load_sd(sd)
+    def load_sd(self, sd):
+        return self.cond_stage_model.load_sd(sd)
 
     def get_sd(self):
         return self.cond_stage_model.state_dict()
@@ -165,10 +155,7 @@ class VAE:
         self.memory_used_encode = lambda shape, dtype: (1767 * shape[2] * shape[3]) * model_management.dtype_size(dtype) #These are for AutoencoderKL and need tweaking (should be lower)
         self.memory_used_decode = lambda shape, dtype: (2178 * shape[2] * shape[3] * 64) * model_management.dtype_size(dtype)
         self.downscale_ratio = 8
-        self.upscale_ratio = 8
         self.latent_channels = 4
-        self.process_input = lambda image: image * 2.0 - 1.0
-        self.process_output = lambda image: torch.clamp((image + 1.0) / 2.0, min=0.0, max=1.0)
 
         if config is None:
             if "decoder.mid.block_1.mix_factor" in sd:
@@ -181,64 +168,25 @@ class VAE:
                                                             decoder_config={'target': "ldm_patched.ldm.modules.temporal_ae.VideoDecoder", 'params': decoder_config})
             elif "taesd_decoder.1.weight" in sd:
                 self.first_stage_model = ldm_patched.taesd.taesd.TAESD()
-            elif "vquantizer.codebook.weight" in sd: #VQGan: stage a of stable cascade
-                self.first_stage_model = StageA()
-                self.downscale_ratio = 4
-                self.upscale_ratio = 4
-                #TODO
-                #self.memory_used_encode
-                #self.memory_used_decode
-                self.process_input = lambda image: image
-                self.process_output = lambda image: image
-            elif "backbone.1.0.block.0.1.num_batches_tracked" in sd: #effnet: encoder for stage c latent of stable cascade
-                self.first_stage_model = StageC_coder()
-                self.downscale_ratio = 32
-                self.latent_channels = 16
-                new_sd = {}
-                for k in sd:
-                    new_sd["encoder.{}".format(k)] = sd[k]
-                sd = new_sd
-            elif "blocks.11.num_batches_tracked" in sd: #previewer: decoder for stage c latent of stable cascade
-                self.first_stage_model = StageC_coder()
-                self.latent_channels = 16
-                new_sd = {}
-                for k in sd:
-                    new_sd["previewer.{}".format(k)] = sd[k]
-                sd = new_sd
-            elif "encoder.backbone.1.0.block.0.1.num_batches_tracked" in sd: #combined effnet and previewer for stable cascade
-                self.first_stage_model = StageC_coder()
-                self.downscale_ratio = 32
-                self.latent_channels = 16
-            elif "decoder.conv_in.weight" in sd:
+            else:
                 #default SD1.x/SD2.x VAE parameters
                 ddconfig = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}
 
-                if 'encoder.down.2.downsample.conv.weight' not in sd and 'decoder.up.3.upsample.conv.weight' not in sd: #Stable diffusion x4 upscaler VAE
+                if 'encoder.down.2.downsample.conv.weight' not in sd: #Stable diffusion x4 upscaler VAE
                     ddconfig['ch_mult'] = [1, 2, 4]
                     self.downscale_ratio = 4
-                    self.upscale_ratio = 4
 
-                self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.weight"].shape[1]
-                if 'quant_conv.weight' in sd:
-                    self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=4)
-                else:
-                    self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "ldm_patched.ldm.models.autoencoder.DiagonalGaussianRegularizer"},
-                                                                encoder_config={'target': "ldm_patched.ldm.modules.diffusionmodules.model.Encoder", 'params': ddconfig},
-                                                                decoder_config={'target': "ldm_patched.ldm.modules.diffusionmodules.model.Decoder", 'params': ddconfig})
-            else:
-                logging.warning("WARNING: No VAE weights detected, VAE not initalized.")
-                self.first_stage_model = None
-                return
+                self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=4)
         else:
             self.first_stage_model = AutoencoderKL(**(config['params']))
         self.first_stage_model = self.first_stage_model.eval()
 
         m, u = self.first_stage_model.load_state_dict(sd, strict=False)
         if len(m) > 0:
-            logging.warning("Missing VAE keys {}".format(m))
+            print("Missing VAE keys", m)
 
         if len(u) > 0:
-            logging.debug("Leftover VAE keys {}".format(u))
+            print("Leftover VAE keys", u)
 
         if device is None:
             device = model_management.vae_device()
@@ -252,27 +200,18 @@ class VAE:
 
         self.patcher = ldm_patched.modules.model_patcher.ModelPatcher(self.first_stage_model, load_device=self.device, offload_device=offload_device)
 
-    def vae_encode_crop_pixels(self, pixels):
-        x = (pixels.shape[1] // self.downscale_ratio) * self.downscale_ratio
-        y = (pixels.shape[2] // self.downscale_ratio) * self.downscale_ratio
-        if pixels.shape[1] != x or pixels.shape[2] != y:
-            x_offset = (pixels.shape[1] % self.downscale_ratio) // 2
-            y_offset = (pixels.shape[2] % self.downscale_ratio) // 2
-            pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :]
-        return pixels
-
     def decode_tiled_(self, samples, tile_x=64, tile_y=64, overlap = 16):
         steps = samples.shape[0] * ldm_patched.modules.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x, tile_y, overlap)
         steps += samples.shape[0] * ldm_patched.modules.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x // 2, tile_y * 2, overlap)
         steps += samples.shape[0] * ldm_patched.modules.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x * 2, tile_y // 2, overlap)
         pbar = ldm_patched.modules.utils.ProgressBar(steps)
 
-        decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).float()
-        output = self.process_output(
-            (ldm_patched.modules.utils.tiled_scale(samples, decode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar) +
-            ldm_patched.modules.utils.tiled_scale(samples, decode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar) +
-             ldm_patched.modules.utils.tiled_scale(samples, decode_fn, tile_x, tile_y, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar))
-            / 3.0)
+        decode_fn = lambda a: (self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)) + 1.0).float()
+        output = torch.clamp((
+            (ldm_patched.modules.utils.tiled_scale(samples, decode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = self.downscale_ratio, output_device=self.output_device, pbar = pbar) +
+            ldm_patched.modules.utils.tiled_scale(samples, decode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = self.downscale_ratio, output_device=self.output_device, pbar = pbar) +
+             ldm_patched.modules.utils.tiled_scale(samples, decode_fn, tile_x, tile_y, overlap, upscale_amount = self.downscale_ratio, output_device=self.output_device, pbar = pbar))
+            / 3.0) / 2.0, min=0.0, max=1.0)
         return output
 
     def encode_tiled_(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
@@ -281,7 +220,7 @@ class VAE:
         steps += pixel_samples.shape[0] * ldm_patched.modules.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x * 2, tile_y // 2, overlap)
         pbar = ldm_patched.modules.utils.ProgressBar(steps)
 
-        encode_fn = lambda a: self.first_stage_model.encode((self.process_input(a)).to(self.vae_dtype).to(self.device)).float()
+        encode_fn = lambda a: self.first_stage_model.encode((2. * a - 1.).to(self.vae_dtype).to(self.device)).float()
         samples = ldm_patched.modules.utils.tiled_scale(pixel_samples, encode_fn, tile_x, tile_y, overlap, upscale_amount = (1/self.downscale_ratio), out_channels=self.latent_channels, output_device=self.output_device, pbar=pbar)
         samples += ldm_patched.modules.utils.tiled_scale(pixel_samples, encode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = (1/self.downscale_ratio), out_channels=self.latent_channels, output_device=self.output_device, pbar=pbar)
         samples += ldm_patched.modules.utils.tiled_scale(pixel_samples, encode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = (1/self.downscale_ratio), out_channels=self.latent_channels, output_device=self.output_device, pbar=pbar)
@@ -296,12 +235,12 @@ class VAE:
             batch_number = int(free_memory / memory_used)
             batch_number = max(1, batch_number)
 
-            pixel_samples = torch.empty((samples_in.shape[0], 3, round(samples_in.shape[2] * self.upscale_ratio), round(samples_in.shape[3] * self.upscale_ratio)), device=self.output_device)
+            pixel_samples = torch.empty((samples_in.shape[0], 3, round(samples_in.shape[2] * self.downscale_ratio), round(samples_in.shape[3] * self.downscale_ratio)), device=self.output_device)
             for x in range(0, samples_in.shape[0], batch_number):
                 samples = samples_in[x:x+batch_number].to(self.vae_dtype).to(self.device)
-                pixel_samples[x:x+batch_number] = self.process_output(self.first_stage_model.decode(samples).to(self.output_device).float())
+                pixel_samples[x:x+batch_number] = torch.clamp((self.first_stage_model.decode(samples).to(self.output_device).float() + 1.0) / 2.0, min=0.0, max=1.0)
         except model_management.OOM_EXCEPTION as e:
-            logging.warning("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
+            print("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
             pixel_samples = self.decode_tiled_(samples_in)
 
         pixel_samples = pixel_samples.to(self.output_device).movedim(1,-1)
@@ -313,7 +252,6 @@ class VAE:
         return output.movedim(1,-1)
 
     def encode(self, pixel_samples):
-        pixel_samples = self.vae_encode_crop_pixels(pixel_samples)
         pixel_samples = pixel_samples.movedim(-1,1)
         try:
             memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype)
@@ -323,17 +261,16 @@ class VAE:
             batch_number = max(1, batch_number)
             samples = torch.empty((pixel_samples.shape[0], self.latent_channels, round(pixel_samples.shape[2] // self.downscale_ratio), round(pixel_samples.shape[3] // self.downscale_ratio)), device=self.output_device)
             for x in range(0, pixel_samples.shape[0], batch_number):
-                pixels_in = self.process_input(pixel_samples[x:x+batch_number]).to(self.vae_dtype).to(self.device)
+                pixels_in = (2. * pixel_samples[x:x+batch_number] - 1.).to(self.vae_dtype).to(self.device)
                 samples[x:x+batch_number] = self.first_stage_model.encode(pixels_in).to(self.output_device).float()
 
         except model_management.OOM_EXCEPTION as e:
-            logging.warning("Warning: Ran out of memory when regular VAE encoding, retrying with tiled VAE encoding.")
+            print("Warning: Ran out of memory when regular VAE encoding, retrying with tiled VAE encoding.")
             samples = self.encode_tiled_(pixel_samples)
 
         return samples
 
     def encode_tiled(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
-        pixel_samples = self.vae_encode_crop_pixels(pixel_samples)
         model_management.load_model_gpu(self.patcher)
         pixel_samples = pixel_samples.movedim(-1,1)
         samples = self.encode_tiled_(pixel_samples, tile_x=tile_x, tile_y=tile_y, overlap=overlap)
@@ -354,17 +291,14 @@ def load_style_model(ckpt_path):
     model_data = ldm_patched.modules.utils.load_torch_file(ckpt_path, safe_load=True)
     keys = model_data.keys()
     if "style_embedding" in keys:
-        model = ldm_patched.modules.t2i_adapter.adapter.StyleAdapter(width=1024, context_dim=768, num_head=8, n_layes=3, num_token=8)
+        model = ldm_patched.t2ia.adapter.StyleAdapter(width=1024, context_dim=768, num_head=8, n_layes=3, num_token=8)
     else:
         raise Exception("invalid style model {}".format(ckpt_path))
     model.load_state_dict(model_data)
     return StyleModel(model)
 
-class CLIPType(Enum):
-    STABLE_DIFFUSION = 1
-    STABLE_CASCADE = 2
 
-def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DIFFUSION):
+def load_clip(ckpt_paths, embedding_directory=None):
     clip_data = []
     for p in ckpt_paths:
         clip_data.append(ldm_patched.modules.utils.load_torch_file(p, safe_load=True))
@@ -374,21 +308,14 @@ def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DI
 
     for i in range(len(clip_data)):
         if "transformer.resblocks.0.ln_1.weight" in clip_data[i]:
-            clip_data[i] = ldm_patched.modules.utils.clip_text_transformers_convert(clip_data[i], "", "")
-        else:
-            if "text_projection" in clip_data[i]:
-                clip_data[i]["text_projection.weight"] = clip_data[i]["text_projection"].transpose(0, 1) #old models saved with the CLIPSave node
+            clip_data[i] = ldm_patched.modules.utils.transformers_convert(clip_data[i], "", "text_model.", 32)
 
     clip_target = EmptyClass()
     clip_target.params = {}
     if len(clip_data) == 1:
         if "text_model.encoder.layers.30.mlp.fc1.weight" in clip_data[0]:
-            if clip_type == CLIPType.STABLE_CASCADE:
-                clip_target.clip = sdxl_clip.StableCascadeClipModel
-                clip_target.tokenizer = sdxl_clip.StableCascadeTokenizer
-            else:
-                clip_target.clip = sdxl_clip.SDXLRefinerClipModel
-                clip_target.tokenizer = sdxl_clip.SDXLTokenizer
+            clip_target.clip = sdxl_clip.SDXLRefinerClipModel
+            clip_target.tokenizer = sdxl_clip.SDXLTokenizer
         elif "text_model.encoder.layers.22.mlp.fc1.weight" in clip_data[0]:
             clip_target.clip = sd2_clip.SD2ClipModel
             clip_target.tokenizer = sd2_clip.SD2Tokenizer
@@ -403,10 +330,10 @@ def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DI
     for c in clip_data:
         m, u = clip.load_sd(c)
         if len(m) > 0:
-            logging.warning("clip missing: {}".format(m))
+            print("clip missing:", m)
 
         if len(u) > 0:
-            logging.debug("clip unexpected: {}".format(u))
+            print("clip unexpected:", u)
     return clip
 
 def load_gligen(ckpt_path):
@@ -417,8 +344,6 @@ def load_gligen(ckpt_path):
     return ldm_patched.modules.model_patcher.ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())
 
 def load_checkpoint(config_path=None, ckpt_path=None, output_vae=True, output_clip=True, embedding_directory=None, state_dict=None, config=None):
-    logging.warning("Warning: The load checkpoint with config function is deprecated and will eventually be removed, please use the other one.")
-    model, clip, vae, _ = load_checkpoint_guess_config(ckpt_path, output_vae=output_vae, output_clip=output_clip, output_clipvision=False, embedding_directory=embedding_directory, output_model=True)
     #TODO: this function is a mess and should be removed eventually
     if config is None:
         with open(config_path, 'r') as stream:
@@ -426,20 +351,81 @@ def load_checkpoint(config_path=None, ckpt_path=None, output_vae=True, output_cl
     model_config_params = config['model']['params']
     clip_config = model_config_params['cond_stage_config']
     scale_factor = model_config_params['scale_factor']
+    vae_config = model_config_params['first_stage_config']
+
+    fp16 = False
+    if "unet_config" in model_config_params:
+        if "params" in model_config_params["unet_config"]:
+            unet_config = model_config_params["unet_config"]["params"]
+            if "use_fp16" in unet_config:
+                fp16 = unet_config.pop("use_fp16")
+                if fp16:
+                    unet_config["dtype"] = torch.float16
+
+    noise_aug_config = None
+    if "noise_aug_config" in model_config_params:
+        noise_aug_config = model_config_params["noise_aug_config"]
+
+    model_type = model_base.ModelType.EPS
 
     if "parameterization" in model_config_params:
         if model_config_params["parameterization"] == "v":
-            m = model.clone()
-            class ModelSamplingAdvanced(ldm_patched.modules.model_sampling.ModelSamplingDiscrete, ldm_patched.modules.model_sampling.V_PREDICTION):
-                pass
-            m.add_object_patch("model_sampling", ModelSamplingAdvanced(model.model.model_config))
-            model = m
+            model_type = model_base.ModelType.V_PREDICTION
 
-    layer_idx = clip_config.get("params", {}).get("layer_idx", None)
-    if layer_idx is not None:
-        clip.clip_layer(layer_idx)
+    clip = None
+    vae = None
 
-    return (model, clip, vae)
+    class WeightsLoader(torch.nn.Module):
+        pass
+
+    if state_dict is None:
+        state_dict = ldm_patched.modules.utils.load_torch_file(ckpt_path)
+
+    class EmptyClass:
+        pass
+
+    model_config = ldm_patched.modules.supported_models_base.BASE({})
+
+    from . import latent_formats
+    model_config.latent_format = latent_formats.SD15(scale_factor=scale_factor)
+    model_config.unet_config = model_detection.convert_config(unet_config)
+
+    if config['model']["target"].endswith("ImageEmbeddingConditionedLatentDiffusion"):
+        model = model_base.SD21UNCLIP(model_config, noise_aug_config["params"], model_type=model_type)
+    else:
+        model = model_base.BaseModel(model_config, model_type=model_type)
+
+    if config['model']["target"].endswith("LatentInpaintDiffusion"):
+        model.set_inpaint()
+
+    if fp16:
+        model = model.half()
+
+    offload_device = model_management.unet_offload_device()
+    model = model.to(offload_device)
+    model.load_model_weights(state_dict, "model.diffusion_model.")
+
+    if output_vae:
+        vae_sd = ldm_patched.modules.utils.state_dict_prefix_replace(state_dict, {"first_stage_model.": ""}, filter_keys=True)
+        vae = VAE(sd=vae_sd, config=vae_config)
+
+    if output_clip:
+        w = WeightsLoader()
+        clip_target = EmptyClass()
+        clip_target.params = clip_config.get("params", {})
+        if clip_config["target"].endswith("FrozenOpenCLIPEmbedder"):
+            clip_target.clip = sd2_clip.SD2ClipModel
+            clip_target.tokenizer = sd2_clip.SD2Tokenizer
+            clip = CLIP(clip_target, embedding_directory=embedding_directory)
+            w.cond_stage_model = clip.cond_stage_model.clip_h
+        elif clip_config["target"].endswith("FrozenCLIPEmbedder"):
+            clip_target.clip = sd1_clip.SD1ClipModel
+            clip_target.tokenizer = sd1_clip.SD1Tokenizer
+            clip = CLIP(clip_target, embedding_directory=embedding_directory)
+            w.cond_stage_model = clip.cond_stage_model.clip_l
+        load_clip_weights(w, state_dict)
+
+    return (ldm_patched.modules.model_patcher.ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=offload_device), clip, vae)
 
 def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True, vae_filename_param=None):
     sd = ldm_patched.modules.utils.load_torch_file(ckpt_path)
@@ -453,12 +439,15 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
     clip_target = None
 
     parameters = ldm_patched.modules.utils.calculate_parameters(sd, "model.diffusion_model.")
+    unet_dtype = model_management.unet_dtype(model_params=parameters)
     load_device = model_management.get_torch_device()
+    manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device)
 
-    model_config = model_detection.model_config_from_unet(sd, "model.diffusion_model.")
-    unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=model_config.supported_inference_dtypes)
-    manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
-    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
+    class WeightsLoader(torch.nn.Module):
+        pass
+
+    model_config = model_detection.model_config_from_unet(sd, "model.diffusion_model.", unet_dtype)
+    model_config.set_manual_cast(manual_cast_dtype)
 
     if model_config is None:
         raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))
@@ -475,7 +464,7 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
 
     if output_vae:
         if vae_filename_param is None:
-            vae_sd = ldm_patched.modules.utils.state_dict_prefix_replace(sd, {k: "" for k in model_config.vae_key_prefix}, filter_keys=True)
+            vae_sd = ldm_patched.modules.utils.state_dict_prefix_replace(sd, {"first_stage_model.": ""}, filter_keys=True)
             vae_sd = model_config.process_vae_state_dict(vae_sd)
         else:
             vae_sd = ldm_patched.modules.utils.load_torch_file(vae_filename_param)
@@ -483,50 +472,41 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
         vae = VAE(sd=vae_sd)
 
     if output_clip:
+        w = WeightsLoader()
         clip_target = model_config.clip_target()
         if clip_target is not None:
-            clip_sd = model_config.process_clip_state_dict(sd)
-            if len(clip_sd) > 0:
-                clip = CLIP(clip_target, embedding_directory=embedding_directory)
-                m, u = clip.load_sd(clip_sd, full_model=True)
-                if len(m) > 0:
-                    m_filter = list(filter(lambda a: ".logit_scale" not in a and ".transformer.text_projection.weight" not in a, m))
-                    if len(m_filter) > 0:
-                        logging.warning("clip missing: {}".format(m))
-                    else:
-                        logging.debug("clip missing: {}".format(m))
-
-                if len(u) > 0:
-                    logging.debug("clip unexpected {}:".format(u))
-            else:
-                logging.warning("no CLIP/text encoder weights in checkpoint, the text encoder model will not be loaded.")
+            clip = CLIP(clip_target, embedding_directory=embedding_directory)
+            w.cond_stage_model = clip.cond_stage_model
+            sd = model_config.process_clip_state_dict(sd)
+            load_model_weights(w, sd)
 
     left_over = sd.keys()
     if len(left_over) > 0:
-        logging.debug("left over keys: {}".format(left_over))
+        print("left over keys:", left_over)
 
     if output_model:
         model_patcher = ldm_patched.modules.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=model_management.unet_offload_device(), current_device=inital_load_device)
         if inital_load_device != torch.device("cpu"):
-            logging.info("loaded straight to GPU")
+            print("loaded straight to GPU")
             model_management.load_model_gpu(model_patcher)
 
-    return (model_patcher, clip, vae, vae_filename, clipvision)
+    return model_patcher, clip, vae, vae_filename, clipvision
 
 
 def load_unet_state_dict(sd): #load unet in diffusers format
     parameters = ldm_patched.modules.utils.calculate_parameters(sd)
     unet_dtype = model_management.unet_dtype(model_params=parameters)
     load_device = model_management.get_torch_device()
+    manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device)
 
-    if "input_blocks.0.0.weight" in sd or 'clf.1.weight' in sd: #ldm or stable cascade
-        model_config = model_detection.model_config_from_unet(sd, "")
+    if "input_blocks.0.0.weight" in sd: #ldm
+        model_config = model_detection.model_config_from_unet(sd, "", unet_dtype)
         if model_config is None:
             return None
         new_sd = sd
 
     else: #diffusers
-        model_config = model_detection.model_config_from_diffusers_unet(sd)
+        model_config = model_detection.model_config_from_diffusers_unet(sd, unet_dtype)
         if model_config is None:
             return None
 
@@ -537,39 +517,33 @@ def load_unet_state_dict(sd): #load unet in diffusers format
             if k in sd:
                 new_sd[diffusers_keys[k]] = sd.pop(k)
             else:
-                logging.warning("{} {}".format(diffusers_keys[k], k))
-
+                print(diffusers_keys[k], k)
     offload_device = model_management.unet_offload_device()
-    unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=model_config.supported_inference_dtypes)
-    manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
-    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
+    model_config.set_manual_cast(manual_cast_dtype)
     model = model_config.get_model(new_sd, "")
     model = model.to(offload_device)
     model.load_model_weights(new_sd, "")
     left_over = sd.keys()
     if len(left_over) > 0:
-        logging.info("left over keys in unet: {}".format(left_over))
+        print("left over keys in unet:", left_over)
     return ldm_patched.modules.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=offload_device)
 
 def load_unet(unet_path):
     sd = ldm_patched.modules.utils.load_torch_file(unet_path)
     model = load_unet_state_dict(sd)
     if model is None:
-        logging.error("ERROR UNSUPPORTED UNET {}".format(unet_path))
+        print("ERROR UNSUPPORTED UNET", unet_path)
         raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path))
     return model
 
-def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, metadata=None, extra_keys={}):
+def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, metadata=None):
     clip_sd = None
     load_models = [model]
     if clip is not None:
         load_models.append(clip.load_model())
         clip_sd = clip.get_sd()
 
-    model_management.load_models_gpu(load_models, force_patch_weights=True)
+    model_management.load_models_gpu(load_models)
     clip_vision_sd = clip_vision.get_sd() if clip_vision is not None else None
     sd = model.model.state_dict_for_saving(clip_sd, vae.get_sd(), clip_vision_sd)
-    for k in extra_keys:
-        sd[k] = extra_keys[k]
-
     ldm_patched.modules.utils.save_torch_file(sd, output_path, metadata=metadata)

ldm_patched/modules/sd1_clip.py

@@ -8,7 +8,6 @@ import zipfile
 from . import model_management
 import ldm_patched.modules.clip_model
 import json
-import logging
 
 def gen_empty_tokens(special_tokens, length):
     start_token = special_tokens.get("start", None)
@@ -68,7 +67,7 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
     ]
     def __init__(self, version="openai/clip-vit-large-patch14", device="cpu", max_length=77,
                  freeze=True, layer="last", layer_idx=None, textmodel_json_config=None, dtype=None, model_class=ldm_patched.modules.clip_model.CLIPTextModel,
-                 special_tokens={"start": 49406, "end": 49407, "pad": 49407}, layer_norm_hidden_state=True, enable_attention_masks=False, return_projected_pooled=True):  # clip-vit-base-patch32
+                 special_tokens={"start": 49406, "end": 49407, "pad": 49407}, layer_norm_hidden_state=True):  # clip-vit-base-patch32
         super().__init__()
         assert layer in self.LAYERS
 
@@ -87,18 +86,16 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
         self.layer = layer
         self.layer_idx = None
         self.special_tokens = special_tokens
-
+        self.text_projection = torch.nn.Parameter(torch.eye(self.transformer.get_input_embeddings().weight.shape[1]))
         self.logit_scale = torch.nn.Parameter(torch.tensor(4.6055))
-        self.enable_attention_masks = enable_attention_masks
+        self.enable_attention_masks = False
 
         self.layer_norm_hidden_state = layer_norm_hidden_state
-        self.return_projected_pooled = return_projected_pooled
-
         if layer == "hidden":
             assert layer_idx is not None
             assert abs(layer_idx) < self.num_layers
-            self.set_clip_options({"layer": layer_idx})
-        self.options_default = (self.layer, self.layer_idx, self.return_projected_pooled)
+            self.clip_layer(layer_idx)
+        self.layer_default = (self.layer, self.layer_idx)
 
     def freeze(self):
         self.transformer = self.transformer.eval()
@@ -106,19 +103,16 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
         for param in self.parameters():
             param.requires_grad = False
 
-    def set_clip_options(self, options):
-        layer_idx = options.get("layer", self.layer_idx)
-        self.return_projected_pooled = options.get("projected_pooled", self.return_projected_pooled)
-        if layer_idx is None or abs(layer_idx) > self.num_layers:
+    def clip_layer(self, layer_idx):
+        if abs(layer_idx) > self.num_layers:
             self.layer = "last"
         else:
             self.layer = "hidden"
             self.layer_idx = layer_idx
 
-    def reset_clip_options(self):
-        self.layer = self.options_default[0]
-        self.layer_idx = self.options_default[1]
-        self.return_projected_pooled = self.options_default[2]
+    def reset_clip_layer(self):
+        self.layer = self.layer_default[0]
+        self.layer_idx = self.layer_default[1]
 
     def set_up_textual_embeddings(self, tokens, current_embeds):
         out_tokens = []
@@ -138,7 +132,7 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
                         tokens_temp += [next_new_token]
                         next_new_token += 1
                     else:
-                        logging.warning("WARNING: shape mismatch when trying to apply embedding, embedding will be ignored {} != {}".format(y.shape[0], current_embeds.weight.shape[1]))
+                        print("WARNING: shape mismatch when trying to apply embedding, embedding will be ignored", y.shape[0], current_embeds.weight.shape[1])
             while len(tokens_temp) < len(x):
                 tokens_temp += [self.special_tokens["pad"]]
             out_tokens += [tokens_temp]
@@ -183,19 +177,23 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
         else:
             z = outputs[1]
 
-        pooled_output = None
-        if len(outputs) >= 3:
-            if not self.return_projected_pooled and len(outputs) >= 4 and outputs[3] is not None:
-                pooled_output = outputs[3].float()
-            elif outputs[2] is not None:
-                pooled_output = outputs[2].float()
+        if outputs[2] is not None:
+            pooled_output = outputs[2].float()
+        else:
+            pooled_output = None
 
+        if self.text_projection is not None and pooled_output is not None:
+            pooled_output = pooled_output.float().to(self.text_projection.device) @ self.text_projection.float()
         return z.float(), pooled_output
 
     def encode(self, tokens):
         return self(tokens)
 
     def load_sd(self, sd):
+        if "text_projection" in sd:
+            self.text_projection[:] = sd.pop("text_projection")
+        if "text_projection.weight" in sd:
+            self.text_projection[:] = sd.pop("text_projection.weight").transpose(0, 1)
         return self.transformer.load_state_dict(sd, strict=False)
 
 def parse_parentheses(string):
@@ -328,9 +326,11 @@ def load_embed(embedding_name, embedding_directory, embedding_size, embed_key=No
                 except:
                     embed_out = safe_load_embed_zip(embed_path)
             else:
-                embed = torch.load(embed_path, map_location="cpu")
+                embed = torch.load(embed_path, map_location="cpu", weights_only=True)
     except Exception as e:
-        logging.warning("{}\n\nerror loading embedding, skipping loading: {}".format(traceback.format_exc(), embedding_name))
+        print(traceback.format_exc())
+        print()
+        print("error loading embedding, skipping loading:", embedding_name)
         return None
 
     if embed_out is None:
@@ -354,12 +354,11 @@ def load_embed(embedding_name, embedding_directory, embedding_size, embed_key=No
     return embed_out
 
 class SDTokenizer:
-    def __init__(self, tokenizer_path=None, max_length=77, pad_with_end=True, embedding_directory=None, embedding_size=768, embedding_key='clip_l', tokenizer_class=CLIPTokenizer, has_start_token=True, pad_to_max_length=True, min_length=None):
+    def __init__(self, tokenizer_path=None, max_length=77, pad_with_end=True, embedding_directory=None, embedding_size=768, embedding_key='clip_l', tokenizer_class=CLIPTokenizer, has_start_token=True, pad_to_max_length=True):
         if tokenizer_path is None:
             tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_tokenizer")
         self.tokenizer = tokenizer_class.from_pretrained(tokenizer_path)
         self.max_length = max_length
-        self.min_length = min_length
 
         empty = self.tokenizer('')["input_ids"]
         if has_start_token:
@@ -421,7 +420,7 @@ class SDTokenizer:
                     embedding_name = word[len(self.embedding_identifier):].strip('\n')
                     embed, leftover = self._try_get_embedding(embedding_name)
                     if embed is None:
-                        logging.warning(f"warning, embedding:{embedding_name} does not exist, ignoring")
+                        print(f"warning, embedding:{embedding_name} does not exist, ignoring")
                     else:
                         if len(embed.shape) == 1:
                             tokens.append([(embed, weight)])
@@ -471,8 +470,6 @@ class SDTokenizer:
         batch.append((self.end_token, 1.0, 0))
         if self.pad_to_max_length:
             batch.extend([(pad_token, 1.0, 0)] * (self.max_length - len(batch)))
-        if self.min_length is not None and len(batch) < self.min_length:
-            batch.extend([(pad_token, 1.0, 0)] * (self.min_length - len(batch)))
 
         if not return_word_ids:
             batched_tokens = [[(t, w) for t, w,_ in x] for x in batched_tokens]
@@ -506,11 +503,11 @@ class SD1ClipModel(torch.nn.Module):
         self.clip = "clip_{}".format(self.clip_name)
         setattr(self, self.clip, clip_model(device=device, dtype=dtype, **kwargs))
 
-    def set_clip_options(self, options):
-        getattr(self, self.clip).set_clip_options(options)
+    def clip_layer(self, layer_idx):
+        getattr(self, self.clip).clip_layer(layer_idx)
 
-    def reset_clip_options(self):
-        getattr(self, self.clip).reset_clip_options()
+    def reset_clip_layer(self):
+        getattr(self, self.clip).reset_clip_layer()
 
     def encode_token_weights(self, token_weight_pairs):
         token_weight_pairs = token_weight_pairs[self.clip_name]

ldm_patched/modules/sd2_clip.py

@@ -1,4 +1,5 @@
 from ldm_patched.modules import sd1_clip
+import torch
 import os
 
 class SD2ClipHModel(sd1_clip.SDClipModel):

ldm_patched/modules/sdxl_clip.py

@@ -40,13 +40,13 @@ class SDXLClipModel(torch.nn.Module):
         self.clip_l = sd1_clip.SDClipModel(layer="hidden", layer_idx=-2, device=device, dtype=dtype, layer_norm_hidden_state=False)
         self.clip_g = SDXLClipG(device=device, dtype=dtype)
 
-    def set_clip_options(self, options):
-        self.clip_l.set_clip_options(options)
-        self.clip_g.set_clip_options(options)
+    def clip_layer(self, layer_idx):
+        self.clip_l.clip_layer(layer_idx)
+        self.clip_g.clip_layer(layer_idx)
 
-    def reset_clip_options(self):
-        self.clip_g.reset_clip_options()
-        self.clip_l.reset_clip_options()
+    def reset_clip_layer(self):
+        self.clip_g.reset_clip_layer()
+        self.clip_l.reset_clip_layer()
 
     def encode_token_weights(self, token_weight_pairs):
         token_weight_pairs_g = token_weight_pairs["g"]
@@ -64,25 +64,3 @@ class SDXLClipModel(torch.nn.Module):
 class SDXLRefinerClipModel(sd1_clip.SD1ClipModel):
     def __init__(self, device="cpu", dtype=None):
         super().__init__(device=device, dtype=dtype, clip_name="g", clip_model=SDXLClipG)
-
-
-class StableCascadeClipGTokenizer(sd1_clip.SDTokenizer):
-    def __init__(self, tokenizer_path=None, embedding_directory=None):
-        super().__init__(tokenizer_path, pad_with_end=True, embedding_directory=embedding_directory, embedding_size=1280, embedding_key='clip_g')
-
-class StableCascadeTokenizer(sd1_clip.SD1Tokenizer):
-    def __init__(self, embedding_directory=None):
-        super().__init__(embedding_directory=embedding_directory, clip_name="g", tokenizer=StableCascadeClipGTokenizer)
-
-class StableCascadeClipG(sd1_clip.SDClipModel):
-    def __init__(self, device="cpu", max_length=77, freeze=True, layer="hidden", layer_idx=-1, dtype=None):
-        textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_config_bigg.json")
-        super().__init__(device=device, freeze=freeze, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype,
-                         special_tokens={"start": 49406, "end": 49407, "pad": 49407}, layer_norm_hidden_state=False, enable_attention_masks=True)
-
-    def load_sd(self, sd):
-        return super().load_sd(sd)
-
-class StableCascadeClipModel(sd1_clip.SD1ClipModel):
-    def __init__(self, device="cpu", dtype=None):
-        super().__init__(device=device, dtype=dtype, clip_name="g", clip_model=StableCascadeClipG)

ldm_patched/modules/supported_models.py

@@ -40,16 +40,11 @@ class SD15(supported_models_base.BASE):
                 state_dict['cond_stage_model.transformer.text_model.embeddings.position_ids'] = ids.round()
 
         replace_prefix = {}
-        replace_prefix["cond_stage_model."] = "clip_l."
-        state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=True)
+        replace_prefix["cond_stage_model."] = "cond_stage_model.clip_l."
+        state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix)
         return state_dict
 
     def process_clip_state_dict_for_saving(self, state_dict):
-        pop_keys = ["clip_l.transformer.text_projection.weight", "clip_l.logit_scale"]
-        for p in pop_keys:
-            if p in state_dict:
-                state_dict.pop(p)
-
         replace_prefix = {"clip_l.": "cond_stage_model."}
         return utils.state_dict_prefix_replace(state_dict, replace_prefix)
 
@@ -65,28 +60,22 @@ class SD20(supported_models_base.BASE):
         "use_temporal_attention": False,
     }
 
-    unet_extra_config = {
-        "num_heads": -1,
-        "num_head_channels": 64,
-        "attn_precision": torch.float32,
-    }
-
     latent_format = latent_formats.SD15
 
     def model_type(self, state_dict, prefix=""):
         if self.unet_config["in_channels"] == 4: #SD2.0 inpainting models are not v prediction
             k = "{}output_blocks.11.1.transformer_blocks.0.norm1.bias".format(prefix)
-            out = state_dict.get(k, None)
-            if out is not None and torch.std(out, unbiased=False) > 0.09: # not sure how well this will actually work. I guess we will find out.
+            out = state_dict[k]
+            if torch.std(out, unbiased=False) > 0.09: # not sure how well this will actually work. I guess we will find out.
                 return model_base.ModelType.V_PREDICTION
         return model_base.ModelType.EPS
 
     def process_clip_state_dict(self, state_dict):
         replace_prefix = {}
-        replace_prefix["conditioner.embedders.0.model."] = "clip_h." #SD2 in sgm format
-        replace_prefix["cond_stage_model.model."] = "clip_h."
-        state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=True)
-        state_dict = utils.clip_text_transformers_convert(state_dict, "clip_h.", "clip_h.transformer.")
+        replace_prefix["conditioner.embedders.0.model."] = "cond_stage_model.model." #SD2 in sgm format
+        state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix)
+
+        state_dict = utils.transformers_convert(state_dict, "cond_stage_model.model.", "cond_stage_model.clip_h.transformer.text_model.", 24)
         return state_dict
 
     def process_clip_state_dict_for_saving(self, state_dict):
@@ -142,10 +131,11 @@ class SDXLRefiner(supported_models_base.BASE):
     def process_clip_state_dict(self, state_dict):
         keys_to_replace = {}
         replace_prefix = {}
-        replace_prefix["conditioner.embedders.0.model."] = "clip_g."
-        state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=True)
 
-        state_dict = utils.clip_text_transformers_convert(state_dict, "clip_g.", "clip_g.transformer.")
+        state_dict = utils.transformers_convert(state_dict, "conditioner.embedders.0.model.", "cond_stage_model.clip_g.transformer.text_model.", 32)
+        keys_to_replace["conditioner.embedders.0.model.text_projection"] = "cond_stage_model.clip_g.text_projection"
+        keys_to_replace["conditioner.embedders.0.model.logit_scale"] = "cond_stage_model.clip_g.logit_scale"
+
         state_dict = utils.state_dict_key_replace(state_dict, keys_to_replace)
         return state_dict
 
@@ -174,18 +164,7 @@ class SDXL(supported_models_base.BASE):
     latent_format = latent_formats.SDXL
 
     def model_type(self, state_dict, prefix=""):
-        if 'edm_mean' in state_dict and 'edm_std' in state_dict: #Playground V2.5
-            self.latent_format = latent_formats.SDXL_Playground_2_5()
-            self.sampling_settings["sigma_data"] = 0.5
-            self.sampling_settings["sigma_max"] = 80.0
-            self.sampling_settings["sigma_min"] = 0.002
-            return model_base.ModelType.EDM
-        elif "edm_vpred.sigma_max" in state_dict:
-            self.sampling_settings["sigma_max"] = float(state_dict["edm_vpred.sigma_max"].item())
-            if "edm_vpred.sigma_min" in state_dict:
-                self.sampling_settings["sigma_min"] = float(state_dict["edm_vpred.sigma_min"].item())
-            return model_base.ModelType.V_PREDICTION_EDM
-        elif "v_pred" in state_dict:
+        if "v_pred" in state_dict:
             return model_base.ModelType.V_PREDICTION
         else:
             return model_base.ModelType.EPS
@@ -200,28 +179,26 @@ class SDXL(supported_models_base.BASE):
         keys_to_replace = {}
         replace_prefix = {}
 
-        replace_prefix["conditioner.embedders.0.transformer.text_model"] = "clip_l.transformer.text_model"
-        replace_prefix["conditioner.embedders.1.model."] = "clip_g."
-        state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=True)
+        replace_prefix["conditioner.embedders.0.transformer.text_model"] = "cond_stage_model.clip_l.transformer.text_model"
+        state_dict = utils.transformers_convert(state_dict, "conditioner.embedders.1.model.", "cond_stage_model.clip_g.transformer.text_model.", 32)
+        keys_to_replace["conditioner.embedders.1.model.text_projection"] = "cond_stage_model.clip_g.text_projection"
+        keys_to_replace["conditioner.embedders.1.model.text_projection.weight"] = "cond_stage_model.clip_g.text_projection"
+        keys_to_replace["conditioner.embedders.1.model.logit_scale"] = "cond_stage_model.clip_g.logit_scale"
 
+        state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix)
         state_dict = utils.state_dict_key_replace(state_dict, keys_to_replace)
-        state_dict = utils.clip_text_transformers_convert(state_dict, "clip_g.", "clip_g.transformer.")
         return state_dict
 
     def process_clip_state_dict_for_saving(self, state_dict):
         replace_prefix = {}
         keys_to_replace = {}
         state_dict_g = diffusers_convert.convert_text_enc_state_dict_v20(state_dict, "clip_g")
+        if "clip_g.transformer.text_model.embeddings.position_ids" in state_dict_g:
+            state_dict_g.pop("clip_g.transformer.text_model.embeddings.position_ids")
         for k in state_dict:
             if k.startswith("clip_l"):
                 state_dict_g[k] = state_dict[k]
 
-        state_dict_g["clip_l.transformer.text_model.embeddings.position_ids"] = torch.arange(77).expand((1, -1))
-        pop_keys = ["clip_l.transformer.text_projection.weight", "clip_l.logit_scale"]
-        for p in pop_keys:
-            if p in state_dict_g:
-                state_dict_g.pop(p)
-
         replace_prefix["clip_g"] = "conditioner.embedders.1.model"
         replace_prefix["clip_l"] = "conditioner.embedders.0"
         state_dict_g = utils.state_dict_prefix_replace(state_dict_g, replace_prefix)
@@ -250,26 +227,6 @@ class Segmind_Vega(SDXL):
         "use_temporal_attention": False,
     }
 
-class KOALA_700M(SDXL):
-    unet_config = {
-        "model_channels": 320,
-        "use_linear_in_transformer": True,
-        "transformer_depth": [0, 2, 5],
-        "context_dim": 2048,
-        "adm_in_channels": 2816,
-        "use_temporal_attention": False,
-    }
-
-class KOALA_1B(SDXL):
-    unet_config = {
-        "model_channels": 320,
-        "use_linear_in_transformer": True,
-        "transformer_depth": [0, 2, 6],
-        "context_dim": 2048,
-        "adm_in_channels": 2816,
-        "use_temporal_attention": False,
-    }
-
 class SVD_img2vid(supported_models_base.BASE):
     unet_config = {
         "model_channels": 320,
@@ -282,12 +239,6 @@ class SVD_img2vid(supported_models_base.BASE):
         "use_temporal_resblock": True
     }
 
-    unet_extra_config = {
-        "num_heads": -1,
-        "num_head_channels": 64,
-        "attn_precision": torch.float32,
-    }
-
     clip_vision_prefix = "conditioner.embedders.0.open_clip.model.visual."
 
     latent_format = latent_formats.SD15
@@ -301,41 +252,6 @@ class SVD_img2vid(supported_models_base.BASE):
     def clip_target(self):
         return None
 
-class SV3D_u(SVD_img2vid):
-    unet_config = {
-        "model_channels": 320,
-        "in_channels": 8,
-        "use_linear_in_transformer": True,
-        "transformer_depth": [1, 1, 1, 1, 1, 1, 0, 0],
-        "context_dim": 1024,
-        "adm_in_channels": 256,
-        "use_temporal_attention": True,
-        "use_temporal_resblock": True
-    }
-
-    vae_key_prefix = ["conditioner.embedders.1.encoder."]
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.SV3D_u(self, device=device)
-        return out
-
-class SV3D_p(SV3D_u):
-    unet_config = {
-        "model_channels": 320,
-        "in_channels": 8,
-        "use_linear_in_transformer": True,
-        "transformer_depth": [1, 1, 1, 1, 1, 1, 0, 0],
-        "context_dim": 1024,
-        "adm_in_channels": 1280,
-        "use_temporal_attention": True,
-        "use_temporal_resblock": True
-    }
-
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.SV3D_p(self, device=device)
-        return out
-
 class Stable_Zero123(supported_models_base.BASE):
     unet_config = {
         "context_dim": 768,
@@ -351,11 +267,6 @@ class Stable_Zero123(supported_models_base.BASE):
         "num_head_channels": -1,
     }
 
-    required_keys = {
-        "cc_projection.weight": None,
-        "cc_projection.bias": None,
-    }
-
     clip_vision_prefix = "cond_stage_model.model.visual."
 
     latent_format = latent_formats.SD15
@@ -395,99 +306,5 @@ class SD_X4Upscaler(SD20):
         out = model_base.SD_X4Upscaler(self, device=device)
         return out
 
-class Stable_Cascade_C(supported_models_base.BASE):
-    unet_config = {
-        "stable_cascade_stage": 'c',
-    }
-
-    unet_extra_config = {}
-
-    latent_format = latent_formats.SC_Prior
-    supported_inference_dtypes = [torch.bfloat16, torch.float32]
-
-    sampling_settings = {
-        "shift": 2.0,
-    }
-
-    vae_key_prefix = ["vae."]
-    text_encoder_key_prefix = ["text_encoder."]
-    clip_vision_prefix = "clip_l_vision."
-
-    def process_unet_state_dict(self, state_dict):
-        key_list = list(state_dict.keys())
-        for y in ["weight", "bias"]:
-            suffix = "in_proj_{}".format(y)
-            keys = filter(lambda a: a.endswith(suffix), key_list)
-            for k_from in keys:
-                weights = state_dict.pop(k_from)
-                prefix = k_from[:-(len(suffix) + 1)]
-                shape_from = weights.shape[0] // 3
-                for x in range(3):
-                    p = ["to_q", "to_k", "to_v"]
-                    k_to = "{}.{}.{}".format(prefix, p[x], y)
-                    state_dict[k_to] = weights[shape_from*x:shape_from*(x + 1)]
-        return state_dict
-
-    def process_clip_state_dict(self, state_dict):
-        state_dict = utils.state_dict_prefix_replace(state_dict, {k: "" for k in self.text_encoder_key_prefix}, filter_keys=True)
-        if "clip_g.text_projection" in state_dict:
-            state_dict["clip_g.transformer.text_projection.weight"] = state_dict.pop("clip_g.text_projection").transpose(0, 1)
-        return state_dict
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.StableCascade_C(self, device=device)
-        return out
-
-    def clip_target(self):
-        return supported_models_base.ClipTarget(sdxl_clip.StableCascadeTokenizer, sdxl_clip.StableCascadeClipModel)
-
-class Stable_Cascade_B(Stable_Cascade_C):
-    unet_config = {
-        "stable_cascade_stage": 'b',
-    }
-
-    unet_extra_config = {}
-
-    latent_format = latent_formats.SC_B
-    supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]
-
-    sampling_settings = {
-        "shift": 1.0,
-    }
-
-    clip_vision_prefix = None
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.StableCascade_B(self, device=device)
-        return out
-
-class SD15_instructpix2pix(SD15):
-    unet_config = {
-        "context_dim": 768,
-        "model_channels": 320,
-        "use_linear_in_transformer": False,
-        "adm_in_channels": None,
-        "use_temporal_attention": False,
-        "in_channels": 8,
-    }
-
-    def get_model(self, state_dict, prefix="", device=None):
-        return model_base.SD15_instructpix2pix(self, device=device)
-
-class SDXL_instructpix2pix(SDXL):
-    unet_config = {
-        "model_channels": 320,
-        "use_linear_in_transformer": True,
-        "transformer_depth": [0, 0, 2, 2, 10, 10],
-        "context_dim": 2048,
-        "adm_in_channels": 2816,
-        "use_temporal_attention": False,
-        "in_channels": 8,
-    }
-
-    def get_model(self, state_dict, prefix="", device=None):
-        return model_base.SDXL_instructpix2pix(self, model_type=self.model_type(state_dict, prefix), device=device)
-
-models = [Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p]
-
+models = [Stable_Zero123, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXLRefiner, SDXL, SSD1B, Segmind_Vega, SD_X4Upscaler]
 models += [SVD_img2vid]

ldm_patched/modules/supported_models_base.py

@@ -16,28 +16,19 @@ class BASE:
         "num_head_channels": 64,
     }
 
-    required_keys = {}
-
     clip_prefix = []
     clip_vision_prefix = None
     noise_aug_config = None
     sampling_settings = {}
     latent_format = latent_formats.LatentFormat
-    vae_key_prefix = ["first_stage_model."]
-    text_encoder_key_prefix = ["cond_stage_model."]
-    supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]
 
     manual_cast_dtype = None
 
     @classmethod
-    def matches(s, unet_config, state_dict=None):
+    def matches(s, unet_config):
         for k in s.unet_config:
-            if k not in unet_config or s.unet_config[k] != unet_config[k]:
+            if s.unet_config[k] != unet_config[k]:
                 return False
-        if state_dict is not None:
-            for k in s.required_keys:
-                if k not in state_dict:
-                    return False
         return True
 
     def model_type(self, state_dict, prefix=""):
@@ -47,8 +38,7 @@ class BASE:
         return self.unet_config["in_channels"] > 4
 
     def __init__(self, unet_config):
-        self.unet_config = unet_config.copy()
-        self.sampling_settings = self.sampling_settings.copy()
+        self.unet_config = unet_config
         self.latent_format = self.latent_format()
         for x in self.unet_extra_config:
             self.unet_config[x] = self.unet_extra_config[x]
@@ -63,7 +53,6 @@ class BASE:
         return out
 
     def process_clip_state_dict(self, state_dict):
-        state_dict = utils.state_dict_prefix_replace(state_dict, {k: "" for k in self.text_encoder_key_prefix}, filter_keys=True)
         return state_dict
 
     def process_unet_state_dict(self, state_dict):
@@ -73,7 +62,7 @@ class BASE:
         return state_dict
 
     def process_clip_state_dict_for_saving(self, state_dict):
-        replace_prefix = {"": self.text_encoder_key_prefix[0]}
+        replace_prefix = {"": "cond_stage_model."}
         return utils.state_dict_prefix_replace(state_dict, replace_prefix)
 
     def process_clip_vision_state_dict_for_saving(self, state_dict):
@@ -87,9 +76,8 @@ class BASE:
         return utils.state_dict_prefix_replace(state_dict, replace_prefix)
 
     def process_vae_state_dict_for_saving(self, state_dict):
-        replace_prefix = {"": self.vae_key_prefix[0]}
+        replace_prefix = {"": "first_stage_model."}
         return utils.state_dict_prefix_replace(state_dict, replace_prefix)
 
-    def set_inference_dtype(self, dtype, manual_cast_dtype):
-        self.unet_config['dtype'] = dtype
+    def set_manual_cast(self, manual_cast_dtype):
         self.manual_cast_dtype = manual_cast_dtype

ldm_patched/modules/types.py

@@ -1,32 +0,0 @@
-import torch
-from typing import Callable, Protocol, TypedDict, Optional, List
-
-
-class UnetApplyFunction(Protocol):
-    """Function signature protocol on comfy.model_base.BaseModel.apply_model"""
-
-    def __call__(self, x: torch.Tensor, t: torch.Tensor, **kwargs) -> torch.Tensor:
-        pass
-
-
-class UnetApplyConds(TypedDict):
-    """Optional conditions for unet apply function."""
-
-    c_concat: Optional[torch.Tensor]
-    c_crossattn: Optional[torch.Tensor]
-    control: Optional[torch.Tensor]
-    transformer_options: Optional[dict]
-
-
-class UnetParams(TypedDict):
-    # Tensor of shape [B, C, H, W]
-    input: torch.Tensor
-    # Tensor of shape [B]
-    timestep: torch.Tensor
-    c: UnetApplyConds
-    # List of [0, 1], [0], [1], ...
-    # 0 means conditional, 1 means conditional unconditional
-    cond_or_uncond: List[int]
-
-
-UnetWrapperFunction = Callable[[UnetApplyFunction, UnetParams], torch.Tensor]

ldm_patched/modules/utils.py

@@ -5,7 +5,6 @@ import ldm_patched.modules.checkpoint_pickle
 import safetensors.torch
 import numpy as np
 from PIL import Image
-import logging
 
 def load_torch_file(ckpt, safe_load=False, device=None):
     if device is None:
@@ -15,14 +14,14 @@ def load_torch_file(ckpt, safe_load=False, device=None):
     else:
         if safe_load:
             if not 'weights_only' in torch.load.__code__.co_varnames:
-                logging.warning("Warning torch.load doesn't support weights_only on this pytorch version, loading unsafely.")
+                print("Warning torch.load doesn't support weights_only on this pytorch version, loading unsafely.")
                 safe_load = False
         if safe_load:
             pl_sd = torch.load(ckpt, map_location=device, weights_only=True)
         else:
             pl_sd = torch.load(ckpt, map_location=device, pickle_module=ldm_patched.modules.checkpoint_pickle)
         if "global_step" in pl_sd:
-            logging.debug(f"Global Step: {pl_sd['global_step']}")
+            print(f"Global Step: {pl_sd['global_step']}")
         if "state_dict" in pl_sd:
             sd = pl_sd["state_dict"]
         else:
@@ -99,22 +98,8 @@ def transformers_convert(sd, prefix_from, prefix_to, number):
                     p = ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"]
                     k_to = "{}encoder.layers.{}.{}.{}".format(prefix_to, resblock, p[x], y)
                     sd[k_to] = weights[shape_from*x:shape_from*(x + 1)]
-
     return sd
 
-def clip_text_transformers_convert(sd, prefix_from, prefix_to):
-    sd = transformers_convert(sd, prefix_from, "{}text_model.".format(prefix_to), 32)
-
-    tp = "{}text_projection.weight".format(prefix_from)
-    if tp in sd:
-        sd["{}text_projection.weight".format(prefix_to)] = sd.pop(tp)
-
-    tp = "{}text_projection".format(prefix_from)
-    if tp in sd:
-        sd["{}text_projection.weight".format(prefix_to)] = sd.pop(tp).transpose(0, 1).contiguous()
-    return sd
-
-
 UNET_MAP_ATTENTIONS = {
     "proj_in.weight",
     "proj_in.bias",
@@ -184,8 +169,6 @@ UNET_MAP_BASIC = {
 }
 
 def unet_to_diffusers(unet_config):
-    if "num_res_blocks" not in unet_config:
-        return {}
     num_res_blocks = unet_config["num_res_blocks"]
     channel_mult = unet_config["channel_mult"]
     transformer_depth = unet_config["transformer_depth"][:]
@@ -295,11 +278,8 @@ def set_attr(obj, attr, value):
     for name in attrs[:-1]:
         obj = getattr(obj, name)
     prev = getattr(obj, attrs[-1])
-    setattr(obj, attrs[-1], value)
-    return prev
-
-def set_attr_param(obj, attr, value):
-    return set_attr(obj, attr, torch.nn.Parameter(value, requires_grad=False))
+    setattr(obj, attrs[-1], torch.nn.Parameter(value, requires_grad=False))
+    del prev
 
 def copy_to_param(obj, attr, value):
     # inplace update tensor instead of replacing it
@@ -433,8 +413,6 @@ def tiled_scale(samples, function, tile_x=64, tile_y=64, overlap = 8, upscale_am
         out_div = torch.zeros((s.shape[0], out_channels, round(s.shape[2] * upscale_amount), round(s.shape[3] * upscale_amount)), device=output_device)
         for y in range(0, s.shape[2], tile_y - overlap):
             for x in range(0, s.shape[3], tile_x - overlap):
-                x = max(0, min(s.shape[-1] - overlap, x))
-                y = max(0, min(s.shape[-2] - overlap, y))
                 s_in = s[:,:,y:y+tile_y,x:x+tile_x]
 
                 ps = function(s_in).to(output_device)

ldm_patched/pfn/architecture/face/codeformer.py

@@ -377,15 +377,15 @@ class VQAutoEncoder(nn.Module):
         )
 
         if model_path is not None:
-            chkpt = torch.load(model_path, map_location="cpu")
+            chkpt = torch.load(model_path, map_location="cpu", weights_only=True)
             if "params_ema" in chkpt:
                 self.load_state_dict(
-                    torch.load(model_path, map_location="cpu")["params_ema"]
+                    torch.load(model_path, map_location="cpu", weights_only=True)["params_ema"]
                 )
                 logger.info(f"vqgan is loaded from: {model_path} [params_ema]")
             elif "params" in chkpt:
                 self.load_state_dict(
-                    torch.load(model_path, map_location="cpu")["params"]
+                    torch.load(model_path, map_location="cpu", weights_only=True)["params"]
                 )
                 logger.info(f"vqgan is loaded from: {model_path} [params]")
             else:

ldm_patched/pfn/architecture/face/gfpgan_bilinear_arch.py

@@ -273,8 +273,8 @@ class GFPGANBilinear(nn.Module):
         if decoder_load_path:
             self.stylegan_decoder.load_state_dict(
                 torch.load(
-                    decoder_load_path, map_location=lambda storage, loc: storage
-                )["params_ema"]
+                    decoder_load_path, map_location=lambda storage, loc: storage, 
+                weights_only=True)["params_ema"]
             )
         # fix decoder without updating params
         if fix_decoder: