Merge 1f6ab7dbfb into ff838657fa

Signed-off-by: bigcat88 <bigcat88@icloud.com>

comfy/ldm/lightricks/model.py

@@ -456,9 +456,8 @@ class LTXVModel(torch.nn.Module):
         x = self.patchify_proj(x)
         timestep = timestep * 1000.0
 
-        attention_mask = 1.0 - attention_mask.to(x.dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1]))
-        attention_mask = attention_mask.masked_fill(attention_mask.to(torch.bool), float("-inf"))  # not sure about this
-        # attention_mask = (context != 0).any(dim=2).to(dtype=x.dtype)
+        if attention_mask is not None and not torch.is_floating_point(attention_mask):
+            attention_mask = (attention_mask - 1).to(x.dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1])) * torch.finfo(x.dtype).max
 
         pe = precompute_freqs_cis(indices_grid, dim=self.inner_dim, out_dtype=x.dtype)
 

main.py

@@ -166,6 +166,14 @@ def prompt_worker(q, server_instance):
         queue_item = q.get(timeout=timeout)
         if queue_item is not None:
             item, item_id = queue_item
+
+            if item[3].get("unload_models"):
+                # For those cases where the flag is set, to clear memory before execution
+                comfy.model_management.unload_all_models()
+                gc.collect()
+                comfy.model_management.soft_empty_cache()
+                last_gc_collect = time.perf_counter()
+
             execution_start_time = time.perf_counter()
             prompt_id = item[1]
             server_instance.last_prompt_id = prompt_id