Fix image corruption in half mode with embeddings.

(DirectML)

modules/devices.py

@@ -135,9 +135,6 @@ def set_cuda_params():
             except Exception:
                 pass
     global dtype, dtype_vae, dtype_unet, unet_needs_upcast # pylint: disable=global-statement
-    if shared.cmd_opts.use_directml and not shared.cmd_opts.experimental: # TODO DirectML does not have full autocast capabilities
-        shared.opts.no_half = True
-        shared.opts.no_half_vae = True
     if shared.opts.cuda_dtype == 'FP32':
         dtype = torch.float32
         dtype_vae = torch.float32

modules/dml/amp/autocast_mode.py

@@ -2,18 +2,25 @@ import importlib
 from typing import Any, Optional
 import torch
 
-ops = ["torch.Tensor.__matmul__", "torch.addbmm", "torch.addmm", "torch.addmv", "torch.addr", "torch.baddbmm", "torch.bmm", "torch.chain_matmul", "torch.linalg.multi_dot", "torch.nn.functional.conv1d", "torch.nn.functional.conv2d", "torch.nn.functional.conv3d", "torch.nn.functional.conv_transpose1d", "torch.nn.functional.conv_transpose2d", "torch.nn.functional.conv_transpose3d", "torch.nn.GRUCell", "torch.nn.functional.linear", "torch.nn.LSTMCell", "torch.matmul", "torch.mm", "torch.mv", "torch.prelu", "torch.nn.RNNCell"]
+ops = ["torch.Tensor.__matmul__", "torch.addbmm", "torch.addmm", "torch.addmv", "torch.addr", "torch.baddbmm", "torch.bmm", "torch.chain_matmul", "torch.linalg.multi_dot", "torch.nn.functional.conv1d", "torch.nn.functional.conv2d", "torch.nn.functional.conv3d", "torch.nn.functional.conv_transpose1d", "torch.nn.functional.conv_transpose2d", "torch.nn.functional.conv_transpose3d", "torch.nn.GRUCell", "torch.nn.functional.linear", "torch.nn.LSTMCell", "torch.matmul", "torch.mm", "torch.mv", "torch.prelu", "torch.nn.RNNCell", "torch.embedding"]
+supported_cast_pairs = {
+    torch.float16: (torch.float32,),
+    torch.float32: (torch.float16,),
+}
 
-def pre_forward(forward, args, kwargs):
+def forward(op, args: tuple, kwargs: dict):
     if not torch.dml.is_autocast_enabled:
-        return forward(*args, **kwargs)
+        return op(*args, **kwargs)
     args = list(map(cast, args))
-    for keyword in kwargs:
-        kwargs[keyword] = cast(kwargs[keyword])
-    return forward(*args, **kwargs)
+    for kwarg in kwargs:
+        kwargs[kwarg] = cast(kwargs[kwarg])
+    return op(*args, **kwargs)
 
-def cast(tensor):
-    if not isinstance(tensor, torch.Tensor) or torch.dml.autocast_gpu_dtype == tensor.dtype:
+def cast(tensor: torch.Tensor):
+    if not torch.is_tensor(tensor):
+        return tensor
+    dtype: torch.dtype = tensor.dtype
+    if dtype not in supported_cast_pairs or (torch.dml.autocast_gpu_dtype != dtype and torch.dml.autocast_gpu_dtype not in supported_cast_pairs[dtype]):
         return tensor
     return tensor.type(torch.dml.autocast_gpu_dtype)
 
@@ -29,21 +36,25 @@ def cond(op: str):
         for attr_name in func_path[i:-1]:
             resolved_obj = getattr(resolved_obj, attr_name)
         op = getattr(resolved_obj, func_path[-1])
-        setattr(resolved_obj, func_path[-1], lambda *args, **kwargs: pre_forward(op, args, kwargs))
+        setattr(resolved_obj, func_path[-1], lambda *args, **kwargs: forward(op, args, kwargs))
 
 for op in ops:
     cond(op)
 
 class autocast:
-    def __init__(self, dtype: Optional[torch.dtype] = None):
-        self.fast_dtype = dtype or torch.dml.autocast_gpu_dtype
+    prev: bool
+
+    fast_dtype: torch.dtype = torch.float16
+    prev_fast_dtype: torch.dtype
+    def __init__(self, dtype: Optional[torch.dtype] = torch.float16):
+        self.fast_dtype = dtype
 
     def __enter__(self):
         self.prev = torch.dml.is_autocast_enabled
-        self.prev_fastdtype = torch.dml.autocast_gpu_dtype
+        self.prev_fast_dtype = torch.dml.autocast_gpu_dtype
         torch.dml.is_autocast_enabled = True
         torch.dml.autocast_gpu_dtype = self.fast_dtype
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):
         torch.dml.is_autocast_enabled = self.prev
-        torch.dml.autocast_gpu_dtype = self.prev_fastdtype
+        torch.dml.autocast_gpu_dtype = self.prev_fast_dtype

modules/dml/backend.py

@@ -61,7 +61,7 @@ class DirectML:
         return f"privateuseone:{get_device(device).index}"
 
     def get_device_name(device: Optional[rDevice]=None) -> str:
-        return torch_directml.device_name(get_device(device))
+        return torch_directml.device_name(get_device(device).index)
 
     def get_device_properties(device: Optional[rDevice]=None) -> DeviceProperties:
         return DeviceProperties(get_device(device))
@@ -84,8 +84,7 @@ class DirectML:
         return (0 if available < 0 else available, DirectML.__gpu_memory_bound)
 
     def memory_allocated(device: Optional[rDevice]=None) -> int:
-        device = get_device(device)
-        return sum(torch_directml.gpu_memory(device.index)) * (1 << 20)
+        return sum(torch_directml.gpu_memory(get_device(device).index)) * (1 << 20)
 
     def max_memory_allocated(device: Optional[rDevice]=None):
         return DirectML.memory_allocated(device) # DirectML does not empty GPU memory

modules/dml/hijack/transformers.py

@@ -1,4 +1,5 @@
 import torch
+from typing import Optional
 import transformers.models.clip.modeling_clip
 
 # Copied from transformers.models.bart.modeling_bart._make_causal_mask
@@ -19,4 +20,25 @@ def _make_causal_mask(
         mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
     return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
 
+def CLIPTextEmbeddings_forward(
+    self: transformers.models.clip.modeling_clip.CLIPTextEmbeddings,
+    input_ids: Optional[torch.LongTensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+) -> torch.Tensor:
+    from modules.devices import dtype
+    seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2]
+
+    if position_ids is None:
+        position_ids = self.position_ids[:, :seq_length]
+
+    if inputs_embeds is None:
+        inputs_embeds = self.token_embedding(input_ids).type(dtype) # Type correction.
+
+    position_embeddings = self.position_embedding(position_ids)
+    embeddings = inputs_embeds + position_embeddings
+
+    return embeddings
+
 transformers.models.clip.modeling_clip._make_causal_mask = _make_causal_mask
+transformers.models.clip.modeling_clip.CLIPTextEmbeddings.forward = CLIPTextEmbeddings_forward