diff --git a/.pylintrc b/.pylintrc index bef8906ce..2c9bbc0c2 100644 --- a/.pylintrc +++ b/.pylintrc @@ -23,6 +23,7 @@ ignore-paths=/usr/lib/.*$, modules/k-diffusion, modules/flex2, modules/ldsr, + modules/hidream, modules/meissonic, modules/mod, modules/omnigen, diff --git a/.ruff.toml b/.ruff.toml index 96fd1323a..ab5e0601c 100644 --- a/.ruff.toml +++ b/.ruff.toml @@ -20,6 +20,7 @@ exclude = [ "modules/meissonic", "modules/mod", "modules/omnigen", + "modules/hidream", "modules/pag", "modules/pixelsmith", "modules/postprocess/aurasr_arch.py", diff --git a/CHANGELOG.md b/CHANGELOG.md index e8842f90e..c19c9d288 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,44 @@ # Change Log for SD.Next +## Update for 2025-05-06 + +- **Features** + - [FramePack](https://vladmandic.github.io/sdnext-docs/FramePack) + add **T2V** mode in addition to **I2V** and **FLF2V** + support for new **F1: forward-only** model variant in addition to regular **bi-directional** + add **prompt enhance** using VLM: it will analyze input image and then create enhanced prompt based on user prompt and image + add **prompt interpolation**, section prompts do not need to match exact video section count + and improved performance + [Docs](https://vladmandic.github.io/sdnext-docs/FramePack) rewrite! + - **Prompt-Enhhance** + add **Qwen3** *0.6B/1.7B/4B* models + add thinking mode support (for models that have it) + - [HiDream-E1](https://huggingface.co/HiDream-ai/HiDream-E1-Full) natural language image-editing model built on HiDream-I1 + available via *networks -> models -> reference* + *note*: right now HiDream-E1 is limited to 768x768 images, so you must force resize image before running it +- **Other** + - CUDA: set default to `torch==2.7.0` with `cuda==12.8` + - ZLUDA: update to `zluda==3.9.4` and `flash-attn-2` + - Docker: pre-install `ffmpeg` + - Wiki: updated pages: *FramePack, Video, ROCm, ZLUDA, Quantization* + - Gallery: support JXL image format + - Scheduler: add sigmoid beta scheduler + - GitHub: updated issue template +- **Fixes** + - FramePack: correct dtype + - NNCF: check dependencies and register quant type + - API: refresh checkpoint list + - API: vlm-api endpoint + - Styles: save style with prompt + - Texture tiling: fix apply when switching models + - Diffusers: slow initial startup + - Gated access: obfuscate and log token used for access + - SDXL refiner workflow + - Control: t2i-adapter workflow + - Control: xs-controlnet workflow + - Control: lllite-workflow + - Control: refiner workflow with multiple control elements + ## Highlights for 2025-04-28 Another major release with *over 120 commits*! diff --git a/TODO.md b/TODO.md index c1e467642..e4d663f18 100644 --- a/TODO.md +++ b/TODO.md @@ -4,6 +4,7 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma ## Current +- Video: API support - ModernUI for custom model loader - ModernUI for history tab @@ -13,10 +14,8 @@ N/A ## Future Candidates -- IPAdapter: negative guidance: - Control: API enhance scripts compatibility -- Video: add generate context menu -- Video: API support +- IPAdapter: negative guidance: - Video: STG: - Video: SmoothCache: https://github.com/huggingface/diffusers/issues/11135 @@ -27,9 +26,7 @@ N/A - control: support scripts via api - fc: autodetect distilled based on model - fc: autodetect tensor format based on model -- hidream: pack latents for remote vae - hypertile: vae breaks when using non-standard sizes -- infotext: handle using regex instead - install: enable ROCm for windows when available - loader: load receipe - loader: save receipe diff --git a/configs/Dockerfile.cuda b/configs/Dockerfile.cuda index 1534647af..ba6d53af0 100644 --- a/configs/Dockerfile.cuda +++ b/configs/Dockerfile.cuda @@ -18,7 +18,7 @@ LABEL org.opencontainers.image.version="latest" # minimum install RUN ["apt-get", "-y", "update"] -RUN ["apt-get", "-y", "install", "git", "build-essential", "google-perftools", "curl"] +RUN ["apt-get", "-y", "install", "git", "build-essential", "google-perftools", "curl", "ffmpeg"] # optional if full cuda-dev is required by some downstream library # RUN ["apt-get", "-y", "nvidia-cuda-toolkit"] RUN ["/usr/sbin/ldconfig"] diff --git a/html/reference.json b/html/reference.json index 822c64598..2441112e6 100644 --- a/html/reference.json +++ b/html/reference.json @@ -364,6 +364,13 @@ "skip": true, "extras": "sampler: Default" }, + "HiDream-E1 Full": { + "path": "HiDream-ai/HiDream-E1-Full", + "desc": "HiDream-E1 is an image editing model built on HiDream-I1.", + "preview": "HiDream-ai--HiDream-I1-Fast.jpg", + "skip": true, + "extras": "sampler: Default" + }, "Kwai Kolors": { "path": "Kwai-Kolors/Kolors-diffusers", diff --git a/installer.py b/installer.py index 9b7d824bb..4a1eafbbe 100644 --- a/installer.py +++ b/installer.py @@ -508,17 +508,23 @@ def get_platform(): # check python version -def check_python(supported_minors=[9, 10, 11, 12], reason=None): +def check_python(supported_minors=[], experimental_minors=[], reason=None): + if supported_minors is None or len(supported_minors) == 0: + supported_minors = [9, 10, 11, 12] + experimental_minors = [] t_start = time.time() if args.quick: return log.info(f'Python: version={platform.python_version()} platform={platform.system()} bin="{sys.executable}" venv="{sys.prefix}"') if not (int(sys.version_info.major) == 3 and int(sys.version_info.minor) in supported_minors): - log.error(f"Python version incompatible: {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro} required 3.{supported_minors}") - if reason is not None: - log.error(reason) - if not args.ignore and not args.experimental: - sys.exit(1) + if (int(sys.version_info.major) == 3 and int(sys.version_info.minor) in experimental_minors): + log.warning(f"Python version experimental: {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro} recommended 3.{supported_minors}") + else: + log.error(f"Python version incompatible: {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro} required 3.{supported_minors}") + if reason is not None: + log.error(reason) + if not args.ignore and not args.experimental: + sys.exit(1) if int(sys.version_info.minor) == 12: os.environ.setdefault('SETUPTOOLS_USE_DISTUTILS', 'local') # hack for python 3.11 setuptools if not args.skip_git: @@ -538,7 +544,7 @@ def check_diffusers(): t_start = time.time() if args.skip_all or args.skip_git or args.experimental: return - sha = 'b4be42282dc9bf9deccc2f60f5db952de772cf42' # diffusers commit hash + sha = '8c661ea586bf11cb2440da740dd3c4cf84679b85' # diffusers commit hash pkg = pkg_resources.working_set.by_key.get('diffusers', None) minor = int(pkg.version.split('.')[1] if pkg is not None else 0) cur = opts.get('diffusers_version', '') if minor > 0 else '' @@ -573,7 +579,8 @@ def install_cuda(): if args.use_nightly: cmd = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu128 --extra-index-url https://download.pytorch.org/whl/nightly/cu126') else: - cmd = os.environ.get('TORCH_COMMAND', 'torch==2.6.0+cu126 torchvision==0.21.0+cu126 --index-url https://download.pytorch.org/whl/cu126') + # cmd = os.environ.get('TORCH_COMMAND', 'torch==2.6.0+cu126 torchvision==0.21.0+cu126 --index-url https://download.pytorch.org/whl/cu126') + cmd = os.environ.get('TORCH_COMMAND', 'torch==2.7.0+cu128 torchvision==0.22.0+cu128 --index-url https://download.pytorch.org/whl/cu128') return cmd @@ -1090,7 +1097,13 @@ def install_submodules(force=True): return '\n'.join(res) -def reload(package): +def reload(package, desired=None): + loaded = package in sys.modules + if not loaded: + return + current = sys.modules[package].__version__ if hasattr(sys.modules[package], "__version__") else None + if desired is not None and current == desired: + return modules = [m for m in sys.modules if m.startswith(package)] for m in modules: del sys.modules[m] @@ -1142,17 +1155,17 @@ def install_optional(): install('pillow-jxl-plugin==1.3.2', ignore=True) install('optimum-quanto==0.2.7', ignore=True) install('torchao==0.10.0', ignore=True) - install('bitsandbytes==0.45.1', ignore=True) + install('bitsandbytes==0.45.5', ignore=True) install('pynvml', ignore=True) install('ultralytics==8.3.40', ignore=True) install('Cython', ignore=True) install('insightface==0.7.3', ignore=True) # problematic build install('albumentations==1.4.3', ignore=True) install('pydantic==1.10.21', ignore=True) - reload('pydantic') - install('nncf==2.16.0', ignore=True, no_deps=True) # requires older pandas - # install('flash-attn', ignore=True) # requires cuda and nvcc to be installed + reload('pydantic', '1.10.21') + install('nncf==2.16.0', ignore=True) # requires older pandas install('gguf', ignore=True) + install('av', ignore=True) try: import gguf scripts_dir = os.path.join(os.path.dirname(gguf.__file__), '..', 'scripts') diff --git a/javascript/gallery.js b/javascript/gallery.js index 39649d5dc..cea21fa12 100644 --- a/javascript/gallery.js +++ b/javascript/gallery.js @@ -126,7 +126,7 @@ class GalleryFile extends HTMLElement { async connectedCallback() { if (this.shadow.children.length > 0) return; const ext = this.name.split('.').pop().toLowerCase(); - if (!['jpg', 'jpeg', 'png', 'gif', 'webp', 'svg', 'mp4'].includes(ext)) return; + if (!['jpg', 'jpeg', 'png', 'gif', 'webp', 'jxl', 'svg', 'mp4'].includes(ext)) return; this.hash = await getHash(`${this.folder}/${this.name}/${this.size}/${this.mtime}`); // eslint-disable-line no-use-before-define const style = document.createElement('style'); const width = opts.browser_fixed_width ? `${opts.extra_networks_card_size}px` : 'unset'; diff --git a/modules/api/endpoints.py b/modules/api/endpoints.py index d5a3d8708..1decd8c71 100644 --- a/modules/api/endpoints.py +++ b/modules/api/endpoints.py @@ -17,8 +17,11 @@ def get_upscalers(): return [{"name": upscaler.name, "model_name": upscaler.scaler.model_name, "model_path": upscaler.data_path, "model_url": None, "scale": upscaler.scale} for upscaler in shared.sd_upscalers] def get_sd_models(): - from modules import sd_models, sd_models_config - return [{"title": x.title, "model_name": x.name, "filename": x.filename, "type": x.type, "hash": x.shorthash, "sha256": x.sha256, "config": sd_models_config.find_checkpoint_config_near_filename(x)} for x in sd_models.checkpoints_list.values()] + from modules import sd_checkpoint, sd_models_config + checkpoints = [] + for v in sd_checkpoint.checkpoints_list.values(): + checkpoints.append({"title": v.title, "model_name": v.name, "filename": v.filename, "type": v.type, "hash": v.shorthash, "sha256": v.sha256, "config": sd_models_config.find_checkpoint_config_near_filename(v)}) + return checkpoints def get_hypernetworks(): return [{"name": name, "path": shared.hypernetworks[name]} for name in shared.hypernetworks] @@ -128,10 +131,12 @@ def post_reload_checkpoint(): return {} def post_refresh_checkpoints(): - return shared.refresh_checkpoints() + shared.refresh_checkpoints() + return {} def post_refresh_vae(): - return shared.refresh_vaes() + shared.refresh_vaes() + return {} def post_refresh_loras(): from modules.lora import lora_load diff --git a/modules/api/helpers.py b/modules/api/helpers.py index da826a81e..21e89c5c1 100644 --- a/modules/api/helpers.py +++ b/modules/api/helpers.py @@ -24,8 +24,8 @@ def decode_base64_to_image(encoding, quiet=False): return image except Exception as e: shared.log.warning(f'API cannot decode image: {e}') - from modules import errors - errors.display(e, 'API cannot decode image') + # from modules import errors + # errors.display(e, 'API cannot decode image') if not quiet: raise HTTPException(status_code=500, detail="Invalid encoded image") from e return None diff --git a/modules/api/models.py b/modules/api/models.py index a67f42b7f..fee8d781d 100644 --- a/modules/api/models.py +++ b/modules/api/models.py @@ -355,6 +355,7 @@ class ReqVQA(BaseModel): image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.") model: str = Field(default="Microsoft Florence 2 Base", title="Model", description="The interrogate model used.") question: str = Field(default="describe the image", title="Question", description="Question to ask the model.") + system: str = Field(default="You are image captioning expert, creative, unbiased and uncensored.", title="System prompt", description="Prompt to shape how the model interprets and responds to user prompts.") class ReqLatentHistory(BaseModel): name: str = Field(title="Name", description="Name of the history item to select") diff --git a/modules/cmd_args.py b/modules/cmd_args.py index d0d7a3217..8907416b7 100644 --- a/modules/cmd_args.py +++ b/modules/cmd_args.py @@ -39,6 +39,7 @@ def main_args(): group_diag.add_argument("--profile", default=os.environ.get("SD_PROFILE", False), action='store_true', help="Run profiler, default: %(default)s") group_diag.add_argument("--monitor", default=os.environ.get("SD_MONITOR", 0), help="Run memory monitor, default: %(default)s") group_diag.add_argument("--status", default=os.environ.get("SD_STATUS", 120), help="Run server is-alive status, default: %(default)s") + group_diag.add_argument('--experimental', default=os.environ.get("SD_EXPERIMENTAL",False), action='store_true', help="Allow unsupported versions of libraries, default: %(default)s") group_http = parser.add_argument_group('HTTP') group_http.add_argument('--theme', type=str, default=os.environ.get("SD_THEME", None), help='Override UI theme') diff --git a/modules/control/run.py b/modules/control/run.py index c1a93fd6c..712838666 100644 --- a/modules/control/run.py +++ b/modules/control/run.py @@ -133,7 +133,7 @@ def check_active(p, unit_type, units): if u.type != unit_type: continue num_units += 1 - debug_log(f'Control unit: i={num_units} type={u.type} enabled={u.enabled}') + debug_log(f'Control unit: i={num_units} type={u.type} enabled={u.enabled} cn={u.controlnet} proc={u.process}') if not u.enabled: if u.controlnet is not None and u.controlnet.model is not None: debug_log(f'Control unit offload: model="{u.controlnet.model_id}" device={devices.cpu}') @@ -692,7 +692,16 @@ def control_run(state: str = '', p.task_args['image'] = p.init_images # need to set explicitly for txt2img del p.init_images if unit_type == 'lite': - p.init_images = [input_image] + if input_type == 0: + shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.TEXT_2_IMAGE) + shared.sd_model.no_task_switch = True + elif input_type == 1: + p.init_images = [input_image] + elif input_type == 2: + if init_image is None: + shared.log.warning('Control: separate init image not provided') + init_image = input_image + p.init_images = [init_image] instance.apply(selected_models, processed_image, control_conditioning) if hasattr(p, 'init_images') and p.init_images is None: # delete empty del p.init_images diff --git a/modules/control/units/lite.py b/modules/control/units/lite.py index 854e5bd56..6da22c642 100644 --- a/modules/control/units/lite.py +++ b/modules/control/units/lite.py @@ -125,6 +125,7 @@ class ControlLLLite(): class ControlLLitePipeline(): def __init__(self, pipeline: Union[StableDiffusionXLPipeline, StableDiffusionPipeline]): self.pipeline = pipeline + # self.pipeline.__class__.__name__ = 'ControlLLLitePipeline' self.nets = [] def apply(self, controlnet: Union[ControlNetLLLite, list[ControlNetLLLite]], image, conditioning): diff --git a/modules/control/units/t2iadapter.py b/modules/control/units/t2iadapter.py index 2c39a54b5..473563d10 100644 --- a/modules/control/units/t2iadapter.py +++ b/modules/control/units/t2iadapter.py @@ -2,7 +2,7 @@ import os import time from typing import Union import threading -from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, T2IAdapter, MultiAdapter, StableDiffusionAdapterPipeline, StableDiffusionXLAdapterPipeline # pylint: disable=unused-import +from diffusers import pipelines, StableDiffusionPipeline, StableDiffusionXLPipeline, T2IAdapter, MultiAdapter, StableDiffusionAdapterPipeline, StableDiffusionXLAdapterPipeline # pylint: disable=unused-import from installer import log from modules import errors, sd_models from modules.control.units import detect @@ -142,6 +142,14 @@ class AdapterPipeline(): if isinstance(adapter, list) and len(adapter) > 1: adapter = MultiAdapter(adapter) adapter.to(device=pipeline.device, dtype=pipeline.dtype) + """ + pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["sd-t2iadapter"] = StableDiffusionAdapterPipeline + pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["sd-t2iadapter"] = StableDiffusionAdapterPipeline + pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["sd-t2iadapter"] = StableDiffusionAdapterPipeline + pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["sdxl-t2iadapter"] = StableDiffusionXLAdapterPipeline + pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["sdxl-t2iadapter"] = StableDiffusionXLAdapterPipeline + pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["sdxl-t2iadapter"] = StableDiffusionXLAdapterPipeline + """ if pipeline.__class__.__name__ == 'StableDiffusionAdapterPipeline' or pipeline.__class__.__name__ == 'StableDiffusionXLAdapterPipeline': pass # already initialized if detect.is_sdxl(pipeline): diff --git a/modules/control/units/xs_model.py b/modules/control/units/xs_model.py index 460ced12e..f4866b58d 100644 --- a/modules/control/units/xs_model.py +++ b/modules/control/units/xs_model.py @@ -631,6 +631,7 @@ class ControlNetXSModel(ModelMixin, ConfigMixin): t_emb = t_emb.to(dtype=sample.dtype) if self.config.learn_embedding: # pylint: disable=no-member + base_model = base_model.to(self.control_model.device) ctrl_temb = self.control_model.time_embedding(t_emb, timestep_cond) base_temb = base_model.time_embedding(t_emb, timestep_cond) interpolation_param = self.config.time_embedding_mix**0.3 # pylint: disable=no-member diff --git a/modules/control/units/xs_pipe.py b/modules/control/units/xs_pipe.py index 30cd8cef0..f178c11b1 100644 --- a/modules/control/units/xs_pipe.py +++ b/modules/control/units/xs_pipe.py @@ -114,6 +114,14 @@ class StableDiffusionXLControlNetXSPipeline( ): super().__init__() + if isinstance(controlnet, list): + if len(controlnet) == 1: + controlnet = controlnet[0] + else: + raise ValueError( + "ControlNetXS pipeline only supports a single ControlNetXS model" + ) + vae_compatible, cnxs_condition_downsample_factor, vae_downsample_factor = controlnet._check_if_vae_compatible( vae ) diff --git a/modules/flash_attn_triton_amd/bwd_prefill.py b/modules/flash_attn_triton_amd/bwd_prefill.py deleted file mode 100644 index 7f5be379b..000000000 --- a/modules/flash_attn_triton_amd/bwd_prefill.py +++ /dev/null @@ -1,606 +0,0 @@ -import torch -import triton -import triton.language as tl -from modules.flash_attn_triton_amd.utils import get_shape_from_layout, get_strides_from_layout - - -@triton.jit -def _bwd_preprocess_use_o( - Out, - DO, - Delta, - stride_oz, stride_oh, stride_om, stride_ok, - stride_doz, stride_doh, stride_dom, stride_dok, # pylint: disable=unused-argument - stride_deltaz, stride_deltah, stride_deltam, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - BLOCK_M: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - ACTUAL_BLOCK_DMODEL: tl.constexpr, - N_CTX_Q: tl.constexpr, - Z: tl.constexpr, # pylint: disable=unused-argument - H: tl.constexpr, - IS_VARLEN: tl.constexpr -): - pid_m = tl.program_id(0) - pid_bh = tl.program_id(1) - - # Compute batch and head indices - off_z = pid_bh // H - off_h = pid_bh % H - - if IS_VARLEN: - # Compute sequence lengths for the current batch - q_start = tl.load(cu_seqlens_q + off_z) - q_end = tl.load(cu_seqlens_q + off_z + 1) - k_start = tl.load(cu_seqlens_k + off_z) - k_end = tl.load(cu_seqlens_k + off_z + 1) - - # Compute actual sequence lengths - N_CTX_Q = q_end - q_start - N_CTX_K = k_end - k_start # pylint: disable=unused-variable - else: - q_start = 0 - k_start = 0 - N_CTX_Q = max_seqlen_q - N_CTX_K = max_seqlen_k # pylint: disable=unused-variable - - off_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - off_d = tl.arange(0, BLOCK_DMODEL) - - # create masks - mask_m = off_m < N_CTX_Q - mask_d = off_d < ACTUAL_BLOCK_DMODEL - - # compute offsets - o_offset = Out + off_z * stride_oz + off_h * stride_oh + q_start * stride_om - do_offset = DO + off_z * stride_oz + off_h * stride_oh + q_start * stride_om - - # compute pointers - out_ptrs = o_offset + off_m[:, None] * stride_om + off_d[None, :] * stride_ok - do_ptrs = do_offset + off_m[:, None] * stride_dom + off_d[None, :] * stride_dok - - # load - o = tl.load(out_ptrs, mask=mask_m[:, None] & mask_d[None, :], other=0.0).to(tl.float32) - do = tl.load(do_ptrs, mask=mask_m[:, None] & mask_d[None, :], other=0.0).to(tl.float32) - - # compute delta - delta = tl.sum(o * do, axis=1) - - # write-back delta - delta_offset = Delta + off_z * stride_deltaz + off_h * stride_deltah + q_start * stride_deltam - delta_ptrs = delta_offset + off_m * stride_deltam - tl.store(delta_ptrs, delta, mask=mask_m) - - -@triton.jit -def _bwd_kernel_one_col_block( - Q, - K, - V, - sm_scale, - Out, DO, DQ, DK, DV, L, D, # pylint: disable=unused-argument - q_offset, - k_offset, - v_offset, - do_offset, - dq_offset, - dk_offset, - dv_offset, - d_offset, - l_offset, - stride_dq_all, stride_qz, stride_qh, # pylint: disable=unused-argument - stride_qm, - stride_qk, - stride_kz, stride_kh, # pylint: disable=unused-argument - stride_kn, - stride_kk, - stride_vz, stride_vh, # pylint: disable=unused-argument - stride_vn, - stride_vk, - stride_deltaz, stride_deltah, # pylint: disable=unused-argument - stride_deltam, - Z, H, # pylint: disable=unused-argument - N_CTX_Q, - N_CTX_K, - off_h, off_z, off_hz, # pylint: disable=unused-argument - start_n, - num_block_m, - num_block_n, # pylint: disable=unused-argument - BLOCK_M: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - ACTUAL_BLOCK_DMODEL: tl.constexpr, - BLOCK_N: tl.constexpr, - SEQUENCE_PARALLEL: tl.constexpr, - CAUSAL: tl.constexpr, - USE_EXP2: tl.constexpr, -): - if CAUSAL: - # TODO: Causal can skip more blocks with something like lo = start_m * BLOCK_M - lo = 0 - else: - lo = 0 - - # initialize col and head offsets - offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N) - offs_d = tl.arange(0, BLOCK_DMODEL) - - # masks - mask_n = offs_n < N_CTX_K - mask_d = offs_d < ACTUAL_BLOCK_DMODEL - kv_mask = mask_n[:, None] & mask_d[None, :] - - # initialize grad accumulators - dv = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32) - dk = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32) - - # load k and v once per column block - k_ptrs = k_offset + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kk - v_ptrs = v_offset + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vk - k = tl.load(k_ptrs, mask=kv_mask, other=0.0) - v = tl.load(v_ptrs, mask=kv_mask, other=0.0) - - # loop over rows - for start_m in range(lo, num_block_m * BLOCK_M, BLOCK_M): - offs_m = start_m + tl.arange(0, BLOCK_M) - q_ptrs = q_offset + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qk - dq_ptrs = dq_offset + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qk - do_ptrs = do_offset + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qk - - # update mask as row block changes - mask_m = offs_m < N_CTX_Q - q_mask = mask_m[:, None] & mask_d[None, :] - - # load q, k, v, do on-chip - q = tl.load(q_ptrs, mask=q_mask, other=0.0) - do = tl.load(do_ptrs, mask=q_mask, other=0.0) - - # recompute p = softmax(qk, dim=-1).T - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - qk += tl.dot(q, tl.trans(k)) - - if CAUSAL: - col_offset = N_CTX_Q - N_CTX_K - causal_mask = offs_m[:, None] >= (col_offset + offs_n[None, :]) - qk = tl.where(causal_mask, qk, float("-inf")) - - l_ptrs = l_offset + offs_m * stride_deltam - l_i = tl.load(l_ptrs, mask=mask_m) - - # compute p - if USE_EXP2: - RCP_LN2: tl.constexpr = 1.4426950408889634 - qk *= sm_scale * RCP_LN2 - l_i *= RCP_LN2 - p = tl.math.exp2(qk - l_i[:, None]) - else: - qk *= sm_scale - p = tl.math.exp(qk - l_i[:, None]) - - # mask block in the cases where the data is smaller the block size - p_mask = mask_m[:, None] & mask_n[None, :] - p = tl.where(p_mask, p, 0.0) - - # compute dv - dv += tl.dot(tl.trans(p.to(Q.dtype.element_ty)), do) - - # compute dp - dp = tl.dot(do, tl.trans(v)) - - # compute ds , ds = p * (dp - delta[:, None]) - d_ptrs = d_offset + offs_m * stride_deltam - Di = tl.load(d_ptrs, mask=mask_m) - ds = (p * (dp - Di[:, None])) * sm_scale - ds = tl.where(p_mask, ds, 0.0).to(Q.dtype.element_ty) - - # compute dk = dot(ds.T, q) - dk += tl.dot(tl.trans(ds), q) - - # compute dq - if SEQUENCE_PARALLEL: - dq = tl.dot(ds, k) - else: - dq = tl.load(dq_ptrs, mask=q_mask, other=0.0) - dq += tl.dot(ds, k) - tl.store(dq_ptrs, dq.to(Q.dtype.element_ty), mask=q_mask) - - # write-back dv and dk - dk_ptrs = dk_offset + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kk - dv_ptrs = dv_offset + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vk - - # write-back - tl.store(dk_ptrs, dk.to(K.dtype.element_ty), mask=kv_mask) - tl.store(dv_ptrs, dv.to(V.dtype.element_ty), mask=kv_mask) - -@triton.jit -def _bwd_kernel( - Q, - K, - V, - sm_scale, - Out, - DO, - DQ, - DK, - DV, - L, - D, - stride_dq_all, - stride_qz, - stride_qh, - stride_qm, - stride_qk, - stride_kz, - stride_kh, - stride_kn, - stride_kk, - stride_vz, - stride_vh, - stride_vn, - stride_vk, - stride_deltaz, - stride_deltah, - stride_deltam, - Z, - H, - num_block_m, - num_block_n, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - BLOCK_M: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - ACTUAL_BLOCK_DMODEL: tl.constexpr, - BLOCK_N: tl.constexpr, - SEQUENCE_PARALLEL: tl.constexpr, - CAUSAL: tl.constexpr, - USE_EXP2: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - # program ids - off_hz = tl.program_id(0) - if SEQUENCE_PARALLEL: - start_n = tl.program_id(1) - off_z = off_hz // H - off_h = off_hz % H - - if IS_VARLEN: - # Compute sequence lengths for the current batch - q_start = tl.load(cu_seqlens_q + off_z) - q_end = tl.load(cu_seqlens_q + off_z + 1) - k_start = tl.load(cu_seqlens_k + off_z) - k_end = tl.load(cu_seqlens_k + off_z + 1) - - # Compute actual sequence lengths - N_CTX_Q = q_end - q_start - N_CTX_K = k_end - k_start - else: - q_start = 0 - k_start = 0 - N_CTX_Q = max_seqlen_q - N_CTX_K = max_seqlen_k - - # input tensor offsets - q_offset = Q + off_z * stride_qz + off_h * stride_qh + q_start * stride_qm - k_offset = K + off_z * stride_kz + off_h * stride_kh + k_start * stride_kn - v_offset = V + off_z * stride_vz + off_h * stride_vh + k_start * stride_vn - do_offset = DO + off_z * stride_qz + off_h * stride_qh + q_start * stride_qm - l_offset = L + off_z * stride_deltaz + off_h * stride_deltah + q_start * stride_deltam - d_offset = D + off_z * stride_deltaz + off_h * stride_deltah + q_start * stride_deltam - - # output tensor offsets - dk_offset = DK + off_z * stride_kz + off_h * stride_kh + k_start * stride_kn - dv_offset = DV + off_z * stride_vz + off_h * stride_vh + k_start * stride_vn - if SEQUENCE_PARALLEL: - dq_offset = DQ + start_n * stride_dq_all + off_z * stride_qz + off_h * stride_qh + q_start * stride_qm - else: - dq_offset = DQ + off_z * stride_qz + off_h * stride_qh + q_start * stride_qm - - # inner loop - if SEQUENCE_PARALLEL: - _bwd_kernel_one_col_block( - Q, - K, - V, - sm_scale, - Out, - DO, - DQ, - DK, - DV, - L, - D, - q_offset, - k_offset, - v_offset, - do_offset, - dq_offset, - dk_offset, - dv_offset, - d_offset, - l_offset, - stride_dq_all, - stride_qz, - stride_qh, - stride_qm, - stride_qk, - stride_kz, - stride_kh, - stride_kn, - stride_kk, - stride_vz, - stride_vh, - stride_vn, - stride_vk, - stride_deltaz, - stride_deltah, - stride_deltam, - Z, - H, - N_CTX_Q, - N_CTX_K, - off_h, - off_z, - off_hz, - start_n, - num_block_m, - num_block_n, - BLOCK_M=BLOCK_M, - BLOCK_DMODEL=BLOCK_DMODEL, - ACTUAL_BLOCK_DMODEL=ACTUAL_BLOCK_DMODEL, - BLOCK_N=BLOCK_N, - SEQUENCE_PARALLEL=SEQUENCE_PARALLEL, - CAUSAL=CAUSAL, - USE_EXP2=USE_EXP2, - ) - else: - for start_n in range(0, num_block_n): - _bwd_kernel_one_col_block( - Q, - K, - V, - sm_scale, - Out, - DO, - DQ, - DK, - DV, - L, - D, - q_offset, - k_offset, - v_offset, - do_offset, - dq_offset, - dk_offset, - dv_offset, - d_offset, - l_offset, - stride_dq_all, - stride_qz, - stride_qh, - stride_qm, - stride_qk, - stride_kz, - stride_kh, - stride_kn, - stride_kk, - stride_vz, - stride_vh, - stride_vn, - stride_vk, - stride_deltaz, - stride_deltah, - stride_deltam, - Z, - H, - N_CTX_Q, - N_CTX_K, - off_h, - off_z, - off_hz, - start_n, - num_block_m, - num_block_n, - BLOCK_M=BLOCK_M, - BLOCK_DMODEL=BLOCK_DMODEL, - ACTUAL_BLOCK_DMODEL=ACTUAL_BLOCK_DMODEL, - BLOCK_N=BLOCK_N, - SEQUENCE_PARALLEL=SEQUENCE_PARALLEL, - CAUSAL=CAUSAL, - USE_EXP2=USE_EXP2, - ) - - -# NOTE: smaller blocks have lower accuracy. more accumlation error probably 128 * 128 seems good but leads to oom. 64 * 64 has accumlation errors but no oom. -def attention_prefill_backward_triton_impl( - do, - q, - k, - v, - o, - softmax_lse, - dq, - dk, - dv, - sm_scale: float, - alibi_slopes, # pylint: disable=unused-argument - causal, - layout: str, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q: int, - max_seqlen_k: int, - use_exp2: bool, - sequence_parallel = True, -): - # make contigious - q = q.contiguous() - k = k.contiguous() - v = v.contiguous() - softmax_lse = softmax_lse.contiguous() - - # get strides and shape - batch, nheads_q, nheads_k, head_size, max_seqlen_q, max_seqlen_k = get_shape_from_layout(q, k, layout, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k) # pylint: disable=unused-variable - q_strides, k_strides, v_strides, o_strides = get_strides_from_layout(q, k, v, o, layout) - stride_qz, stride_qh, stride_qm, stride_qk = q_strides - stride_kz, stride_kh, stride_kn, stride_kk = k_strides - stride_vz, stride_vh, stride_vn, stride_vk = v_strides - stride_oz, stride_oh, stride_om, stride_ok = o_strides - batch_headsize = batch * nheads_q - is_varlen = layout == "thd" - - # FIXME: some configs lead to oom for some reason when using 64 x 64 blocks - if max_seqlen_q <= 32 or max_seqlen_k <= 32: - BLOCK_M = 32 - BLOCK_N = 32 - else: - BLOCK_M = 64 - BLOCK_N = 64 - num_warps = 4 # NOTE: originial is 8. changing it to 1 caused issues be careful - num_stages = 1 - waves_per_eu = 1 - - # divide up the problem - num_blocks_m = triton.cdiv(max_seqlen_q, BLOCK_M) - num_blocks_n = triton.cdiv(max_seqlen_k, BLOCK_N) - - # get closest power of 2 over or equal to 32. - padded_d_model = 1 << (head_size - 1).bit_length() - padded_d_model = max(padded_d_model, 16) - BLOCK_DMODEL = padded_d_model - ACTUAL_BLOCK_DMODEL = head_size - - do = do.contiguous() - # NOTE: we might need to copy the output tensor if they are not continuous or have other issues - copy_back = {"dq": False, "dk": False, "dv": False} - - dq_og = None - # deal with dq - if dq is None: - if sequence_parallel: - dq = torch.zeros((num_blocks_n,) + q.shape, device=q.device, dtype=q.dtype) - else: - dq = torch.zeros(q.shape, device=q.device, dtype=q.dtype) - else: - dq_og = dq - if not dq.is_contiguous(): - dq = dq.contiguous() - copy_back["dq"] = True - - if sequence_parallel: - dq = torch.zeros((num_blocks_n,) + q.shape, device=q.device, dtype=q.dtype) - copy_back["dq"] = True - else: - # NOTE: the kernel does inplace accumlation so dq has to be zeros. This avoids the case where we are passed empty dq and it is not all zeros - dq.zero_() - stride_dq_all = dq.stride()[0] - - dk_og = None - dv_og = None - # deal with dk, dv - if (dk is None) or (dv is None): - dk = torch.empty_like(k) - dv = torch.empty_like(v) - else: - if not dk.is_contiguous(): - dk_og = dk - dk = dk.contiguous() - copy_back["dk"] = True - - if not dv.is_contiguous(): - dv_og = dv - dv = dv.contiguous() - copy_back["dv"] = True - - # assert contigious - assert do.is_contiguous() - assert q.is_contiguous() - assert k.is_contiguous() - assert v.is_contiguous() - assert o.is_contiguous() - assert softmax_lse.is_contiguous() - - # init delta - delta = torch.empty_like(softmax_lse) - if is_varlen: - stride_deltam, stride_deltah = delta.stride() - stride_deltaz = 0 - else: - stride_deltaz, stride_deltah, stride_deltam = delta.stride() - - _bwd_preprocess_use_o[(num_blocks_m, batch_headsize)]( - o, - do, - delta, - stride_oz, stride_oh, stride_om, stride_ok, - stride_oz, stride_oh, stride_om, stride_ok, - stride_deltaz, stride_deltah, stride_deltam, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - BLOCK_M=BLOCK_M, - BLOCK_DMODEL=BLOCK_DMODEL, - ACTUAL_BLOCK_DMODEL=ACTUAL_BLOCK_DMODEL, - N_CTX_Q=max_seqlen_q, - Z=batch, - H=nheads_q, - IS_VARLEN=is_varlen - ) - - _bwd_kernel[(batch_headsize, num_blocks_n if sequence_parallel else 1)]( - q, - k, - v, - sm_scale, - o, - do, - dq, - dk, - dv, - softmax_lse, - delta, - stride_dq_all, - stride_qz, stride_qh, stride_qm, stride_qk, - stride_kz, stride_kh, stride_kn, stride_kk, - stride_vz, stride_vh, stride_vn, stride_vk, - stride_deltaz, stride_deltah, stride_deltam, - batch, - nheads_q, - num_blocks_m, - num_blocks_n, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - BLOCK_M=BLOCK_M, - BLOCK_N=BLOCK_N, - BLOCK_DMODEL=BLOCK_DMODEL, - ACTUAL_BLOCK_DMODEL=ACTUAL_BLOCK_DMODEL, - SEQUENCE_PARALLEL=sequence_parallel, - CAUSAL=causal, - USE_EXP2=use_exp2, - num_warps=num_warps, - num_stages=num_stages, - waves_per_eu = waves_per_eu, - IS_VARLEN=is_varlen - ) - - if sequence_parallel: - dq = dq.sum(dim=0) - - if copy_back["dq"]: - dq_og.copy_(dq) - dq = dq_og - if copy_back["dk"]: - dk_og.copy_(dk) - dk = dk_og - if copy_back["dv"]: - dv_og.copy_(dv) - dv = dv_og - - return dq, dk, dv, delta, None, None diff --git a/modules/flash_attn_triton_amd/fwd_decode.py b/modules/flash_attn_triton_amd/fwd_decode.py deleted file mode 100644 index 7a2a234d6..000000000 --- a/modules/flash_attn_triton_amd/fwd_decode.py +++ /dev/null @@ -1,700 +0,0 @@ -import torch -import triton -import triton.language as tl -from modules.flash_attn_triton_amd.utils import _strides, get_padded_headsize - - -@triton.jit -def _fwd_kernel_splitK( - Q, - K, - V, - sm_scale, - Out_splitK, # [B, H, split_k, Mq, K] - Metadata, # [B, H, 2, split_k, M_ceil] contains [mi, li] - K_new, - V_new, - Cache_seqlens, - Cache_batch_idx, - Alibi_slopes, - stride_qz, - stride_qm, - stride_qg, - stride_qh, - stride_qd, - stride_kz, - stride_kn, - stride_kg, - stride_kh, - stride_kd, - stride_vz, - stride_vn, - stride_vg, - stride_vh, - stride_vd, - stride_osk_zhg, - stride_osk_s, - stride_osk_m, - stride_osk_d, # pylint: disable=unused-argument - stride_mzhg, - stride_m2, - stride_ms, - stride_mm, # pylint: disable=unused-argument - stride_kn_z, - stride_kn_n, - stride_kn_g, - stride_kn_h, - stride_kn_d, - stride_vn_z, - stride_vn_n, - stride_vn_g, - stride_vn_h, - stride_vn_d, - stride_az, - stride_ah, - Z, # pylint: disable=unused-argument - N_CTX_Q, - N_CTX_K, - N_CTX_NEW, - BLOCK_N_PER_SPLIT, - H_q: tl.constexpr, - H_kv: tl.constexpr, - G_q: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_DMODEL: tl.constexpr, - ACTUAL_BLOCK_DMODEL: tl.constexpr, - BLOCK_N: tl.constexpr, - BOUNDS_CHECKS_N: tl.constexpr, - USE_CACHE_SEQLENs: tl.constexpr, - USE_CACHE_BATCH_IDX: tl.constexpr, - NEW_KV: tl.constexpr, - IS_GQA: tl.constexpr, - IS_CAUSAL: tl.constexpr, - USE_ALIBI: tl.constexpr, -): - # Padding - PADDED_HEAD: tl.constexpr = ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL - if PADDED_HEAD: - d_mask = tl.arange(0, BLOCK_DMODEL) < ACTUAL_BLOCK_DMODEL - - start_m = tl.program_id(0) - off_zhg = tl.program_id(1) - off_z = off_zhg // (H_q * G_q) - off_h_q = (off_zhg // G_q) % H_q - off_g_q = off_zhg % G_q - splitk_idx = tl.program_id(2) - - # pick batch index - if USE_CACHE_BATCH_IDX: - cache_batch_idx = tl.load(Cache_batch_idx + off_z) - else: - cache_batch_idx = off_z - - # Load ALiBi slope if enabled - if USE_ALIBI: - a_offset = off_z * stride_az + off_h_q * stride_ah - alibi_slope = tl.load(Alibi_slopes + a_offset) - else: - alibi_slope = None - - lo = splitk_idx * BLOCK_N_PER_SPLIT - if USE_CACHE_SEQLENs: - cache_seqlen_last_idx = tl.load(Cache_seqlens + off_z) - if NEW_KV: - kv_len = cache_seqlen_last_idx + N_CTX_NEW - else: - kv_len = cache_seqlen_last_idx - else: - kv_len = N_CTX_K - hi = tl.minimum((splitk_idx + 1) * BLOCK_N_PER_SPLIT, kv_len) - - HEAD_RATIO: tl.constexpr = H_q // H_kv - if IS_GQA: - k_head_idx = off_h_q // HEAD_RATIO - v_head_idx = k_head_idx - else: - k_head_idx = off_h_q - v_head_idx = off_h_q - - # calculate base offset - k_base = K + k_head_idx * stride_kh + cache_batch_idx * stride_kz + off_g_q * stride_kg - v_base = V + v_head_idx * stride_vh + cache_batch_idx * stride_vz + off_g_q * stride_vg - - # Copy new Keys and Values into Cache - if NEW_KV: - knew_base = K_new + k_head_idx * stride_kn_h + off_z * stride_kn_z + off_g_q * stride_kn_g - - # Determine the starting position for new data in the cache - if USE_CACHE_SEQLENs: - start_idx = tl.load(Cache_seqlens + off_z) - else: - start_idx = N_CTX_K - N_CTX_NEW - - # Copy new Keys - for i in range(0, N_CTX_NEW, BLOCK_N): - # Load from K_new - k_new_block = tl.load( - knew_base + - tl.arange(0, BLOCK_DMODEL)[:, None] * stride_kn_d + - (tl.arange(0, BLOCK_N) + i)[None, :] * stride_kn_n, - mask=(tl.arange(0, BLOCK_N)[None, :] + i < N_CTX_NEW) & - (tl.arange(0, BLOCK_DMODEL)[:, None] < ACTUAL_BLOCK_DMODEL), - other=0 - ) - - # Store to K - tl.store( - k_base + - tl.arange(0, BLOCK_DMODEL)[:, None] * stride_kd + - (tl.arange(0, BLOCK_N) + i + start_idx)[None, :] * stride_kn, - k_new_block, - mask=(tl.arange(0, BLOCK_N)[None, :] + i < N_CTX_NEW) & - (tl.arange(0, BLOCK_DMODEL)[:, None] < ACTUAL_BLOCK_DMODEL), - ) - - # Copy new Values - vnew_base = V_new + v_head_idx * stride_vn_h + off_z * stride_vn_z + off_g_q * stride_vn_g - for i in range(0, N_CTX_NEW, BLOCK_N): - # Load from V_new - v_new_block = tl.load( - vnew_base + - (tl.arange(0, BLOCK_N) + i)[:, None] * stride_vn_n + - tl.arange(0, BLOCK_DMODEL)[None, :] * stride_vn_d, - mask=(tl.arange(0, BLOCK_N)[:, None] + i < N_CTX_NEW) & - (tl.arange(0, BLOCK_DMODEL)[None, :] < ACTUAL_BLOCK_DMODEL), - other=0 - ) - - # Store to V - tl.store( - v_base + - (tl.arange(0, BLOCK_N) + i + start_idx)[:, None] * stride_vn + - tl.arange(0, BLOCK_DMODEL)[None, :] * stride_vd, - v_new_block, - mask=(tl.arange(0, BLOCK_N)[:, None] + i < N_CTX_NEW) & - (tl.arange(0, BLOCK_DMODEL)[None, :] < ACTUAL_BLOCK_DMODEL), - ) - - Q_block_ptr = tl.make_block_ptr( - base=Q + off_h_q * stride_qh + off_z * stride_qz + off_g_q * stride_qg, - shape=(N_CTX_Q, ACTUAL_BLOCK_DMODEL), - strides=(stride_qm, stride_qd), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, BLOCK_DMODEL), - order=(1, 0), - ) - - K_block_ptr = tl.make_block_ptr( - base=k_base, - shape=(ACTUAL_BLOCK_DMODEL, hi), - strides=(stride_kd, stride_kn), - offsets=(0, lo), - block_shape=(BLOCK_DMODEL, BLOCK_N), - order=(0, 1), - ) - V_block_ptr = tl.make_block_ptr( - base=v_base, - shape=(hi, ACTUAL_BLOCK_DMODEL), - strides=(stride_vn, stride_vd), - offsets=(lo, 0), - block_shape=(BLOCK_N, BLOCK_DMODEL), - order=(1, 0), - ) - - K_scale_shift_block_ptr = None - V_scale_shift_block_ptr = None - - # initialize pointer to m and l - m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) - l_i = tl.zeros([BLOCK_M], dtype=tl.float32) - - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) # noqa: F821 - - # scale sm_scale by log_2(e) and use - # 2^x instead of exp in the loop because CSE and LICM - # don't work as expected with `exp` in the loop - qk_scale = sm_scale * 1.44269504 - # load q: it will stay in SRAM throughout - q = tl.load( # noqa: F821 - tl.advance(Q_block_ptr, (0, 0)), boundary_check=(0, )) - q = (q * qk_scale).to(q.dtype) - if PADDED_HEAD: - q = tl.where(d_mask[None, :], q, 0.0) - - # loop over k, v and update accumulator - for start_n in range(lo, hi, BLOCK_N): - k, v = load_k_v_group( - K_block_ptr, - V_block_ptr, - K_scale_shift_block_ptr, - V_scale_shift_block_ptr, - BOUNDS_CHECKS_N, - 1, - BLOCK_DMODEL, - ACTUAL_BLOCK_DMODEL, - Q.dtype.element_ty, - 0, - ) - if PADDED_HEAD: - k = tl.where(d_mask[:, None], k, 0.0) - v = tl.where(d_mask[None, :], v, 0.0) - - # -- compute qk --- - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) - qk += tl.dot(q, k) # noqa: F821 - - if USE_ALIBI: - row_idx = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - col_idx = start_n + tl.arange(0, BLOCK_N) - - # Compute relative positions - relative_pos = row_idx[:, None] + kv_len - (N_CTX_Q + col_idx[None, :]) - relative_pos = tl.abs(relative_pos) - - # Compute ALiBi bias - alibi_bias = -1 * alibi_slope * relative_pos - qk += (alibi_bias * 1.44269504) - - # Apply causal mask if IS_CAUSAL is True - if IS_CAUSAL: - row_idx = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - col_idx = start_n + tl.arange(0, BLOCK_N) - - # create a N_CTX_Q x kv_len causal mask - col_offset = N_CTX_Q - kv_len - causal_mask = row_idx[:, None] >= (col_offset + col_idx[None, :]) - - # Apply the mask - qk = tl.where(causal_mask, qk, float("-inf")) - - # TODO: This is slow, and only needed at the last iteration. - # Maybe we can unroll the last iteration instead? - if BOUNDS_CHECKS_N: - qk = tl.where(tl.arange(0, BLOCK_N) < hi - start_n, qk, float("-inf")) - - # -- compute scaling constant --- - m_i_new = tl.maximum(m_i, tl.max(qk, 1)) - if IS_CAUSAL: - alpha = tl.math.exp2(tl.where(m_i > float("-inf"), m_i - m_i_new, float("-inf"))) - else: - alpha = tl.math.exp2(m_i - m_i_new) - # cause of nan because subtracting infs - if IS_CAUSAL: - qk = tl.where(qk > float("-inf"), qk - m_i_new[:, None], float("-inf")) - else: - qk = qk - m_i_new[:, None] - - p = tl.math.exp2(qk) - - # -- update m_i and l_i -- - l_i = l_i * alpha + tl.sum(p, 1) - m_i = m_i_new - p = p.to(Q.dtype.element_ty) - - # -- scale and update acc -- - acc *= alpha[:, None] - acc += tl.dot(p.to(v.dtype), v) - - # update pointers - K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N)) - V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0)) - - # write back O - O_block_ptr = tl.make_block_ptr( - base=Out_splitK + off_zhg * stride_osk_zhg + splitk_idx * stride_osk_s, - shape=(N_CTX_Q, BLOCK_DMODEL), - strides=(stride_osk_m, 1), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, BLOCK_DMODEL), - order=(1, 0), - ) - tl.store( - tl.advance(O_block_ptr, (0, 0)), - acc, - boundary_check=(0, ), - ) - # Write metadata for split-K reduction - Metadata_ptr = (Metadata + off_zhg * stride_mzhg + splitk_idx * stride_ms + start_m * BLOCK_M + - tl.arange(0, BLOCK_M)) - tl.store(Metadata_ptr, m_i) - tl.store(Metadata_ptr + stride_m2, l_i) - - -@triton.jit -def load_k_v_group( - K_block_ptr, - V_block_ptr, - K_scale_shift_block_ptr, V_scale_shift_block_ptr, # pylint: disable=unused-argument - BOUNDS_CHECKS_N: tl.constexpr, - PACKED_PER_VAL: tl.constexpr, BLOCK_DMODEL: tl.constexpr, # pylint: disable=unused-argument - ACTUAL_BLOCK_DMODEL: tl.constexpr, - dtype: tl.constexpr, # pylint: disable=unused-argument - group_id: tl.constexpr, -): - # Load K/V for a given block - # Advance to the current quantization group - K_block_ptr = tl.advance(K_block_ptr, (ACTUAL_BLOCK_DMODEL * group_id, 0)) - V_block_ptr = tl.advance(V_block_ptr, (0, ACTUAL_BLOCK_DMODEL * group_id)) - - # -- load k, v -- - k = tl.load(K_block_ptr, boundary_check=(1, ) if BOUNDS_CHECKS_N else ()) - v = tl.load(V_block_ptr, boundary_check=(0, ) if BOUNDS_CHECKS_N else ()) - - return k, v - - -@triton.jit -def cast_uint32_to_half2(scale_shift): - # Extract two float16 packed into one int32 - scale = scale_shift & 0xFFFF - shift = scale_shift >> 16 - scale = scale.to(tl.uint16).to(tl.float16, bitcast=True) - shift = shift.to(tl.uint16).to(tl.float16, bitcast=True) - return scale, shift - - -@triton.jit -def dequantize( - x_, - scale, - shift, - PACKED_PER_VAL: tl.constexpr = 8, -): - # PACKED_PER_VAL is the number of values packed into - # each element x_. For example, for int4 quantization - #and x_ of type int32, PACKED_PER_VAL is 8. - - BLOCK_N: tl.constexpr = x_.shape[0] - BLOCK_DMODEL_PACKED: tl.constexpr = x_.shape[1] - offsets = tl.arange(0, PACKED_PER_VAL) * 4 - quant_offset = (x_[:, None, :] >> offsets[None, :, None]) # (BLOCK_N, PACKED_PER_VAL, D // PACKED_PER_VAL) - - quant_offset = tl.view(quant_offset, (BLOCK_N, BLOCK_DMODEL_PACKED * PACKED_PER_VAL)) - # Trick - instead of converting int4 to float16 we view it as float16 - # and then multiply by 32768 * 512 == 2**24 - quant_offset = (quant_offset & 0xF).to(tl.uint16).to(tl.float16, bitcast=True) - quant_offset = (quant_offset * 32768.0).to(tl.float16) - scale_512 = scale * 512 - - dequant = quant_offset * scale_512 + shift - return dequant - - -@triton.jit -def _splitK_reduce( - Out_splitK, # [B, H, split_k, Mq, K] - Metadata, # [B, H, 2, split_k, M_ceil] contains [mi, li] - Out, # [B, H, M, K] - LSE, # [B, H, M] - stride_osk_zhg, - stride_osk_s, - stride_osk_m, - stride_osk_k, - stride_mzhg, - stride_m2, - stride_ms, - stride_mm, - stride_oz, - stride_oh, - stride_og, - stride_om, - stride_ok, # pylint: disable=unused-argument - stride_lse_zhg, - stride_lse_m, M_ceil: tl.constexpr, # pylint: disable=unused-argument - BLOCK_SIZE: tl.constexpr, - H: tl.constexpr, - G: tl.constexpr, - split_k: tl.constexpr, - splitK_pow2: tl.constexpr, - use_mask: tl.constexpr, - IS_CAUSAL: tl.constexpr, -): - off_zhg = tl.program_id(0) - off_z = off_zhg // (H * G) - off_h = (off_zhg // G) % H - off_g = off_zhg % G - off_m = tl.program_id(1) - off_k = tl.program_id(2) - - # read chunk - spk_idx = tl.arange(0, splitK_pow2) - kidx = tl.arange(0, BLOCK_SIZE) - - Metadata_ptr = Metadata + stride_mzhg * off_zhg + spk_idx * stride_ms + off_m * stride_mm - - o_ptr = (Out_splitK + off_zhg * stride_osk_zhg + stride_osk_m * off_m + off_k * BLOCK_SIZE + - stride_osk_s * spk_idx[:, None] + kidx[None, :] * stride_osk_k) - - # read max values of each splitK - if use_mask: - spk_mask = spk_idx < split_k - l_m = tl.load(Metadata_ptr, mask=spk_mask, other=float("-inf")) - l_sum = tl.load(Metadata_ptr + stride_m2, mask=spk_mask, other=0.0) - acc = tl.load(o_ptr, mask=spk_mask[:, None], other=0.0) - else: - l_m = tl.load(Metadata_ptr) - l_sum = tl.load(Metadata_ptr + stride_m2) - acc = tl.load(o_ptr) - - g_m = tl.max(l_m, axis=0) - - if IS_CAUSAL: - l_m_offset = l_m - g_m - alpha = tl.where(l_m_offset > float("-inf"), tl.math.exp2(l_m_offset), 0.0) - else: - alpha = tl.math.exp2(l_m - g_m) - - # read sum - l_sum *= alpha - g_sum = tl.sum(l_sum, axis=0) - acc = acc * alpha[:, None] - - if IS_CAUSAL: - # Avoid division by zero - g_sum_safe = tl.where(g_sum > 0, g_sum, 1.0) - acc_out = tl.sum(acc, axis=0) / g_sum_safe - else: - acc_out = tl.sum(acc, axis=0) / g_sum - - # Store output - Out_ptr = (Out + stride_oz * off_z + stride_oh * off_h + stride_og * off_g + stride_om * off_m + - off_k * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)) - tl.store(Out_ptr, acc_out) - - # Store lse - l_ptrs = LSE + off_zhg * stride_lse_zhg + off_m - if IS_CAUSAL: - lse = tl.where(g_sum > 0, (g_m + tl.math.log2(g_sum)) / 1.44269504, g_m) - tl.store(l_ptrs, lse) - else: - tl.store(l_ptrs, (g_m + tl.math.log2(g_sum)) / 1.44269504) - - -def quantize_kv_int4(k: torch.Tensor, num_groups: int = 1) -> torch.Tensor: - # Scale and shift are such that quantization linearly maps - # int4 values range [0..15] to input values range min(k)..max(k) - # individually for every row - k = k.reshape(*k.shape[:-1], num_groups, k.shape[-1] // num_groups) - max_vals = torch.max(k, dim=-1, keepdim=True).values - min_vals = torch.min(k, dim=-1, keepdim=True).values - scale_k: torch.Tensor = (max_vals - min_vals) / 15 - - shift_k = torch.min(k, dim=-1, keepdim=True).values - scale_k = scale_k.to(torch.float16) - shift_k = shift_k.to(torch.float16) - - in_bytes = ((k - shift_k.expand(k.shape)) / scale_k.expand(k.shape)) + 0.5 - in_bytes = in_bytes.to(torch.uint8) - in_int4 = in_bytes & 0xF - in_int4_packed = in_int4[..., ::2] + (in_int4[..., 1::2] << 4) - scale_shift = torch.concat([scale_k.view(torch.uint8), shift_k.view(torch.uint8)], dim=-1) - k_quant = torch.concat( - [ - scale_shift.flatten(start_dim=-2), - in_int4_packed.flatten(start_dim=-2), - ], - dim=-1, - ).view(torch.int16) - return k_quant - - -def dequantize_kv_fp16(quant_k: torch.Tensor, num_groups: int = 1) -> torch.Tensor: - k_i16 = quant_k.view(torch.int16) - k_ui8 = k_i16.view(torch.uint8) - - ss_size = num_groups * 4 - scale_shift_ui8 = k_ui8[..., 0:ss_size] - scale_shift_ui8 = scale_shift_ui8.reshape(*scale_shift_ui8.shape[:-1], num_groups, 4) - scale = scale_shift_ui8[..., 0:2].view(torch.float16) - shift = scale_shift_ui8[..., 2:4].view(torch.float16) - - kv_ui8 = k_ui8[..., ss_size:] - k_ui8 = kv_ui8.reshape(*kv_ui8.shape[:-1], num_groups, -1) - k1_i4 = k_ui8 & 0xF - k2_i4 = (k_ui8 & 0xF0) >> 4 - k_shape = k1_i4.shape - k1_f16 = k1_i4.to(torch.float16) * scale.expand(k_shape) + shift.expand(k_shape) - k2_f16 = k2_i4.to(torch.float16) * scale.expand(k_shape) + shift.expand(k_shape) - - out = torch.empty((*k1_f16.shape[:-1], k1_f16.shape[-1] * 2), dtype=torch.float16, device=quant_k.device) - out[..., ::2] = k1_f16 - out[..., 1::2] = k2_f16 - out = out.reshape(*k_shape[:-2], -1) - - return out - - -def get_split_k(B: int, G: int, H: int, Mk: int) -> int: - """Heuristic for the number of splits""" - bh = max(B * H, 1) # NOTE: Handle B*h=0 case - split_k = max(Mk, 1024) // bh - max_chunk_size = 64 - while split_k > 0 and Mk / split_k < max_chunk_size: - split_k = split_k // 2 - while B * H * G * split_k >= 1024: - split_k = split_k // 2 - split_k = min(split_k, 512) - split_k = max(split_k, 1) - return split_k - -def attention_decode_forward_triton_impl(q, k, v, sm_scale, causal, alibi_slopes, layout, cache_seqlens, cache_batch_idx, new_kv, k_new, v_new): - # kernel config - BLOCK_M = 16 - BLOCK_N = 64 - SPLIT_K = None - NUM_QUANT_GROUPS = 1 # pylint: disable=unused-variable - - # kernels expects "bsghd" - original_layout = layout - if layout == "bshd": - q = q.unsqueeze(2) - k = k.unsqueeze(2) - v = v.unsqueeze(2) - if new_kv: - k_new = k_new.unsqueeze(2) - v_new = v_new.unsqueeze(2) - layout = "bsghd" - elif layout == "bhsd": - q = q.permute(0, 2, 1, 3).unsqueeze(2) - k = k.permute(0, 2, 1, 3).unsqueeze(2) - v = v.permute(0, 2, 1, 3).unsqueeze(2) - if new_kv: - k_new = k_new.permute(0, 2, 1, 3).unsqueeze(2) - v_new = v_new.permute(0, 2, 1, 3).unsqueeze(2) - layout = "bsghd" - elif layout == "bsghd": - pass - elif layout is None: - raise ValueError("Layout not given") - assert layout == "bsghd" - - # get dims - batch_size, seqlen_q, n_group_q, heads_per_group_q, dim_q = q.shape - _, seqlen_k, n_group_k, heads_per_group_k, dim_k = k.shape # pylint: disable=unused-variable - _, seqlen_v, n_group_v, heads_per_group_v, dim_v = v.shape # pylint: disable=unused-variable - - assert dim_q == dim_k == dim_v, f"Dimensions must match: {dim_q}, {dim_k}, {dim_v}" - - # get padded size - dim_padded = get_padded_headsize(dim_k) - - # Handle MQA/GQA case - if heads_per_group_q > heads_per_group_k: - is_gqa = True - elif heads_per_group_q < heads_per_group_k: - raise ValueError("heads_per_group_q < heads_per_group_k") - else: - is_gqa = False - - assert dim_k == dim_q, f"Keys have head dim {dim_k} but queries have head dim {dim_q}" - - if SPLIT_K is not None: - split_k = SPLIT_K - else: - # Use heuristics - split_k = get_split_k(batch_size, n_group_q, heads_per_group_q, seqlen_k) # NOTE: should the split think about seqlens? - - seqlen_q_ceil = (seqlen_q + BLOCK_M - 1) // BLOCK_M * BLOCK_M - out_splitk = torch.empty([batch_size * n_group_q * heads_per_group_q, split_k, seqlen_q_ceil, dim_padded], dtype=torch.float32, device=q.device) - metadata = torch.empty([batch_size * n_group_q * heads_per_group_q, 2, split_k, seqlen_q_ceil], dtype=torch.float32, device=q.device) - lse = torch.empty((batch_size * n_group_q * heads_per_group_q, seqlen_q), device=q.device, dtype=torch.float32) - grid = (triton.cdiv(seqlen_q, BLOCK_M), batch_size * n_group_q * heads_per_group_q, split_k) - - num_warps = 1 - split_size = (seqlen_k + split_k - 1) // split_k - use_cache_seqlens = cache_seqlens is not None - - # TODO: enable quantization - _fwd_kernel_splitK[grid]( - Q=q, - K=k, - V=v, - sm_scale=sm_scale, - Out_splitK=out_splitk, - Metadata=metadata, - K_new = k_new, - V_new = v_new, - Cache_seqlens=cache_seqlens, - Cache_batch_idx=cache_batch_idx, - Alibi_slopes=alibi_slopes, - **_strides(q, "qz", "qm", "qg", "qh", "qd"), - **_strides(k, "kz", "kn", "kg", "kh", "kd"), - **_strides(v, "vz", "vn", "vg", "vh", "vd"), - **_strides(out_splitk, "osk_zhg", "osk_s", "osk_m", "osk_d"), - **_strides(metadata, "mzhg", "m2", "ms", "mm"), - **_strides(k_new, "kn_z", "kn_n", "kn_g", "kn_h", "kn_d"), - **_strides(v_new, "vn_z", "vn_n", "vn_g", "vn_h", "vn_d"), - **_strides(alibi_slopes, "az", "ah"), - Z=batch_size, - H_q=heads_per_group_q, - H_kv=heads_per_group_k, - G_q=n_group_q, - N_CTX_Q=seqlen_q, - N_CTX_K=seqlen_k, - N_CTX_NEW=k_new.shape[1] if new_kv else None, - BLOCK_N_PER_SPLIT=split_size, - BLOCK_M=BLOCK_M, - BLOCK_N=BLOCK_N, - BLOCK_DMODEL=dim_padded, - ACTUAL_BLOCK_DMODEL=dim_k, - BOUNDS_CHECKS_N=(split_size % BLOCK_N) > 0 or use_cache_seqlens, - USE_CACHE_SEQLENs=use_cache_seqlens, - USE_CACHE_BATCH_IDX=cache_batch_idx is not None, - NEW_KV=new_kv, - IS_GQA=is_gqa, - IS_CAUSAL=causal, - USE_ALIBI=False if alibi_slopes is None else True, - num_warps=num_warps, - num_stages=1, - ) - - out = torch.empty((batch_size, seqlen_q, n_group_q, heads_per_group_q, dim_padded), device=q.device, dtype=q.dtype) - - # Merge together - splitK_pow2 = triton.next_power_of_2(split_k) - use_mask = splitK_pow2 > split_k - if batch_size * n_group_q * heads_per_group_q * seqlen_q >= 512: - k_block_num = 1 - else: - k_block_num = 2 - assert dim_padded % k_block_num == 0 - k_block_size = dim_padded // k_block_num - grid = (batch_size * n_group_q * heads_per_group_q, seqlen_q, k_block_num) - - _splitK_reduce[grid]( - out_splitk, - metadata, - out, - lse, - **_strides(out_splitk, "osk_zhg", "osk_s", "osk_m", "osk_k"), - **_strides(metadata, "mzhg", "m2", "ms", "mm"), - **_strides(out, "oz", "om", "og", "oh", "ok"), - **_strides(lse, "lse_zhg", "lse_m"), - M_ceil=seqlen_q_ceil, - BLOCK_SIZE=k_block_size, - G=n_group_q, - H=heads_per_group_q, - # TODO: Tune num_warps - split_k=split_k, - splitK_pow2=splitK_pow2, - use_mask=use_mask, - IS_CAUSAL=causal, - num_warps=4) - - lse = lse.reshape([batch_size, n_group_q, heads_per_group_q, seqlen_q]) - if q.ndim == 4: - # BMGHK -> BMHK - assert n_group_q == 1 - out = out[:, :, 0] - lse = lse[:, 0] - if seqlen_k == 0: - out.zero_() - out = out.reshape(batch_size, heads_per_group_q * n_group_q, -1, dim_padded).contiguous() - - # output is batch_size, heads_per_group_q * group_q, seqlen_q, dim_q - if original_layout == "bshd": - # out=out.transpose(1, 2).contiguous() # this screws up heads and data. - # the data is laid out properly. Just need to reshape dims - out = out.reshape(batch_size, seqlen_q, -1, dim_padded) - - return out.narrow(-1, 0, dim_k), lse diff --git a/modules/flash_attn_triton_amd/fwd_prefill.py b/modules/flash_attn_triton_amd/fwd_prefill.py index ca9fdc507..cf78cd0ee 100644 --- a/modules/flash_attn_triton_amd/fwd_prefill.py +++ b/modules/flash_attn_triton_amd/fwd_prefill.py @@ -1,33 +1,8 @@ +from typing import Literal, Optional, Union import torch import triton import triton.language as tl -from modules.flash_attn_triton_amd.utils import get_shape_from_layout, get_strides_from_layout, is_cdna, is_rdna, AUTOTUNE - - -@triton.jit -def cdiv_fn(x, y): - return (x + y - 1) // y - - -@triton.jit -def dropout_offsets(philox_seed, philox_offset, dropout_p, m, n, stride): # pylint: disable=unused-argument - ms = tl.arange(0, m) - ns = tl.arange(0, n) - return philox_offset + ms[:, None] * stride + ns[None, :] - - -@triton.jit -def dropout_rng(philox_seed, philox_offset, dropout_p, m, n, stride): - rng_offsets = dropout_offsets(philox_seed, philox_offset, dropout_p, m, n, stride).to(tl.uint32) - # TODO: use tl.randint for better performance - return tl.rand(philox_seed, rng_offsets) - - -@triton.jit -def dropout_mask(philox_seed, philox_offset, dropout_p, m, n, stride): - rng_output = dropout_rng(philox_seed, philox_offset, dropout_p, m, n, stride) - rng_keep = rng_output > dropout_p - return rng_keep +from modules.flash_attn_triton_amd.utils import AUTOTUNE, compute_alibi_block, get_shapes_from_layout, get_strides_from_layout, is_cdna, is_rdna # Convenience function to load with optional boundary checks. @@ -50,47 +25,14 @@ def load_fn(ptrs, offset_first, offset_second, boundary_first, boundary_second): @triton.jit -def compute_alibi_block(alibi_slope, seqlen_q, seqlen_k, offs_m, offs_n, transpose=False): - # when seqlen_k and seqlen_q are different we want the diagonal to stick to the bottom right of the attention matrix - # for casual mask we want something like this where (1 is kept and 0 is masked) - # seqlen_q = 2 and seqlen_k = 5 - # 1 1 1 1 0 - # 1 1 1 1 1 - # seqlen_q = 5 and seqlen_k = 2 - # 0 0 - # 0 0 - # 0 0 - # 1 0 - # 1 1 - # for alibi the diagonal is 0 indicating no penalty for attending to that spot and increasing penalty for attending further from the diagonal - # e.g. alibi_slope = 1, seqlen_q = 2, seqlen_k = 5, offs_m = [0, 1, 2, 3], offs_n = [0, 1, 2, 3, 4], transpose = False - # 1. offs_m[:,None] = [[0], - # [1], - # 2. offs_m[:,None] + seqlen_k = [[5], - # [6], - # 3. offs_m[:,None] + seqlen_k - seqlen_q = [[3], - # [4], - # 4. offs_m[:,None] + seqlen_k - seqlen_q - offs_n[None,:] = [[3], - [[0, 1, 2, 3, 4]] = [[ 3, 2, 1, 0,-1], - # [4], [ 4, 3, 2, 1, 0]] - # 5. -1 * alibi_slope * tl.abs(relative_pos_block) = [[ -3, -2, -1, 0,-1], - # [ -4, -3, -2, -1, 0]], - relative_pos_block = offs_m[:, None] + seqlen_k - seqlen_q - offs_n[None, :] - alibi_block = -1 * alibi_slope * tl.abs(relative_pos_block) - if transpose: - return alibi_block.T - else: - return alibi_block - - -@triton.jit -def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn, start_m, - actual_seqlen_k, actual_seqlen_q, dropout_p, philox_seed, batch_philox_offset, exp_scores_ptrs, - block_min, block_max, offs_n_causal, masked_blocks, n_extra_tokens, alibi_slope, score_ptrs, scores_scaled_shifted_ptrs, # pylint: disable=unused-argument +def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn, stride_sn, start_m, + actual_seqlen_k, actual_seqlen_q, dropout_p, philox_seed, philox_ptrs, sd_mask_ptrs, dropout_mask_ptrs, + block_min, block_max, offs_n_causal, masked_blocks, n_extra_tokens, alibi_slope, # pylint: disable=unused-argument IS_CAUSAL: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr, OFFS_M: tl.constexpr, OFFS_N: tl.constexpr, PRE_LOAD_V: tl.constexpr, MASK_STEPS: tl.constexpr, ENABLE_DROPOUT: tl.constexpr, PADDED_HEAD: tl.constexpr, - ACTUAL_BLOCK_DMODEL: tl.constexpr, SM_SCALE: tl.constexpr, USE_EXP2: tl.constexpr, - RETURN_SCORES: tl.constexpr): + ACTUAL_BLOCK_DMODEL: tl.constexpr, SM_SCALE: tl.constexpr, USE_ALIBI: tl.constexpr, USE_EXP2: tl.constexpr, + RETURN_SCORES: tl.constexpr, ACCUMULATOR_TYPE): if USE_EXP2: RCP_LN2: tl.constexpr = 1.4426950408889634 @@ -107,7 +49,7 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri if PRE_LOAD_V: # We can use the same offsets as k, just with dims transposed. v = load_fn(v_ptrs, k_offs_n, k_offs_k, actual_seqlen_k, ACTUAL_BLOCK_DMODEL) - qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=ACCUMULATOR_TYPE) # We start from end of seqlen_k so only the first iteration would need # to be checked for padding if it is not a multiple of block_n # TODO: This can be optimized to only be true for the padded block. @@ -117,18 +59,20 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri # a solution is to always do BLOCK_M // BLOCK_N + 1 steps if not is_modulo_mn. # last step might get wasted but that is okay. check if this masking works For # that case. - if (start_n + BLOCK_N == block_max) and (n_extra_tokens != 0): + if start_n + BLOCK_N == block_max and n_extra_tokens != 0: boundary_m = tl.full([BLOCK_M], actual_seqlen_k, dtype=tl.int32) size_n = start_n + OFFS_N[None, :] mask = size_n < boundary_m[:, None] qk = tl.where(mask, qk, float("-inf")) + # compute masks + q_mask = OFFS_M[:, None] < actual_seqlen_q + k_mask = (start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k + p_mask = q_mask & k_mask + # -- compute qk ---- qk += tl.dot(q, k) - qk_scaled = qk * SM_SCALE - if RETURN_SCORES: - score_mask = (OFFS_M[:, None] < actual_seqlen_q) & ((start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k) - tl.store(score_ptrs, qk_scaled, mask=score_mask) + qk_scaled = qk * SM_SCALE if IS_CAUSAL: causal_boundary = start_n + offs_n_causal @@ -139,8 +83,8 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri bias = load_fn(bias_ptrs, OFFS_M, bias_offs_n, actual_seqlen_q, actual_seqlen_k) qk_scaled += bias - if alibi_slope is not None: - # Compute the global position of each token within the sequence + if USE_ALIBI: + # compute the global position of each token within the sequence global_m_positions = start_m * BLOCK_M + tl.arange(0, BLOCK_M) global_n_positions = start_n + tl.arange(0, BLOCK_N) alibi_block = compute_alibi_block(alibi_slope, actual_seqlen_q, actual_seqlen_k, global_m_positions, @@ -151,10 +95,6 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri # scale and subtract max q_shifted = qk_scaled - m_ij[:, None] - if RETURN_SCORES: - # NOTE: the returned score is not the same as the reference because we need to adjust as we find new maxes per block. We are not doing that - scores_scaled_shifted_mask = (OFFS_M[:, None] < actual_seqlen_q) & ((start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k) - tl.store(scores_scaled_shifted_ptrs, q_shifted, mask=scores_scaled_shifted_mask) # Compute scaled QK and softmax probabilities if USE_EXP2: @@ -165,17 +105,18 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri # CAVEAT: Must update l_ij before applying dropout l_ij = tl.sum(p, 1) if ENABLE_DROPOUT: - philox_offset = batch_philox_offset + start_m * BLOCK_M * actual_seqlen_k + start_n - BLOCK_N - keep = dropout_mask(philox_seed, philox_offset, dropout_p, BLOCK_M, BLOCK_N, actual_seqlen_k) - if RETURN_SCORES: - # NOTE: the returned score is not the same as the reference because we need to adjust as we find new maxes per block. We are not doing that - exp_score_mask = (OFFS_M[:, None] < actual_seqlen_q) & ((start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k) - tl.store(exp_scores_ptrs, tl.where(keep, p, -p), mask=exp_score_mask) - p = tl.where(keep, p, 0.0) + rng_output = tl.rand(philox_seed, philox_ptrs) # TODO: use tl.randint for better performance + dropout_mask = rng_output > dropout_p + + # return scores with negative values for dropped vals + sd_mask = tl.where(dropout_mask, p, -p) + tl.store(sd_mask_ptrs, sd_mask, mask=p_mask) + + # apply dropout mask in place + p = tl.where(dropout_mask, p, 0.0) elif RETURN_SCORES: # NOTE: the returned score is not the same as the reference because we need to adjust as we find new maxes per block. We are not doing that - exp_score_mask = (OFFS_M[:, None] < actual_seqlen_q) & ((start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k) - tl.store(exp_scores_ptrs, p, mask=exp_score_mask) + tl.store(sd_mask_ptrs, p, mask=p_mask) # -- update output accumulator -- # alpha is an adjustment factor for acc and li as we loop and find new maxes @@ -186,7 +127,6 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri else: alpha = tl.math.exp(m_diff) acc = acc * alpha[:, None] - v = None if not PRE_LOAD_V: v = load_fn(v_ptrs, k_offs_n, k_offs_k, actual_seqlen_k, ACTUAL_BLOCK_DMODEL) # -- update m_i and l_i @@ -199,9 +139,11 @@ def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stri if bias_ptrs is not None: bias_ptrs += BLOCK_N * stride_bn if RETURN_SCORES: - score_ptrs += BLOCK_N - scores_scaled_shifted_ptrs += BLOCK_N - exp_scores_ptrs += BLOCK_N + sd_mask_ptrs += BLOCK_N * stride_sn + + if ENABLE_DROPOUT: + dropout_mask_ptrs += BLOCK_N * stride_sn + philox_ptrs += BLOCK_N * stride_sn return acc, l_i, m_i @@ -222,7 +164,7 @@ def get_cdna_autotune_configs(): # Fall-back config. triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1, num_warps=4), - ], ['IS_CAUSAL', 'dropout_p', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K', 'ACTUAL_BLOCK_DMODEL', 'VARLEN', 'HQ', 'HK'] + ], ['IS_CAUSAL', 'dropout_p', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K', 'ACTUAL_BLOCK_DMODEL', 'IS_VARLEN', 'HQ', 'HK'] def get_rdna_autotune_configs(): @@ -242,7 +184,7 @@ def get_rdna_autotune_configs(): # Fall-back config. triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1, num_warps=2), - ], ['IS_CAUSAL', 'dropout_p', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K', 'ACTUAL_BLOCK_DMODEL', 'VARLEN', 'HQ', 'HK'] + ], ['IS_CAUSAL', 'dropout_p', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K', 'ACTUAL_BLOCK_DMODEL', 'IS_VARLEN', 'HQ', 'HK'] def get_autotune_configs(): @@ -266,7 +208,7 @@ def get_autotune_configs(): "MAX_SEQLENS_Q", "MAX_SEQLENS_K", "ACTUAL_BLOCK_DMODEL", - "VARLEN", + "IS_VARLEN", "HQ", "HK", ] @@ -277,37 +219,47 @@ autotune_configs, autotune_keys = get_autotune_configs() @triton.autotune( configs=autotune_configs, key=autotune_keys, - # use_cuda_graph=True, + use_cuda_graph=True, ) @triton.jit -def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_qh, stride_qm, stride_qk, +def attn_fwd(Q, K, V, bias, Cache_seqlens, Cache_batch_idx, # pylint: disable=unused-argument + SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_qh, stride_qm, stride_qk, stride_kz, stride_kh, stride_kn, stride_kk, stride_vz, stride_vh, stride_vk, stride_vn, stride_oz, stride_oh, stride_om, stride_on, stride_bz, stride_bh, stride_bm, stride_bn, stride_az, stride_ah, # pylint: disable=unused-argument stride_sz, stride_sh, stride_sm, stride_sn, stride_lse_z, stride_lse_h, stride_lse_m, cu_seqlens_q, cu_seqlens_k, - dropout_p, philox_seed, philox_offset_base, scores, scores_scaled_shifted, exp_scores, alibi_slopes, HQ: tl.constexpr, + dropout_p, philox_seed, philox_offset_base, sd_mask, dropout_mask, alibi_slopes, HQ: tl.constexpr, HK: tl.constexpr, ACTUAL_BLOCK_DMODEL: tl.constexpr, MAX_SEQLENS_Q: tl.constexpr, - MAX_SEQLENS_K: tl.constexpr, VARLEN: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_M: tl.constexpr, + MAX_SEQLENS_K: tl.constexpr, IS_VARLEN: tl.constexpr, IS_INFERENCE: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr, PRE_LOAD_V: tl.constexpr, USE_BIAS: tl.constexpr, ENABLE_DROPOUT: tl.constexpr, RETURN_SCORES: tl.constexpr, USE_ALIBI: tl.constexpr, USE_EXP2: tl.constexpr): + # set params + ACCUMULATOR_TYPE = tl.float32 + + # compute offsets start_m = tl.program_id(0) off_h_q = tl.program_id(1) off_z = tl.program_id(2) offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = tl.arange(0, BLOCK_N) offs_d = tl.arange(0, BLOCK_DMODEL) - if VARLEN: + + # handle seqlen + if IS_VARLEN: cu_seqlens_q_start = tl.load(cu_seqlens_q + off_z) cu_seqlens_q_end = tl.load(cu_seqlens_q + off_z + 1) - # print("cu_seqlens_q_start:", cu_seqlens_q_start) - seqlen_q = cu_seqlens_q_end - cu_seqlens_q_start - # We have a one-size-fits-all grid in id(0). Some seqlens might be too - # small for all start_m so for those we return early. + + # we have a one-size-fits-all grid in id(0). Some seqlens might be too small for all start_m so for those we return early. if start_m * BLOCK_M > seqlen_q: return cu_seqlens_k_start = tl.load(cu_seqlens_k + off_z) cu_seqlens_k_end = tl.load(cu_seqlens_k + off_z + 1) seqlen_k = cu_seqlens_k_end - cu_seqlens_k_start + elif IS_INFERENCE: + cu_seqlens_q_start = 0 + cu_seqlens_k_start = 0 + seqlen_q = MAX_SEQLENS_Q + seqlen_k = tl.load(Cache_seqlens + off_z) else: cu_seqlens_q_start = 0 cu_seqlens_k_start = 0 @@ -320,14 +272,14 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_ # inf written to LSE. We don't need to do any GEMMs in this case. # This block of code determines what N is, and if this WG is operating # on those M rows. - n_blocks = cdiv_fn(seqlen_k, BLOCK_N) + n_blocks = tl.cdiv(seqlen_k, BLOCK_N) if IS_CAUSAL: # If seqlen_q == seqlen_k, the attn scores are a square matrix. # If seqlen_q != seqlen_k, attn scores are rectangular which means # the causal mask boundary is bottom right aligned, and ends at either # the top edge (seqlen_q < seqlen_k) or left edge. # This captures the decrease in n_blocks if we have a rectangular attn matrix - n_blocks_seqlen = cdiv_fn((start_m + 1) * BLOCK_M + seqlen_k - seqlen_q, BLOCK_N) + n_blocks_seqlen = tl.cdiv((start_m + 1) * BLOCK_M + seqlen_k - seqlen_q, BLOCK_N) # This is what adjusts the block_max for the current WG, only # if IS_CAUSAL. Otherwise we want to always iterate through all n_blocks n_blocks = min(n_blocks, n_blocks_seqlen) @@ -345,7 +297,7 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_ l_offset = LSE + off_z * stride_lse_z + off_h_q * stride_lse_h + cu_seqlens_q_start * stride_lse_m l_ptrs = l_offset + offs_m * stride_lse_m - l = tl.full([BLOCK_M], value=0.0, dtype=tl.float32) + l = tl.full([BLOCK_M], value=0.0, dtype=ACCUMULATOR_TYPE) # mask_m_offsets = start_m + tl.arange(0, BLOCK_M) # lse_mask = mask_m_offsets < causal_start_idx @@ -371,7 +323,7 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_ n_extra_tokens = BLOCK_N - seqlen_k elif seqlen_k % BLOCK_N: n_extra_tokens = seqlen_k % BLOCK_N - PADDED_HEAD: tl.constexpr = ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL + PADDED_HEAD: tl.constexpr = (ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL) # Compute pointers for all the tensors used in this kernel. q_offset = Q + off_z * stride_qz + off_h_q * stride_qh + cu_seqlens_q_start * stride_qm @@ -394,28 +346,23 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_ alibi_slope = None if RETURN_SCORES: - scores_offset = scores + off_z * stride_sz + off_h_q * stride_sh + cu_seqlens_q_start * stride_sm - score_ptrs = scores_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn - - scores_scaled_shifted_offset = scores_scaled_shifted + off_z * stride_sz + off_h_q * stride_sh + cu_seqlens_q_start * stride_sm - scores_scaled_shifted_ptrs = scores_scaled_shifted_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn - - exp_scores_offset = exp_scores + off_z * stride_sz + off_h_q * stride_sh + cu_seqlens_q_start * stride_sm - exp_scores_ptrs = exp_scores_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn + sd_mask_offset = sd_mask + off_z * stride_sz + off_h_q * stride_sh #+ cu_seqlens_q_start * stride_sm + sd_mask_ptrs = sd_mask_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn else: - score_ptrs = None - scores_scaled_shifted_ptrs = None - exp_scores_ptrs = None + sd_mask_ptrs = None if ENABLE_DROPOUT: - off_hz = off_z * HQ + off_h_q - batch_philox_offset = philox_offset_base + off_hz * seqlen_q * seqlen_k + dropout_mask_offset = dropout_mask + off_z * stride_sz + off_h_q * stride_sh #+ cu_seqlens_q_start * stride_sm + dropout_mask_ptrs = dropout_mask_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn + batch_philox_offset = philox_offset_base + off_z * stride_sz + off_h_q * stride_sh #+ cu_seqlens_q_start * stride_sm + philox_ptrs = batch_philox_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn else: - batch_philox_offset = 0 + dropout_mask_ptrs = None + philox_ptrs = 0 # initialize pointer to m and l - m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) - l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32) - acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + m_i = tl.full([BLOCK_M], float("-inf"), dtype=ACCUMULATOR_TYPE) + l_i = tl.full([BLOCK_M], 1.0, dtype=ACCUMULATOR_TYPE) + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=ACCUMULATOR_TYPE) # Q is loaded once at the beginning and shared by all N blocks. q_ptrs_mask = offs_m[:, None] < seqlen_q if PADDED_HEAD: @@ -442,16 +389,16 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_ # value because there is no masking. Similarly we do not need padding. if n_full_blocks > 0: block_max = (n_blocks - masked_blocks) * BLOCK_N - acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn, - start_m, seqlen_k, seqlen_q, dropout_p, philox_seed, batch_philox_offset, - exp_scores_ptrs, + acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn, stride_sn, + start_m, seqlen_k, seqlen_q, dropout_p, philox_seed, philox_ptrs, + sd_mask_ptrs, dropout_mask_ptrs, # _, _, offs_n_causal, masked_blocks, n_extra_tokens, _ - block_min, block_max, 0, 0, 0, alibi_slope, score_ptrs, scores_scaled_shifted_ptrs, + block_min, block_max, 0, 0, 0, alibi_slope, # IS_CAUSAL, .... False, BLOCK_M, BLOCK_DMODEL, BLOCK_N, offs_m, offs_n, # _, MASK_STEPS, ... PRE_LOAD_V, False, ENABLE_DROPOUT, PADDED_HEAD, - ACTUAL_BLOCK_DMODEL, SM_SCALE, USE_EXP2=USE_EXP2, RETURN_SCORES=RETURN_SCORES) + ACTUAL_BLOCK_DMODEL, SM_SCALE, USE_ALIBI=USE_ALIBI, USE_EXP2=USE_EXP2, RETURN_SCORES=RETURN_SCORES, ACCUMULATOR_TYPE=ACCUMULATOR_TYPE) block_min = block_max block_max = n_blocks * BLOCK_N @@ -467,23 +414,25 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_ if USE_BIAS: bias_ptrs += n_full_blocks * BLOCK_N * stride_bn if RETURN_SCORES: - score_ptrs += n_full_blocks * BLOCK_N - scores_scaled_shifted_ptrs += n_full_blocks * BLOCK_N - exp_scores_ptrs += n_full_blocks * BLOCK_N - acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn, - start_m, seqlen_k, seqlen_q, dropout_p, philox_seed, batch_philox_offset, - exp_scores_ptrs, block_min, block_max, offs_n_causal, masked_blocks, - n_extra_tokens, alibi_slope, score_ptrs, scores_scaled_shifted_ptrs, + sd_mask_ptrs += n_full_blocks * BLOCK_N * stride_sn + if ENABLE_DROPOUT: + dropout_mask_ptrs += n_full_blocks * BLOCK_N * stride_sn + philox_ptrs += n_full_blocks * BLOCK_N * stride_sn + acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn, stride_sn, + start_m, seqlen_k, seqlen_q, dropout_p, philox_seed, philox_ptrs, + sd_mask_ptrs, dropout_mask_ptrs, block_min, block_max, offs_n_causal, masked_blocks, + n_extra_tokens, alibi_slope, IS_CAUSAL, BLOCK_M, BLOCK_DMODEL, BLOCK_N, offs_m, offs_n, # _, MASK_STEPS, ... PRE_LOAD_V, True, ENABLE_DROPOUT, PADDED_HEAD, - ACTUAL_BLOCK_DMODEL, SM_SCALE, USE_EXP2=USE_EXP2, RETURN_SCORES=RETURN_SCORES) + ACTUAL_BLOCK_DMODEL, SM_SCALE, USE_ALIBI=USE_ALIBI, USE_EXP2=USE_EXP2, RETURN_SCORES=RETURN_SCORES, ACCUMULATOR_TYPE=ACCUMULATOR_TYPE) # epilogue # This helps the compiler do Newton Raphson on l_i vs on acc which is much larger. l_recip = 1 / l_i[:, None] acc = acc * l_recip if ENABLE_DROPOUT: - acc = acc / (1 - dropout_p) + dropout_scale = 1 / (1 - dropout_p) + acc = acc * dropout_scale # If seqlen_q > seqlen_k but the delta is not a multiple of BLOCK_M, # then we have one block with a row of all NaNs which come from computing # softmax over a row of all -infs (-inf - inf = NaN). We check for that here @@ -491,13 +440,12 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_ end_m_idx = (start_m + 1) * BLOCK_M start_m_idx = start_m * BLOCK_M causal_start_idx = seqlen_q - seqlen_k - acc = acc.to(Out.type.element_ty) if IS_CAUSAL: if causal_start_idx > start_m_idx and causal_start_idx < end_m_idx: out_mask_boundary = tl.full((BLOCK_DMODEL, ), causal_start_idx, dtype=tl.int32) mask_m_offsets = start_m_idx + tl.arange(0, BLOCK_M) out_ptrs_mask = mask_m_offsets[:, None] >= out_mask_boundary[None, :] - z: tl.tensor = 0.0 + z = 0.0 acc = tl.where(out_ptrs_mask, acc, z.to(acc.type.element_ty)) # write back LSE(Log Sum Exponents), the log of the normalization constant @@ -541,30 +489,43 @@ def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_ def attention_prefill_forward_triton_impl( - q, - k, - v, - o, - sm_scale, - alibi_slopes, - causal, - bias, - dropout_p, - layout, - cu_seqlens_q, - cu_seqlens_k, - max_seqlens_q, - max_seqlens_k, - return_scores, - use_exp2): - # check if varlen + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + sm_scale: float, + alibi_slopes: Optional[torch.Tensor], + causal: bool, + bias: Optional[torch.Tensor], + layout: Literal["bshd", "bhsd", "thd"], + # varlen + cu_seqlens_q: Optional[torch.Tensor], + cu_seqlens_k: Optional[torch.Tensor], + max_seqlens_q: int, + max_seqlens_k: int, + # inference + cache_seqlens: Optional[Union[(int, torch.Tensor)]], + cache_batch_idx: Optional[torch.Tensor], + # dropout + dropout_p: float, + philox_seed: Optional[int], + philox_offset: Optional[int], + # misc + return_softmax: bool, + use_exp2: bool, +): + # check flags is_varlen = layout == "thd" + use_alibi, (stride_az, stride_ah) = (True, alibi_slopes.stride()) if alibi_slopes is not None else (False, (0, 0)) + is_inference = cache_seqlens is not None + if is_inference: + assert layout == "bshd", f"{layout} layout is not supported with inference. Use bshd layout" # NOTE: a large bias tensor leads to overflow during pointer arithmetic - if bias is not None: - assert bias.numel() < 2**31 + if (bias is not None): + assert (bias.numel() < 2**31) - batch, nheads_q, nheads_k, head_size, seqlen_q, seqlen_k = get_shape_from_layout(q, k, layout, cu_seqlens_q, cu_seqlens_k, max_seqlens_q, max_seqlens_k) # pylint: disable=unused-variable + batch, nheads_q, nheads_k, head_size, _, _ = get_shapes_from_layout(q, k, layout, cu_seqlens_q, cu_seqlens_k, max_seqlens_q, max_seqlens_k) q_strides, k_strides, v_strides, o_strides = get_strides_from_layout(q, k, v, o, layout) # Get closest power of 2 over or equal to 32. @@ -573,59 +534,45 @@ def attention_prefill_forward_triton_impl( # kernel is padded - there is no padding in memory for any dims. padded_d_model = max(padded_d_model, 16) - grid = lambda META: (triton.cdiv(max_seqlens_q, META['BLOCK_M']), nheads_q, batch) # pylint: disable=unnecessary-lambda-assignment + grid = lambda META: (triton.cdiv(max_seqlens_q, META['BLOCK_M']), nheads_q, batch) - if return_scores: - scores = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device, - dtype=torch.float32) - scores_scaled_shifted = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device, - dtype=torch.float32) - scores_strides = (scores.stride(0), scores.stride(1), scores.stride(2), scores.stride(3)) - else: - scores = None - scores_scaled_shifted = None - scores_strides = (0, 0 , 0 , 0) - - # exp_scores is used to validate dropout behavior vs the PyTorch SDPA math backend reference. We zero this out + # sd_mask is used to validate dropout behavior vs the PyTorch SDPA math backend reference. We zero this out # to give a consistent starting point and then populate it with the output of softmax with the sign bit set according # to the dropout mask. The resulting return allows this mask to be fed into the reference implementation for testing - # only. This return holds no useful output aside from debugging. - if return_scores: - exp_scores = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device, - dtype=torch.float32) + # only. This return holds no useful output aside from debugging. + use_dropout = (dropout_p > 0.0) + if use_dropout or return_softmax: + sd_mask = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device, + dtype=torch.float32) + dropout_mask = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device, + dtype=torch.float32) + scores_strides = (sd_mask.stride(0), sd_mask.stride(1), sd_mask.stride(2), sd_mask.stride(3)) else: - exp_scores = None + sd_mask = None + dropout_mask = None + scores_strides = (0, 0, 0, 0) # stores LSE the log of the normalization constant / sum of expoential score(unnormalzied probablities) if is_varlen: - softmax_lse = torch.empty((q.shape[0], nheads_q), device=q.device, dtype=torch.float32) + softmax_lse = torch.zeros((q.shape[0], nheads_q), device=q.device, dtype=torch.float32) stride_lse_m, stride_lse_h = softmax_lse.stride() stride_lse_z = 0 else: - softmax_lse = torch.empty((batch, nheads_q, max_seqlens_q), device=q.device, dtype=torch.float32) + softmax_lse = torch.zeros((batch, nheads_q, max_seqlens_q), device=q.device, dtype=torch.float32) stride_lse_z, stride_lse_h, stride_lse_m = softmax_lse.stride() - # Seed the RNG so we get reproducible results for testing. - philox_seed = 0x1BF52 - philox_offset = 0x1D4B42 - if bias is not None: bias_strides = (bias.stride(0), bias.stride(1),bias.stride(2), bias.stride(3)) else: bias_strides = (0, 0, 0, 0) - if alibi_slopes is not None: - alibi_strides = (alibi_slopes.stride(0), alibi_slopes.stride(1)) - else: - alibi_strides = (0, 0) - - attn_fwd[grid](q, k, v, bias, sm_scale, softmax_lse, o, *q_strides, *k_strides, *v_strides, *o_strides, - *bias_strides, *alibi_strides, *scores_strides, stride_lse_z, stride_lse_h, stride_lse_m, cu_seqlens_q, cu_seqlens_k, - dropout_p=dropout_p, philox_seed=philox_seed, philox_offset_base=philox_offset, scores=scores, - scores_scaled_shifted=scores_scaled_shifted, exp_scores=exp_scores, alibi_slopes=alibi_slopes, + attn_fwd[grid](q, k, v, bias, cache_seqlens, cache_batch_idx, + sm_scale, softmax_lse, o, *q_strides, *k_strides, *v_strides, *o_strides, + *bias_strides, stride_az, stride_ah, *scores_strides, stride_lse_z, stride_lse_h, stride_lse_m, cu_seqlens_q, cu_seqlens_k, + dropout_p=dropout_p, philox_seed=philox_seed, philox_offset_base=philox_offset, sd_mask=sd_mask, dropout_mask=dropout_mask, alibi_slopes=alibi_slopes, HQ=nheads_q, HK=nheads_k, ACTUAL_BLOCK_DMODEL=head_size, MAX_SEQLENS_Q=max_seqlens_q, - MAX_SEQLENS_K=max_seqlens_k, IS_CAUSAL=causal, VARLEN=is_varlen, + MAX_SEQLENS_K=max_seqlens_k, IS_CAUSAL=causal, IS_VARLEN=is_varlen, IS_INFERENCE=is_inference, BLOCK_DMODEL=padded_d_model, USE_BIAS=False if bias is None else True, - USE_ALIBI=False if alibi_slopes is None else True, ENABLE_DROPOUT=dropout_p - > 0.0, USE_EXP2=use_exp2, RETURN_SCORES=return_scores) + USE_ALIBI=use_alibi, ENABLE_DROPOUT=dropout_p + > 0.0, USE_EXP2=use_exp2, RETURN_SCORES=return_softmax) diff --git a/modules/flash_attn_triton_amd/interface_fa.py b/modules/flash_attn_triton_amd/interface_fa.py index 0bd54e783..573c96017 100644 --- a/modules/flash_attn_triton_amd/interface_fa.py +++ b/modules/flash_attn_triton_amd/interface_fa.py @@ -1,22 +1,16 @@ import torch from modules.flash_attn_triton_amd.fwd_prefill import attention_prefill_forward_triton_impl -from modules.flash_attn_triton_amd.bwd_prefill import attention_prefill_backward_triton_impl -from modules.flash_attn_triton_amd.fwd_decode import attention_decode_forward_triton_impl -from modules.flash_attn_triton_amd.utils import MetaData, get_shape_from_layout +from modules.flash_attn_triton_amd.utils import MetaData -def fwd(q, - k, - v, - dropout_p, - softmax_scale, - causal, +def fwd(q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + dropout_p: float, + softmax_scale: float, + causal: bool ): - if dropout_p != 0.0: - raise ValueError("dropout is not supported on AMD's Triton Backend yet") - - o = torch.empty_like(q) - # Setup metadata metadata = MetaData(sm_scale=softmax_scale) metadata.max_seqlens_q = q.shape[1] @@ -24,251 +18,35 @@ def fwd(q, metadata.layout = "bshd" if causal: - metadata.need_causal() - - #if dropout_p > 0.0: - # metadata.need_dropout(dropout_p, False) - - # Check arguments - metadata.check_args(q, k, v, o) - - attention_prefill_forward_triton_impl( - q, - k, - v, - o, - metadata.sm_scale, - metadata.alibi_slopes, - metadata.causal, - metadata.bias, - metadata.dropout_p, - metadata.layout, - metadata.cu_seqlens_q, - metadata.cu_seqlens_k, - metadata.max_seqlens_q, - metadata.max_seqlens_k, - metadata.return_scores, - metadata.use_exp2) - - return o - - -def bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - alibi_slopes, - dropout_p, - softmax_scale, - causal, - window_size_left, window_size_right, softcap, deterministic, gen_, rng_state, # pylint: disable=unused-argument -): - if dropout_p != 0.0: - raise ValueError("dropout is not supported on AMD yet") - - dq_triton, dk_triton, dv_triton, delta_triton, _, _ = attention_prefill_backward_triton_impl( # pylint: disable=unused-variable - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - softmax_scale, - alibi_slopes, - causal, - "bshd", - None, - None, - None, - None, - False, - ) - delta = delta_triton - - return dq, dk, dv, delta - - -def varlen_fwd( - q, - k, - v, - o, - cu_seqlens_q, - cu_seqlens_k, - seqused_k, leftpad_k, block_table_, # pylint: disable=unused-argument - alibi_slopes,\ - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, # pylint: disable=unused-argument - causal, - window_size_left, window_size_right, softcap, # pylint: disable=unused-argument - return_softmax, - gen_ # pylint: disable=unused-argument -): - if dropout_p != 0.0: - raise ValueError("dropout is not supported on AMD's Triton Backend yet") - - if o is None: - o = torch.empty_like(q) - - # Setup metadata - metadata = MetaData(sm_scale=softmax_scale) - if return_softmax: - metadata.return_scores = True - metadata.set_varlen_params(cu_seqlens_q, cu_seqlens_k) # set layout to "thd" and other metdata - - # get shapes - batch, nheads_q, nheads_k, head_size , seqlen_q, seqlen_k = get_shape_from_layout(q, k, metadata.layout, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k) # pylint: disable=unused-variable - - if causal: - metadata.need_causal() - - if alibi_slopes is not None: - metadata.need_alibi(alibi_slopes, batch, nheads_q) + metadata.need_causal(True) if dropout_p > 0.0: - metadata.need_dropout(dropout_p, return_softmax) + metadata.need_dropout(dropout_p) - # Check arguments - metadata.check_args(q, k, v, o) - if o is None: - o = torch.empty_like(q, dtype=v.dtype) + # check arguments + metadata.check_args(q, k, v, out) + # call implementation attention_prefill_forward_triton_impl( q, k, v, - o, + out, metadata.sm_scale, metadata.alibi_slopes, metadata.causal, - metadata.bias, - metadata.dropout_p, + None, metadata.layout, metadata.cu_seqlens_q, metadata.cu_seqlens_k, metadata.max_seqlens_q, metadata.max_seqlens_k, - metadata.return_scores, + metadata.cache_seqlens, + metadata.cache_batch_idx, + metadata.dropout_p, + metadata.philox_seed, + metadata.philox_offset, + False, metadata.use_exp2) - return o - - -def varlen_bwd( - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - cu_seqlens_q, - cu_seqlens_k, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - dropout_p, - softmax_scale, - zero_tensors, # pylint: disable=unused-argument - causal, - window_size_left, window_size_right, softcap, deterministic, gen_, rng_state, # pylint: disable=unused-argument -): - if dropout_p != 0.0: - raise ValueError("dropout is not supported on AMD yet") - - dq_triton, dk_triton, dv_triton, delta_triton, _, _ = attention_prefill_backward_triton_impl( # pylint: disable=unused-variable - dout, - q, - k, - v, - out, - softmax_lse, - dq, - dk, - dv, - softmax_scale, - alibi_slopes, - causal, - "thd", - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - False, - ) - delta = delta_triton - - return dq, dk, dv, delta - - -def fwd_kvcache( - q, - k_cache, - v_cache, - k, - v, - cache_seqlens, - rotary_cos, rotary_sin, # pylint: disable=unused-argument - cache_batch_idx, - cache_leftpad, block_table, # pylint: disable=unused-argument - alibi_slopes, - out, - softmax_scale, - causal, - window_size_left, window_size_right, softcap, rotary_interleaved, num_splits, # pylint: disable=unused-argument -): - if out is None: - out = torch.empty_like(q) - - # fill metadata - metadata = MetaData(sm_scale=softmax_scale) - metadata.layout = "bshd" - metadata.max_seqlens_q = q.shape[1] - metadata.max_seqlens_k = k_cache.shape[1] - metadata.cache_seqlens = cache_seqlens - metadata.cache_batch_idx = cache_batch_idx - - if k is not None and v is not None: - metadata.new_kv = True - metadata.seqlen_new = k.shape[1] - metadata.k_new = k - metadata.v_new = v - - if causal: - metadata.need_causal() - - if alibi_slopes is not None: - batch, _ , nheads_q, _= q.shape - metadata.need_alibi(alibi_slopes, batch, nheads_q) - - # launch kernel - # TODO: pass output as an arg. Maybe we are copying output which is causing slow down - output, softmax_lse = attention_decode_forward_triton_impl( - q, - k_cache, - v_cache, - metadata.sm_scale, - metadata.causal, - metadata.alibi_slopes, - metadata.layout, - metadata.cache_seqlens, - metadata.cache_batch_idx, - metadata.new_kv, - metadata.k_new, - metadata.v_new, - ) - return output, softmax_lse +# varlen diff --git a/modules/flash_attn_triton_amd/utils.py b/modules/flash_attn_triton_amd/utils.py index 793e29341..b9912aa49 100644 --- a/modules/flash_attn_triton_amd/utils.py +++ b/modules/flash_attn_triton_amd/utils.py @@ -1,33 +1,48 @@ -import os +import csv +import math import torch +import os +import random +import functools import triton +import triton.language as tl +from typing import Literal, Optional, Union from modules.rocm import Agent, MicroArchitecture AUTOTUNE = os.environ.get('FLASH_ATTENTION_TRITON_AMD_AUTOTUNE', '0').lower() in ('1', 'true', 'yes') +USE_REF = os.environ.get('FLASH_ATTENTION_TRITON_AMD_REF', '0').lower() in ('1', 'true', 'yes') PERF = os.environ.get('FLASH_ATTENTION_TRITON_AMD_PERF', '0').lower() in ('1', 'true', 'yes') +# ------------------------------- +# Metadata +# ------------------------------- class MetaData(): - cu_seqlens_q = None - cu_seqlens_k = None - max_seqlens_q = 0 - max_seqlens_k = 0 - bias = None - alibi_slopes = None - causal = False + cu_seqlens_q: Optional[torch.Tensor] = None + cu_seqlens_k: Optional[torch.Tensor] = None + max_seqlens_q: int = 0 + max_seqlens_k: int = 0 + bias: Optional[torch.Tensor] = None + alibi_slopes: Optional[torch.Tensor] = None + causal: bool = False num_contexts = 0 - varlen = False - layout = None - cache_seqlens = None + varlen: bool = False + layout: Optional[Literal["bshd", "bhsd", "thd"]] = None + cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None cache_batch_idx = None - new_kv = False - seqlen_new = None - k_new = None - v_new = None - dropout_p, return_scores= 0.0, False + packing: Optional[bool] = None + return_scores: bool = False + dropout_p: float = 0.0 + philox_seed: Optional[int] = None + philox_offset : Optional[int]= None # if dropout_p > 0.0 seed the RNG so we get reproducible results for testing. # NOTE: scale sm_scale by log_2(e) and use 2^x in the loop as we do not have native e^x support in HW. - use_exp2 = False + use_exp2: bool = False + rotary_sin: Optional[torch.Tensor] = None + rotary_cos: Optional[torch.Tensor] = None + rotary_interleaved: bool = False + rotary_conjunction: bool = False + def __repr__(self) -> str: return (f"MetaData(\n" @@ -44,10 +59,6 @@ class MetaData(): f" layout={self.layout},\n" f" cache_seqlens={self.cache_seqlens},\n" f" cache_batch_idx={self.cache_batch_idx},\n" - f" new_kv={self.new_kv},\n" - f" seqlen_new={self.seqlen_new},\n" - f" k_new={self.k_new},\n" - f" v_new={self.v_new},\n" f" dropout_p={self.dropout_p},\n" f" return_scores={self.return_scores}\n" f")") @@ -55,20 +66,19 @@ class MetaData(): def __init__(self, sm_scale=1.0): self.sm_scale = sm_scale - def set_varlen_params(self, cu_seqlens_q, cu_seqlens_k): + def set_varlen_params(self, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k): self.varlen = True self.layout = 'thd' self.cu_seqlens_q = cu_seqlens_q self.cu_seqlens_k = cu_seqlens_k + self.max_seqlens_q = max_seqlen_q + self.max_seqlens_k = max_seqlen_k + # Without "varlen", there should still be one sequence. assert len(cu_seqlens_q) >= 2 assert len(cu_seqlens_q) == len(cu_seqlens_k) - self.num_contexts = len(cu_seqlens_q) - 1 - for i in range(0, self.num_contexts): - self.max_seqlens_q = max(cu_seqlens_q[i + 1].item() - cu_seqlens_q[i].item(), self.max_seqlens_q) - self.max_seqlens_k = max(cu_seqlens_k[i + 1].item() - cu_seqlens_k[i].item(), self.max_seqlens_k) - def need_bias(self, bias, batch, nheads, seqlen_q, seqlen_k): # pylint: disable=unused-argument + def need_bias(self, bias, batch, nheads, seqlen_q, seqlen_k): assert bias.is_cuda assert bias.dim() == 4 assert bias.shape[0] == 1 @@ -82,17 +92,25 @@ class MetaData(): assert alibi_slopes.shape[1] == nheads self.alibi_slopes = alibi_slopes - def need_causal(self): - self.causal = True + def need_causal(self, causal): + self.causal = causal - def need_dropout(self, dropout_p, return_scores): - self.dropout_p = dropout_p - self.return_scores = return_scores + def need_rotary(self, sin, cos, rotary_interleaved, rotary_conjunction=False): + self.rotary_sin = sin + self.rotary_cos = cos + self.rotary_interleaved = rotary_interleaved + self.rotary_conjunction = rotary_conjunction + + def need_dropout(self, dropout_p, return_scores = True): + if dropout_p > 0.0: + self.dropout_p = dropout_p + self.return_scores = return_scores + self.philox_seed, self.philox_offset = 0x1BF58, 0x1D4B49 def check_args(self, q, k, v, o): assert q.dim() == k.dim() and q.dim() == v.dim() - batch, nheads_q, nheads_k, head_size, _, _ = get_shape_from_layout(q, k, self.layout, self.cu_seqlens_q, self.cu_seqlens_k, self.max_seqlens_q, self.max_seqlens_k) # pylint: disable=unused-variable + batch, nheads_q, nheads_k, head_size, _, _ = get_shapes_from_layout(q, k, self.layout, self.cu_seqlens_q, self.cu_seqlens_k, self.max_seqlens_q, self.max_seqlens_k) if self.varlen: assert q.dim() == 3 assert self.cu_seqlens_q is not None @@ -100,8 +118,6 @@ class MetaData(): assert len(self.cu_seqlens_q) == len(self.cu_seqlens_k) # TODO: Remove once bias is supported with varlen assert self.bias is None - # TODO:Remove once dropout is supported with varlen - assert self.dropout_p == 0.0 # assert not self.return_scores else: assert q.dim() == 4 @@ -111,138 +127,286 @@ class MetaData(): assert q.shape[-1] == k.shape[-1] and q.shape[-1] == v.shape[-1] # TODO: Change assert if we support qkl f8 and v f16 assert q.dtype == k.dtype and q.dtype == v.dtype - assert head_size <= 256 assert o.shape == q.shape assert (nheads_q % nheads_k) == 0 assert self.layout is not None assert self.layout == 'thd' or not self.varlen -def input_helper(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, layout, device="cuda", DEBUG_INPUT=False): +# ------------------------------- +# Input Helper +# ------------------------------- +def random_seqlens_composition(SEQ_LEN, BATCH): + # generate a random composition of N into Z positive parts. + idx = torch.randperm(SEQ_LEN - 1)[: BATCH - 1] + 1 + idx, _ = torch.sort(idx) + breakpoints = torch.cat([ + torch.tensor([0], dtype=torch.long), + idx, + torch.tensor([SEQ_LEN], dtype=torch.long), + ]) + seqlens = (breakpoints[1:] - breakpoints[:-1]).to(torch.int32) + return seqlens + +def generate_varlen_tensor( + total_seqlen: int, + num_heads: int, + head_size: int, + batch_size: Optional[int] = None, + equal_seqlens: bool = False, + device: str = "cuda", + dtype: torch.dtype = torch.float32, + DEBUG_INPUT: bool = False +): + # get valid batch_size + if batch_size is None: + valid_batch_sizes = [bs for bs in [1, 2, 4, 8, 16, 32, 64] if bs <= total_seqlen] + batch_size = random.choice(valid_batch_sizes) + + # get seqlens + if equal_seqlens: + seqlens = torch.full( + (batch_size,), + total_seqlen // batch_size, + dtype=torch.int32, + device=device + ) + seqlens[-1] += total_seqlen % batch_size + else: + seqlens = random_seqlens_composition(total_seqlen, batch_size).to(device=device) + + # create cumulative sequence lengths + cu_seqlens = torch.cat([torch.tensor([0], dtype=torch.int32, device=device), seqlens.cumsum(dim=0)]).to(torch.int32).to(device=device) + max_seqlen = torch.max(seqlens).to(torch.int32).item() + + # create varlen tensor + if DEBUG_INPUT: + x = torch.zeros(total_seqlen, num_heads, head_size, dtype=dtype, device=device) + for i in range(batch_size): + start = cu_seqlens[i].item() + end = cu_seqlens[i+1].item() + length = end - start + + x[start:end, :, :] = ( + torch.arange(length, dtype=dtype, device=device) + .view(length, 1, 1) + .expand(length, num_heads, head_size) + ) + else: + x = torch.randn((total_seqlen, num_heads, head_size), dtype=dtype, device=device) + + x.requires_grad_() + return x, cu_seqlens, max_seqlen + +def generate_bshd_tensor(BATCH, SEQ_LEN, NUM_HEADS, D_HEAD, dtype, device="cuda", DEBUG_INPUT=False): + # gen tensor + tensor_shape = (BATCH, SEQ_LEN, NUM_HEADS, D_HEAD) + if DEBUG_INPUT: + x = torch.arange(SEQ_LEN, dtype=dtype, device=device).view(1, SEQ_LEN, 1, 1).expand(*tensor_shape).contiguous() + else: + x = torch.randn(tensor_shape, dtype=dtype, device=device) + + x.requires_grad_() + return x + +def generate_bhsd_tensor(BATCH, NUM_HEADS, SEQ_LEN, D_HEAD, dtype, device="cuda", DEBUG_INPUT=False): + # gen tensor + tensor_shape = (BATCH, NUM_HEADS, SEQ_LEN, D_HEAD) + if DEBUG_INPUT: + x = torch.arange(SEQ_LEN, dtype=dtype, device=device).view(1, 1, SEQ_LEN, 1).expand(*tensor_shape).contiguous() + else: + x = torch.randn(tensor_shape, dtype=dtype, device=device) + + x.requires_grad_() + return x + +def input_helper( + BATCH: int, + HQ: int, + HK: int, + N_CTX_Q: int, + N_CTX_K: int, + D_HEAD: int, + CAUSAL: bool, + DROPOUT_P: float, + dtype: torch.dtype, + layout: Literal["bshd", "bhsd", "thd"], + packing: Optional[Literal["kv", "qkv"]] = None, + device: Literal["cpu", "cuda"] = "cuda", + DEBUG_INPUT: bool = False, +): torch.manual_seed(20) - # Initialize q, k, v - if layout == 'bhsd': - q_tensor_shape = (Z, HQ, N_CTX_Q, D_HEAD) - k_tensor_shape = (Z, HK, N_CTX_K, D_HEAD) - elif layout == 'bshd': - q_tensor_shape = (Z, N_CTX_Q, HQ, D_HEAD) - k_tensor_shape = (Z, N_CTX_K, HK, D_HEAD) + if layout == "thd": + # set params + TOTAL_SEQLENS_Q = BATCH * N_CTX_Q + TOTAL_SEQLENS_K = BATCH * N_CTX_K + equal_seqlens=False + + # gen tensors + # TODO: the gen functions should maybe have different gen modes like random, ones, increasing seqlen + q, cu_seqlens_q, max_seqlen_q = generate_varlen_tensor(TOTAL_SEQLENS_Q, HQ, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens, DEBUG_INPUT=DEBUG_INPUT) + k, cu_seqlens_k, max_seqlen_k = generate_varlen_tensor(TOTAL_SEQLENS_K, HK, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens, DEBUG_INPUT=DEBUG_INPUT) + v, _, _ = generate_varlen_tensor(TOTAL_SEQLENS_K, HK, D_HEAD, batch_size=BATCH, dtype=dtype, device=device, equal_seqlens=equal_seqlens, DEBUG_INPUT=DEBUG_INPUT) + do = torch.ones_like(q) if DEBUG_INPUT else torch.randn_like(q) + + # setup metadata + if DEBUG_INPUT: + sm_scale = 1 + else: + sm_scale = D_HEAD**-0.5 + metadata = MetaData(sm_scale=sm_scale) + metadata.set_varlen_params(cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k) + metadata.need_causal(CAUSAL) + metadata.need_dropout(DROPOUT_P) + elif layout == 'bshd' or layout == "bhsd": + # gen tensors + if layout == "bshd": + q = generate_bshd_tensor(BATCH, N_CTX_Q, HQ, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) + k = generate_bshd_tensor(BATCH, N_CTX_K, HK, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) + v = generate_bshd_tensor(BATCH, N_CTX_K, HK, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) + do = torch.ones_like(q) if DEBUG_INPUT else torch.randn_like(q) + elif layout == "bhsd": + q = generate_bhsd_tensor(BATCH, HQ, N_CTX_Q, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) + k = generate_bhsd_tensor(BATCH, HK, N_CTX_K, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) + v = generate_bhsd_tensor(BATCH, HK, N_CTX_K, D_HEAD, dtype=dtype, device=device, DEBUG_INPUT=DEBUG_INPUT) + do = torch.ones_like(q) if DEBUG_INPUT else torch.randn_like(q) + + # setup metadata + if DEBUG_INPUT: + sm_scale = 1 + else: + sm_scale = D_HEAD**-0.5 + metadata = MetaData(sm_scale=sm_scale) + metadata.max_seqlens_q = N_CTX_Q + metadata.max_seqlens_k = N_CTX_K + metadata.layout = layout + metadata.need_causal(CAUSAL) + metadata.need_dropout(DROPOUT_P) else: - assert False, f'Got unsupported tensor layout: {layout}' + raise ValueError(f"Unknown layout: {layout}") - q = None - k = None - v = None - - if DEBUG_INPUT: - if layout == "bhsd": - q = torch.arange(N_CTX_Q, dtype=dtype, device=device).view(1, 1, N_CTX_Q, 1).expand(*q_tensor_shape).contiguous().requires_grad_() - k = torch.arange(N_CTX_K, dtype=dtype, device=device).view(1, 1, N_CTX_K, 1).expand(*k_tensor_shape).contiguous().requires_grad_() - v = torch.arange(N_CTX_K, dtype=dtype, device=device).view(1, 1, N_CTX_K, 1).expand(*k_tensor_shape).contiguous().requires_grad_() + # deal with packing + if packing is None: + return q, k, v, do, metadata + elif packing == "kv": + # pack k and v + if layout in ["bhsd", "thd"]: + kv = torch.stack([k, v], dim=1) elif layout == "bshd": - q = torch.arange(N_CTX_Q, dtype=dtype, device=device).view(1, N_CTX_Q, 1, 1).expand(*q_tensor_shape).contiguous().requires_grad_() - k = torch.arange(N_CTX_K, dtype=dtype, device=device).view(1, N_CTX_K, 1, 1).expand(*k_tensor_shape).contiguous().requires_grad_() - v = torch.arange(N_CTX_K, dtype=dtype, device=device).view(1, N_CTX_K, 1, 1).expand(*k_tensor_shape).contiguous().requires_grad_() + kv = torch.stack([k, v], dim=2) + else: + raise ValueError(f"Unknown layout: {layout}") + + return q, kv, do, metadata + elif packing == "qkv": + # qkv packing - requires same sequence length for q and k + assert N_CTX_Q == N_CTX_K, "For QKV packing, Q and K must have same sequence length" + assert HQ == HK, "For QKV packing, Q and K must have same number of heads" + + # pack q, k, and v + if layout in ["bhsd", "thd"]: + qkv = torch.stack([q, k, v], dim=1) + elif layout == "bshd": + qkv = torch.stack([q, k, v], dim=2) + else: + raise ValueError(f"Unknown layout: {layout}") + + return qkv, do, metadata else: - q = torch.randn(q_tensor_shape, dtype=dtype, device=device, requires_grad=True) - k = torch.randn(k_tensor_shape, dtype=dtype, device=device, requires_grad=True) - v = torch.randn(k_tensor_shape, dtype=dtype, device=device, requires_grad=True) + assert False, f"Unsupported packing mode: {packing}" - if DEBUG_INPUT: - sm_scale = 1 +# ------------------------------- +# Alibi +# ------------------------------- +@triton.jit +def compute_alibi_block(alibi_slope, seqlen_q, seqlen_k, offs_m, offs_n, transpose=False): + # when seqlen_k and seqlen_q are different we want the diagonal to stick to the bottom right of the attention matrix + # for casual mask we want something like this where (1 is kept and 0 is masked) + # seqlen_q = 2 and seqlen_k = 5 + # 1 1 1 1 0 + # 1 1 1 1 1 + # seqlen_q = 5 and seqlen_k = 2 + # 0 0 + # 0 0 + # 0 0 + # 1 0 + # 1 1 + # for alibi the diagonal is 0 indicating no penalty for attending to that spot and increasing penalty for attending further from the diagonal + # e.g. alibi_slope = 1, seqlen_q = 2, seqlen_k = 5, offs_m = [0, 1, 2, 3], offs_n = [0, 1, 2, 3, 4], transpose = False + # 1. offs_m[:,None] = [[0], + # [1], + # 2. offs_m[:,None] + seqlen_k = [[5], + # [6], + # 3. offs_m[:,None] + seqlen_k - seqlen_q = [[3], + # [4], + # 4. offs_m[:,None] + seqlen_k - seqlen_q - offs_n[None,:] = [[3], - [[0, 1, 2, 3, 4]] = [[ 3, 2, 1, 0,-1], + # [4], [ 4, 3, 2, 1, 0]] + # 5. -1 * alibi_slope * tl.abs(relative_pos_block) = [[ -3, -2, -1, 0,-1], + # [ -4, -3, -2, -1, 0]], + relative_pos_block = offs_m[:, None] + seqlen_k - seqlen_q - offs_n[None, :] + alibi_block = -1 * alibi_slope * tl.abs(relative_pos_block) + if transpose: + return alibi_block.T else: - sm_scale = D_HEAD**-0.5 - input_metadata = MetaData(sm_scale=sm_scale) - input_metadata.max_seqlens_q = N_CTX_Q - input_metadata.max_seqlens_k = N_CTX_K - input_metadata.layout = layout - return q, k, v, input_metadata + return alibi_block - -def varlen_input_helper(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, device="cuda", equal_seqlens=False, DEBUG_INPUT=False): - torch.manual_seed(20) - - # Random or equal sequence lengths based on 'equal_seqlens' flag - if not equal_seqlens: - max_seqlens_q = N_CTX_Q // Z - max_seqlens_k = N_CTX_K // Z - seqlens_q = torch.randint(1, max_seqlens_q + 1, (Z,), dtype=torch.int32) - seqlens_k = torch.randint(1, max_seqlens_k + 1, (Z,), dtype=torch.int32) - else: - seqlens_q = torch.full((Z,), N_CTX_Q // Z, dtype=torch.int32) - seqlens_k = torch.full((Z,), N_CTX_K // Z, dtype=torch.int32) - - # Calculate cumulative sequence lengths - cu_seqlens_q = torch.cat([torch.tensor([0], dtype=torch.int32), seqlens_q.cumsum(dim=0)]) - cu_seqlens_k = torch.cat([torch.tensor([0], dtype=torch.int32), seqlens_k.cumsum(dim=0)]) - cu_seqlens_q = cu_seqlens_q.to(device=device).to(torch.int32) - cu_seqlens_k = cu_seqlens_k.to(device=device).to(torch.int32) - - # Total lengths - total_q = cu_seqlens_q[-1].item() - total_k = cu_seqlens_k[-1].item() - - if DEBUG_INPUT: - # Initialize q, k, v with deterministic values - q = torch.arange(total_q, dtype=dtype, device=device).view(total_q, 1, 1) - q = q.expand(total_q, HQ, D_HEAD).contiguous().requires_grad_() - k = torch.arange(total_k, dtype=dtype, device=device).view(total_k, 1, 1) - k = k.expand(total_k, HK, D_HEAD).contiguous().requires_grad_() - v = torch.arange(total_k, dtype=dtype, device=device).view(total_k, 1, 1) - v = v.expand(total_k, HK, D_HEAD).contiguous().requires_grad_() - sm_scale = 1 - else: - # Initialize q, k, v with random values - q = torch.randn((total_q, HQ, D_HEAD), dtype=dtype, device=device).requires_grad_() - k = torch.randn((total_k, HK, D_HEAD), dtype=dtype, device=device).requires_grad_() - v = torch.randn((total_k, HK, D_HEAD), dtype=dtype, device=device).requires_grad_() - sm_scale = D_HEAD ** -0.5 - - input_metadata = MetaData(sm_scale=sm_scale) - input_metadata.set_varlen_params(cu_seqlens_q, cu_seqlens_k) - return q, k, v, input_metadata - - -def get_shape_from_layout(q, k, layout, cu_seqlens_q = None, cu_seqlens_k = None, max_seqlen_q=None, max_seqlen_k=None): +# ------------------------------- +# Misc +# ------------------------------- +def get_shape_from_layout( + x: torch.Tensor, + layout: Literal["bshd", "bhsd", "thd"], + cu_seqlens: Optional[torch.Tensor] = None, + max_seqlen: Optional[int] = None, +) -> tuple[int, int, int, int]: if layout == 'bhsd': - batch_q, nheads_q, max_seqlen_q, head_size_q = q.shape - batch_k, nheads_k, max_seqlen_k, head_size_k = k.shape + batch, num_heads, max_seqlen_final, head_dim = x.shape elif layout == 'bshd': - batch_q, max_seqlen_q, nheads_q, head_size_q = q.shape - batch_k, max_seqlen_k, nheads_k, head_size_k = k.shape + batch, max_seqlen_final, num_heads, head_dim = x.shape elif layout == 'thd': - batch_q, max_seqlen_q, nheads_q, head_size_q = len(cu_seqlens_q) - 1, max_seqlen_q, q.shape[1], q.shape[2] # pylint: disable=self-assigning-variable - batch_k, max_seqlen_k, nheads_k, head_size_k = len(cu_seqlens_k) - 1, max_seqlen_k, k.shape[1], k.shape[2] # pylint: disable=self-assigning-variable + total_seqlen, num_heads, head_dim = x.shape + if cu_seqlens is None: + raise ValueError("cu_seqlens must be provided for varlen (thd) layout") + if max_seqlen is None: + raise ValueError("max_seqlen must be provided for varlen (thd) layout") + + batch, max_seqlen_final, num_heads, head_dim = len(cu_seqlens) - 1, max_seqlen, num_heads, head_dim else: assert False, "Got unsupported layout." + return batch, max_seqlen_final, num_heads, head_dim + + +def get_shapes_from_layout(q, k, layout, cu_seqlens_q = None, cu_seqlens_k = None, max_seqlen_q=None, max_seqlen_k=None): + batch_q, seqlen_q, nheads_q, head_size_q = get_shape_from_layout(q, layout, cu_seqlens_q, max_seqlen_q) + batch_k, seqlen_k, nheads_k, head_size_k = get_shape_from_layout(k, layout, cu_seqlens_k, max_seqlen_k) + # assert assert batch_q == batch_k assert head_size_q == head_size_k - return batch_q, nheads_q, nheads_k, head_size_q, max_seqlen_q, max_seqlen_k + return batch_q, nheads_q, nheads_k, head_size_q, seqlen_q, seqlen_k - -def get_strides_from_layout(q, k, v, o, layout): +def get_stride_from_layout(x: torch.Tensor, layout:Literal["bshd", "bhsd", "thd"]): if layout == 'thd': - q_strides = (0, q.stride(1), q.stride(0), q.stride(2)) - k_strides = (0, k.stride(1), k.stride(0), k.stride(2)) - v_strides = (0, v.stride(1), v.stride(0), v.stride(2)) - o_strides = (0, o.stride(1), o.stride(0), o.stride(2)) + strides = (0, x.stride(1), x.stride(0), x.stride(2)) elif layout == 'bhsd': - q_strides = (q.stride(0), q.stride(1), q.stride(2), q.stride(3)) - k_strides = (k.stride(0), k.stride(1), k.stride(2), k.stride(3)) - v_strides = (v.stride(0), v.stride(1), v.stride(2), v.stride(3)) - o_strides = (o.stride(0), o.stride(1), o.stride(2), o.stride(3)) + strides = (x.stride(0), x.stride(1), x.stride(2), x.stride(3)) elif layout == 'bshd': - q_strides = (q.stride(0), q.stride(2), q.stride(1), q.stride(3)) - k_strides = (k.stride(0), k.stride(2), k.stride(1), k.stride(3)) - v_strides = (v.stride(0), v.stride(2), v.stride(1), v.stride(3)) - o_strides = (o.stride(0), o.stride(2), o.stride(1), o.stride(3)) + strides = (x.stride(0), x.stride(2), x.stride(1), x.stride(3)) else: assert False, 'Got unsupported layout.' - return q_strides, k_strides, v_strides, o_strides + return strides +def get_shape_and_strides_from_layout(x: torch.Tensor, layout: Literal["bshd", "bhsd", "thd"], cu_seqlens: Optional[torch.Tensor] = None, max_seqlen: Optional[int] = None): + return get_shape_from_layout(x, layout, cu_seqlens, max_seqlen), get_stride_from_layout(x, layout) + +def get_strides_from_layout(q, k, v, o, layout): + q_strides = get_stride_from_layout(q, layout) + k_strides = get_stride_from_layout(k, layout) + v_strides = get_stride_from_layout(v, layout) + o_strides = get_stride_from_layout(o, layout) + return q_strides, k_strides, v_strides, o_strides def get_padded_headsize(size): # Get closest power of 2 over or equal to 32. @@ -252,24 +416,79 @@ def get_padded_headsize(size): padded_d_model = max(padded_d_model, 16) return padded_d_model +def compute_alibi_tensor_ref(alibi_slopes, seqlen_q, seqlen_k): + q_idx = torch.arange(seqlen_q, dtype=torch.int32, device="cuda").unsqueeze(-1) # (N_CTX_Q, 1) + k_idx = torch.arange(seqlen_k, dtype=torch.int32, device="cuda").unsqueeze(0) # (1, N_CTX_K) + relative_pos = torch.abs(q_idx + seqlen_k - seqlen_q - k_idx) # (N_CTX_Q, N_CTX_K) + return -1 * alibi_slopes.unsqueeze(-1).unsqueeze(-1) * relative_pos # (Z, H, N_CTX_Q, N_CTX_K) -def _strides(x: torch.Tensor, *stride_names: str): - if x is None: - return {f"stride_{s}": 0 for i, s in enumerate(stride_names)} +# ------------------------------- +# Dropouts +# ------------------------------- +def create_dropout_mask(dropout_p, shape, seed): + device = "cuda" + rand_vals = torch.rand(shape, generator=torch.Generator(device=device).manual_seed(seed), device=device, dtype=torch.float32) + return rand_vals > dropout_p - assert x.ndim == len(stride_names) - return {f"stride_{s}": x.stride(i) for i, s in enumerate(stride_names)} +def create_dropout_mask_varlen(dropout_p, batch, nheads_q, cu_seqlens_q, cu_seqlens_k, philox_seed): + device = "cuda" + qlens = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]) + klens = (cu_seqlens_k[1:] - cu_seqlens_k[:-1]) + max_qlen = qlens.max() + max_klen = klens.max() + dropout_mask = torch.zeros((batch, nheads_q, max_qlen, max_klen), device=device) + for b in range(batch): + qlen = qlens[b] + klen = klens[b] + rand_vals = torch.rand((nheads_q, qlen, klen), generator=torch.Generator(device=device).manual_seed(philox_seed), device=device, dtype=torch.float32) + submask = rand_vals > dropout_p + dropout_mask[b, :, :qlen, :klen] = submask + return dropout_mask -def get_input_shapes(): - cases = [(max(1, 2**(16 - i)), 1, 2**i, 16, 1, 128) - for i in range(8, 18)] + [(max(1, 2**(16 - i)), 1, 2**i, 16, 2, 128) for i in range(8, 18)] - return cases +def write_dropout_mask(x, tensor_name = "tensor"): + batch, head, seqlen_m, seqlen_n = x.shape + x = x.tolist() + with open(f'{tensor_name}.csv', 'w') as f: + writer = csv.writer(f) + for b in range(batch): + for h in range(head): + dropout_mask = x[b][h] + if True: + BLOCK_M = 64 + BLOCK_N = 64 + + # Calculate number of blocks in each dimension + m_blocks = math.ceil(seqlen_m / BLOCK_M) + n_blocks = math.ceil(seqlen_n / BLOCK_N) + + # Process each block + for m_block in range(m_blocks): + # Calculate row range for current block + row_start = m_block * BLOCK_M + row_end = min(row_start + BLOCK_M, seqlen_m) + + for n_block in range(n_blocks): + # Calculate column range for current block + col_start = n_block * BLOCK_N + col_end = min(col_start + BLOCK_N, seqlen_n) + + # Extract and write the current block + for row_idx in range(row_start, row_end): + row_data = dropout_mask[row_idx][col_start:col_end] + writer.writerow(row_data) + else: + writer.writerows(dropout_mask) +# ------------------------------- +# Runtime info +# ------------------------------- +@functools.cache def is_cdna(): return Agent(triton.runtime.driver.active.get_current_target().arch).arch == MicroArchitecture.CDNA +@functools.cache def is_rdna(): return Agent(triton.runtime.driver.active.get_current_target().arch).arch == MicroArchitecture.RDNA diff --git a/modules/hidream/pipeline_hidream_image_editing.py b/modules/hidream/pipeline_hidream_image_editing.py new file mode 100644 index 000000000..c410ec8ac --- /dev/null +++ b/modules/hidream/pipeline_hidream_image_editing.py @@ -0,0 +1,1153 @@ +import inspect +from typing import Any, Callable, Dict, List, Optional, Union +import PIL + +import torch +from transformers import ( + CLIPTextModelWithProjection, + CLIPTokenizer, + LlamaForCausalLM, + PreTrainedTokenizerFast, + T5EncoderModel, + T5Tokenizer, +) + +from diffusers.image_processor import VaeImageProcessor, PipelineImageInput +from diffusers.loaders import HiDreamImageLoraLoaderMixin +from diffusers.models import AutoencoderKL, HiDreamImageTransformer2DModel +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler +from diffusers.utils import deprecate, is_torch_xla_available, logging, replace_example_docstring +from diffusers.utils.torch_utils import randn_tensor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.pipelines.hidream_image.pipeline_output import HiDreamImagePipelineOutput + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from transformers import PreTrainedTokenizerFast, LlamaForCausalLM + >>> from diffusers import UniPCMultistepScheduler + >>> from pipeline_hidream_image_editing import HiDreamImageEditingPipeline + >>> from PIL import Image + + + >>> tokenizer_4 = PreTrainedTokenizerFast.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct") + >>> text_encoder_4 = LlamaForCausalLM.from_pretrained( + ... "meta-llama/Meta-Llama-3.1-8B-Instruct", + ... output_hidden_states=True, + ... output_attentions=True, + ... torch_dtype=torch.bfloat16, + ... ) + + >>> pipe = HiDreamImageEditingPipeline.from_pretrained( + ... "HiDream-ai/HiDream-E1-Full", + ... tokenizer_4=tokenizer_4, + ... text_encoder_4=text_encoder_4, + ... torch_dtype=torch.bfloat16, + ... ) + >>> pipe.enable_model_cpu_offload() + + >>> # Load input image for editing + >>> input_image = Image.open("your_image.jpg") + >>> input_image = input_image.resize((768, 768)) + + >>> # Edit the image based on instructions + >>> image = pipe( + ... prompt='Editing Instruction: Convert the image into a Ghibli style. Target Image Description: A person in a light pink t-shirt with short dark hair, depicted in a Ghibli style against a plain background.', + ... negative_prompt="low resolution, blur", + ... image=input_image, + ... guidance_scale=5.0, + ... image_guidance_scale=4.0, + ... num_inference_steps=28, + ... generator=torch.Generator("cuda").manual_seed(3), + ... ).images[0] + >>> image.save("edited_output.png") + ``` +""" + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + timesteps: Optional[List[int]] = None, + sigmas: Optional[List[float]] = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class HiDreamImageEditingPipeline(DiffusionPipeline, HiDreamImageLoraLoaderMixin): + model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->text_encoder_4->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds_t5", "prompt_embeds_llama3", "pooled_prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKL, + text_encoder: CLIPTextModelWithProjection, + tokenizer: CLIPTokenizer, + text_encoder_2: CLIPTextModelWithProjection, + tokenizer_2: CLIPTokenizer, + text_encoder_3: T5EncoderModel, + tokenizer_3: T5Tokenizer, + text_encoder_4: LlamaForCausalLM, + tokenizer_4: PreTrainedTokenizerFast, + transformer: HiDreamImageTransformer2DModel, + ): + super().__init__() + + self.register_modules( + vae=vae, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + text_encoder_3=text_encoder_3, + text_encoder_4=text_encoder_4, + tokenizer=tokenizer, + tokenizer_2=tokenizer_2, + tokenizer_3=tokenizer_3, + tokenizer_4=tokenizer_4, + scheduler=scheduler, + transformer=transformer, + ) + self.vae_scale_factor = ( + 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8 + ) + # HiDreamImage latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible + # by the patch size. So the vae scale factor is multiplied by the patch size to account for this + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + self.default_sample_size = 128 + if getattr(self, "tokenizer_4", None) is not None: + self.tokenizer_4.pad_token = self.tokenizer_4.eos_token + + def _get_t5_prompt_embeds( + self, + prompt: Union[str, List[str]] = None, + max_sequence_length: int = 128, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder_3.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + + text_inputs = self.tokenizer_3( + prompt, + padding="max_length", + max_length=min(max_sequence_length, self.tokenizer_3.model_max_length), + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + attention_mask = text_inputs.attention_mask + untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = self.tokenizer_3.batch_decode( + untruncated_ids[:, min(max_sequence_length, self.tokenizer_3.model_max_length) - 1 : -1] + ) + logger.warning( + "The following part of your input was truncated because `max_sequence_length` is set to " + f" {min(max_sequence_length, self.tokenizer_3.model_max_length)} tokens: {removed_text}" + ) + + prompt_embeds = self.text_encoder_3(text_input_ids.to(device), attention_mask=attention_mask.to(device))[0] + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + return prompt_embeds + + def _get_clip_prompt_embeds( + self, + tokenizer, + text_encoder, + prompt: Union[str, List[str]], + max_sequence_length: int = 128, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + device = device or self._execution_device + dtype = dtype or text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + + text_inputs = tokenizer( + prompt, + padding="max_length", + max_length=min(max_sequence_length, 218), + truncation=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids + untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = tokenizer.batch_decode(untruncated_ids[:, 218 - 1 : -1]) + logger.warning( + "The following part of your input was truncated because CLIP can only handle sequences up to" + f" {218} tokens: {removed_text}" + ) + prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True) + + # Use pooled output of CLIPTextModel + prompt_embeds = prompt_embeds[0] + prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) + return prompt_embeds + + def _get_llama3_prompt_embeds( + self, + prompt: Union[str, List[str]] = None, + max_sequence_length: int = 128, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder_4.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + + text_inputs = self.tokenizer_4( + prompt, + padding="max_length", + max_length=min(max_sequence_length, self.tokenizer_4.model_max_length), + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + attention_mask = text_inputs.attention_mask + untruncated_ids = self.tokenizer_4(prompt, padding="longest", return_tensors="pt").input_ids + + if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): + removed_text = self.tokenizer_4.batch_decode( + untruncated_ids[:, min(max_sequence_length, self.tokenizer_4.model_max_length) - 1 : -1] + ) + logger.warning( + "The following part of your input was truncated because `max_sequence_length` is set to " + f" {min(max_sequence_length, self.tokenizer_4.model_max_length)} tokens: {removed_text}" + ) + + outputs = self.text_encoder_4( + text_input_ids.to(device), + attention_mask=attention_mask.to(device), + output_hidden_states=True, + output_attentions=True, + ) + + prompt_embeds = outputs.hidden_states[1:] + prompt_embeds = torch.stack(prompt_embeds, dim=0) + return prompt_embeds + + def encode_prompt( + self, + prompt: Optional[Union[str, List[str]]] = None, + prompt_2: Optional[Union[str, List[str]]] = None, + prompt_3: Optional[Union[str, List[str]]] = None, + prompt_4: Optional[Union[str, List[str]]] = None, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + num_images_per_prompt: int = 1, + do_classifier_free_guidance: bool = True, + negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt_2: Optional[Union[str, List[str]]] = None, + negative_prompt_3: Optional[Union[str, List[str]]] = None, + negative_prompt_4: Optional[Union[str, List[str]]] = None, + prompt_embeds_t5: Optional[List[torch.FloatTensor]] = None, + prompt_embeds_llama3: Optional[List[torch.FloatTensor]] = None, + negative_prompt_embeds_t5: Optional[List[torch.FloatTensor]] = None, + negative_prompt_embeds_llama3: Optional[List[torch.FloatTensor]] = None, + pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + max_sequence_length: int = 128, + lora_scale: Optional[float] = None, + ): + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = pooled_prompt_embeds.shape[0] + + device = device or self._execution_device + + if pooled_prompt_embeds is None: + pooled_prompt_embeds_1 = self._get_clip_prompt_embeds( + self.tokenizer, self.text_encoder, prompt, max_sequence_length, device, dtype + ) + + if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + + if len(negative_prompt) > 1 and len(negative_prompt) != batch_size: + raise ValueError(f"negative_prompt must be of length 1 or {batch_size}") + + negative_pooled_prompt_embeds_1 = self._get_clip_prompt_embeds( + self.tokenizer, self.text_encoder, negative_prompt, max_sequence_length, device, dtype + ) + + if negative_pooled_prompt_embeds_1.shape[0] == 1 and batch_size > 1: + negative_pooled_prompt_embeds_1 = negative_pooled_prompt_embeds_1.repeat(batch_size, 1) + + if pooled_prompt_embeds is None: + prompt_2 = prompt_2 or prompt + prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 + + if len(prompt_2) > 1 and len(prompt_2) != batch_size: + raise ValueError(f"prompt_2 must be of length 1 or {batch_size}") + + pooled_prompt_embeds_2 = self._get_clip_prompt_embeds( + self.tokenizer_2, self.text_encoder_2, prompt_2, max_sequence_length, device, dtype + ) + + if pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1: + pooled_prompt_embeds_2 = pooled_prompt_embeds_2.repeat(batch_size, 1) + + if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: + negative_prompt_2 = negative_prompt_2 or negative_prompt + negative_prompt_2 = [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 + + if len(negative_prompt_2) > 1 and len(negative_prompt_2) != batch_size: + raise ValueError(f"negative_prompt_2 must be of length 1 or {batch_size}") + + negative_pooled_prompt_embeds_2 = self._get_clip_prompt_embeds( + self.tokenizer_2, self.text_encoder_2, negative_prompt_2, max_sequence_length, device, dtype + ) + + if negative_pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1: + negative_pooled_prompt_embeds_2 = negative_pooled_prompt_embeds_2.repeat(batch_size, 1) + + if pooled_prompt_embeds is None: + pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_1, pooled_prompt_embeds_2], dim=-1) + + if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: + negative_pooled_prompt_embeds = torch.cat( + [negative_pooled_prompt_embeds_1, negative_pooled_prompt_embeds_2], dim=-1 + ) + + if prompt_embeds_t5 is None: + prompt_3 = prompt_3 or prompt + prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3 + + if len(prompt_3) > 1 and len(prompt_3) != batch_size: + raise ValueError(f"prompt_3 must be of length 1 or {batch_size}") + + prompt_embeds_t5 = self._get_t5_prompt_embeds(prompt_3, max_sequence_length, device, dtype) + + if prompt_embeds_t5.shape[0] == 1 and batch_size > 1: + prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1) + + if do_classifier_free_guidance and negative_prompt_embeds_t5 is None: + negative_prompt_3 = negative_prompt_3 or negative_prompt + negative_prompt_3 = [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3 + + if len(negative_prompt_3) > 1 and len(negative_prompt_3) != batch_size: + raise ValueError(f"negative_prompt_3 must be of length 1 or {batch_size}") + + negative_prompt_embeds_t5 = self._get_t5_prompt_embeds( + negative_prompt_3, max_sequence_length, device, dtype + ) + + if negative_prompt_embeds_t5.shape[0] == 1 and batch_size > 1: + negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1) + + if prompt_embeds_llama3 is None: + prompt_4 = prompt_4 or prompt + prompt_4 = [prompt_4] if isinstance(prompt_4, str) else prompt_4 + + if len(prompt_4) > 1 and len(prompt_4) != batch_size: + raise ValueError(f"prompt_4 must be of length 1 or {batch_size}") + + prompt_embeds_llama3 = self._get_llama3_prompt_embeds(prompt_4, max_sequence_length, device, dtype) + + if prompt_embeds_llama3.shape[0] == 1 and batch_size > 1: + prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1) + + if do_classifier_free_guidance and negative_prompt_embeds_llama3 is None: + negative_prompt_4 = negative_prompt_4 or negative_prompt + negative_prompt_4 = [negative_prompt_4] if isinstance(negative_prompt_4, str) else negative_prompt_4 + + if len(negative_prompt_4) > 1 and len(negative_prompt_4) != batch_size: + raise ValueError(f"negative_prompt_4 must be of length 1 or {batch_size}") + + negative_prompt_embeds_llama3 = self._get_llama3_prompt_embeds( + negative_prompt_4, max_sequence_length, device, dtype + ) + + if negative_prompt_embeds_llama3.shape[0] == 1 and batch_size > 1: + negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1) + + # duplicate pooled_prompt_embeds for each generation per prompt + pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt) + pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) + + # duplicate t5_prompt_embeds for batch_size and num_images_per_prompt + bs_embed, seq_len, _ = prompt_embeds_t5.shape + if bs_embed == 1 and batch_size > 1: + prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1) + elif bs_embed > 1 and bs_embed != batch_size: + raise ValueError(f"cannot duplicate prompt_embeds_t5 of batch size {bs_embed}") + prompt_embeds_t5 = prompt_embeds_t5.repeat(1, num_images_per_prompt, 1) + prompt_embeds_t5 = prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1) + + # duplicate llama3_prompt_embeds for batch_size and num_images_per_prompt + _, bs_embed, seq_len, dim = prompt_embeds_llama3.shape + if bs_embed == 1 and batch_size > 1: + prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1) + elif bs_embed > 1 and bs_embed != batch_size: + raise ValueError(f"cannot duplicate prompt_embeds_llama3 of batch size {bs_embed}") + prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1) + prompt_embeds_llama3 = prompt_embeds_llama3.view(-1, batch_size * num_images_per_prompt, seq_len, dim) + + if do_classifier_free_guidance: + # duplicate negative_pooled_prompt_embeds for batch_size and num_images_per_prompt + bs_embed, seq_len = negative_pooled_prompt_embeds.shape + if bs_embed == 1 and batch_size > 1: + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(batch_size, 1) + elif bs_embed > 1 and bs_embed != batch_size: + raise ValueError(f"cannot duplicate negative_pooled_prompt_embeds of batch size {bs_embed}") + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt) + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) + + # duplicate negative_t5_prompt_embeds for batch_size and num_images_per_prompt + bs_embed, seq_len, _ = negative_prompt_embeds_t5.shape + if bs_embed == 1 and batch_size > 1: + negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1) + elif bs_embed > 1 and bs_embed != batch_size: + raise ValueError(f"cannot duplicate negative_prompt_embeds_t5 of batch size {bs_embed}") + negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(1, num_images_per_prompt, 1) + negative_prompt_embeds_t5 = negative_prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1) + + # duplicate negative_prompt_embeds_llama3 for batch_size and num_images_per_prompt + _, bs_embed, seq_len, dim = negative_prompt_embeds_llama3.shape + if bs_embed == 1 and batch_size > 1: + negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1) + elif bs_embed > 1 and bs_embed != batch_size: + raise ValueError(f"cannot duplicate negative_prompt_embeds_llama3 of batch size {bs_embed}") + negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1) + negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.view( + -1, batch_size * num_images_per_prompt, seq_len, dim + ) + + return ( + prompt_embeds_t5, + negative_prompt_embeds_t5, + prompt_embeds_llama3, + negative_prompt_embeds_llama3, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + ) + + def enable_vae_slicing(self): + r""" + Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to + compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. + """ + self.vae.enable_slicing() + + def disable_vae_slicing(self): + r""" + Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to + computing decoding in one step. + """ + self.vae.disable_slicing() + + def enable_vae_tiling(self): + r""" + Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to + compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow + processing larger images. + """ + self.vae.enable_tiling() + + def disable_vae_tiling(self): + r""" + Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to + computing decoding in one step. + """ + self.vae.disable_tiling() + + def check_inputs( + self, + prompt, + prompt_2, + prompt_3, + prompt_4, + negative_prompt=None, + negative_prompt_2=None, + negative_prompt_3=None, + negative_prompt_4=None, + prompt_embeds_t5=None, + prompt_embeds_llama3=None, + negative_prompt_embeds_t5=None, + negative_prompt_embeds_llama3=None, + pooled_prompt_embeds=None, + negative_pooled_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + ): + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and pooled_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `pooled_prompt_embeds`: {pooled_prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt_2 is not None and pooled_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt_2`: {prompt_2} and `pooled_prompt_embeds`: {pooled_prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt_3 is not None and prompt_embeds_t5 is not None: + raise ValueError( + f"Cannot forward both `prompt_3`: {prompt_3} and `prompt_embeds_t5`: {prompt_embeds_t5}. Please make sure to" + " only forward one of the two." + ) + elif prompt_4 is not None and prompt_embeds_llama3 is not None: + raise ValueError( + f"Cannot forward both `prompt_4`: {prompt_4} and `prompt_embeds_llama3`: {prompt_embeds_llama3}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and pooled_prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `pooled_prompt_embeds`. Cannot leave both `prompt` and `pooled_prompt_embeds` undefined." + ) + elif prompt is None and prompt_embeds_t5 is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds_t5`. Cannot leave both `prompt` and `prompt_embeds_t5` undefined." + ) + elif prompt is None and prompt_embeds_llama3 is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds_llama3`. Cannot leave both `prompt` and `prompt_embeds_llama3` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)): + raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}") + elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)): + raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}") + elif prompt_4 is not None and (not isinstance(prompt_4, str) and not isinstance(prompt_4, list)): + raise ValueError(f"`prompt_4` has to be of type `str` or `list` but is {type(prompt_4)}") + + if negative_prompt is not None and negative_pooled_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_pooled_prompt_embeds`:" + f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two." + ) + elif negative_prompt_2 is not None and negative_pooled_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_pooled_prompt_embeds`:" + f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two." + ) + elif negative_prompt_3 is not None and negative_prompt_embeds_t5 is not None: + raise ValueError( + f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} and `negative_prompt_embeds_t5`:" + f" {negative_prompt_embeds_t5}. Please make sure to only forward one of the two." + ) + elif negative_prompt_4 is not None and negative_prompt_embeds_llama3 is not None: + raise ValueError( + f"Cannot forward both `negative_prompt_4`: {negative_prompt_4} and `negative_prompt_embeds_llama3`:" + f" {negative_prompt_embeds_llama3}. Please make sure to only forward one of the two." + ) + + if pooled_prompt_embeds is not None and negative_pooled_prompt_embeds is not None: + if pooled_prompt_embeds.shape != negative_pooled_prompt_embeds.shape: + raise ValueError( + "`pooled_prompt_embeds` and `negative_pooled_prompt_embeds` must have the same shape when passed directly, but" + f" got: `pooled_prompt_embeds` {pooled_prompt_embeds.shape} != `negative_pooled_prompt_embeds`" + f" {negative_pooled_prompt_embeds.shape}." + ) + if prompt_embeds_t5 is not None and negative_prompt_embeds_t5 is not None: + if prompt_embeds_t5.shape != negative_prompt_embeds_t5.shape: + raise ValueError( + "`prompt_embeds_t5` and `negative_prompt_embeds_t5` must have the same shape when passed directly, but" + f" got: `prompt_embeds_t5` {prompt_embeds_t5.shape} != `negative_prompt_embeds_t5`" + f" {negative_prompt_embeds_t5.shape}." + ) + if prompt_embeds_llama3 is not None and negative_prompt_embeds_llama3 is not None: + if prompt_embeds_llama3.shape != negative_prompt_embeds_llama3.shape: + raise ValueError( + "`prompt_embeds_llama3` and `negative_prompt_embeds_llama3` must have the same shape when passed directly, but" + f" got: `prompt_embeds_llama3` {prompt_embeds_llama3.shape} != `negative_prompt_embeds_llama3`" + f" {negative_prompt_embeds_llama3.shape}." + ) + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + ): + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + + shape = (batch_size, num_channels_latents, height, width) + + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + if latents.shape != shape: + raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") + latents = latents.to(device) + return latents + + + def prepare_image_latents( + self, image, batch_size, num_images_per_prompt, dtype, device, do_classifier_free_guidance, generator=None + ): + if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): + raise ValueError( + f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" + ) + + image = image.to(device=device, dtype=dtype) + + batch_size = batch_size * num_images_per_prompt + + if image.shape[1] == 4: + image_latents = image + else: + image_latents = retrieve_latents(self.vae.encode(image), sample_mode="argmax") + image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor + if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: + # expand image_latents for batch_size + deprecation_message = ( + f"You have passed {batch_size} text prompts (`prompt`), but only {image_latents.shape[0]} initial" + " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" + " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" + " your script to pass as many initial images as text prompts to suppress this warning." + ) + deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) + additional_image_per_prompt = batch_size // image_latents.shape[0] + image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0) + elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0: + raise ValueError( + f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts." + ) + else: + image_latents = torch.cat([image_latents], dim=0) + + if do_classifier_free_guidance: + uncond_image_latents = torch.zeros_like(image_latents) + image_latents = torch.cat([uncond_image_latents, image_latents, image_latents], dim=0) + + return image_latents + + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def image_guidance_scale(self): + return self._image_guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + prompt_2: Optional[Union[str, List[str]]] = None, + prompt_3: Optional[Union[str, List[str]]] = None, + prompt_4: Optional[Union[str, List[str]]] = None, + image: PipelineImageInput = None, + num_inference_steps: int = 50, + sigmas: Optional[List[float]] = None, + guidance_scale: float = 5.0, + image_guidance_scale: float = 2.0, + negative_prompt: Optional[Union[str, List[str]]] = None, + negative_prompt_2: Optional[Union[str, List[str]]] = None, + negative_prompt_3: Optional[Union[str, List[str]]] = None, + negative_prompt_4: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds_t5: Optional[torch.FloatTensor] = None, + prompt_embeds_llama3: Optional[torch.FloatTensor] = None, + negative_prompt_embeds_t5: Optional[torch.FloatTensor] = None, + negative_prompt_embeds_llama3: Optional[torch.FloatTensor] = None, + pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + attention_kwargs: Optional[Dict[str, Any]] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 128, + refine_strength: float = 0.0, + reload_keys: Any = None, + **kwargs, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is + will be used instead. + prompt_3 (`str` or `List[str]`, *optional*): + The prompt or prompts to be sent to `tokenizer_3` and `text_encoder_3`. If not defined, `prompt` is + will be used instead. + prompt_4 (`str` or `List[str]`, *optional*): + The prompt or prompts to be sent to `tokenizer_4` and `text_encoder_4`. If not defined, `prompt` is + will be used instead. + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. This is set to 1024 by default for the best results. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. This is set to 1024 by default for the best results. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to 3.5): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is + not greater than `1`). + negative_prompt_2 (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and + `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders. + negative_prompt_3 (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation to be sent to `tokenizer_3` and + `text_encoder_3`. If not defined, `negative_prompt` is used in all the text-encoders. + negative_prompt_4 (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation to be sent to `tokenizer_4` and + `text_encoder_4`. If not defined, `negative_prompt` is used in all the text-encoders. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will ge generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. + If not provided, pooled text embeddings will be generated from `prompt` input argument. + negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt` + input argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple. + attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + callback_on_step_end (`Callable`, *optional*): + A function that calls at the end of each denoising steps during the inference. The function is called + with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, + callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by + `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + max_sequence_length (`int` defaults to 128): Maximum sequence length to use with the `prompt`. + + Examples: + + Returns: + [`~pipelines.hidream_image.HiDreamImagePipelineOutput`] or `tuple`: + [`~pipelines.hidream_image.HiDreamImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When + returning a tuple, the first element is a list with the generated. images. + """ + + prompt_embeds = kwargs.get("prompt_embeds", None) + negative_prompt_embeds = kwargs.get("negative_prompt_embeds", None) + + if prompt_embeds is not None: + deprecation_message = "The `prompt_embeds` argument is deprecated. Please use `prompt_embeds_t5` and `prompt_embeds_llama3` instead." + deprecate("prompt_embeds", "0.35.0", deprecation_message) + prompt_embeds_t5 = prompt_embeds[0] + prompt_embeds_llama3 = prompt_embeds[1] + + if negative_prompt_embeds is not None: + deprecation_message = "The `negative_prompt_embeds` argument is deprecated. Please use `negative_prompt_embeds_t5` and `negative_prompt_embeds_llama3` instead." + deprecate("negative_prompt_embeds", "0.35.0", deprecation_message) + negative_prompt_embeds_t5 = negative_prompt_embeds[0] + negative_prompt_embeds_llama3 = negative_prompt_embeds[1] + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + prompt_2, + prompt_3, + prompt_4, + negative_prompt=negative_prompt, + negative_prompt_2=negative_prompt_2, + negative_prompt_3=negative_prompt_3, + negative_prompt_4=negative_prompt_4, + prompt_embeds_t5=prompt_embeds_t5, + prompt_embeds_llama3=prompt_embeds_llama3, + negative_prompt_embeds_t5=negative_prompt_embeds_t5, + negative_prompt_embeds_llama3=negative_prompt_embeds_llama3, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + ) + + self._guidance_scale = guidance_scale + self._image_guidance_scale = image_guidance_scale + self._attention_kwargs = attention_kwargs + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + elif pooled_prompt_embeds is not None: + batch_size = pooled_prompt_embeds.shape[0] + + device = self._execution_device + + # 3. Encode prompt + lora_scale = self.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None + ( + prompt_embeds_t5, + negative_prompt_embeds_t5, + prompt_embeds_llama3, + negative_prompt_embeds_llama3, + pooled_prompt_embeds, + negative_pooled_prompt_embeds, + ) = self.encode_prompt( + prompt=prompt, + prompt_2=prompt_2, + prompt_3=prompt_3, + prompt_4=prompt_4, + negative_prompt=negative_prompt, + negative_prompt_2=negative_prompt_2, + negative_prompt_3=negative_prompt_3, + negative_prompt_4=negative_prompt_4, + do_classifier_free_guidance=self.do_classifier_free_guidance, + prompt_embeds_t5=prompt_embeds_t5, + prompt_embeds_llama3=prompt_embeds_llama3, + negative_prompt_embeds_t5=negative_prompt_embeds_t5, + negative_prompt_embeds_llama3=negative_prompt_embeds_llama3, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + lora_scale=lora_scale, + ) + + if prompt is not None and "Target Image Description:" in prompt: + target_prompt = prompt.split("Target Image Description:")[1].strip() + ( + target_prompt_embeds_t5, + target_negative_prompt_embeds_t5, + target_prompt_embeds_llama3, + target_negative_prompt_embeds_llama3, + target_pooled_prompt_embeds, + target_negative_pooled_prompt_embeds, + ) = self.encode_prompt( + prompt=target_prompt, + prompt_2=None, + prompt_3=None, + prompt_4=None, + negative_prompt=negative_prompt, + negative_prompt_2=None, + negative_prompt_3=None, + negative_prompt_4=None, + do_classifier_free_guidance=self.do_classifier_free_guidance, + prompt_embeds_t5=None, + prompt_embeds_llama3=None, + negative_prompt_embeds_t5=None, + negative_prompt_embeds_llama3=None, + pooled_prompt_embeds=None, + negative_pooled_prompt_embeds=None, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + lora_scale=lora_scale, + ) + else: + target_prompt_embeds_t5 = prompt_embeds_t5 + target_negative_prompt_embeds_t5 = negative_prompt_embeds_t5 + target_prompt_embeds_llama3 = prompt_embeds_llama3 + target_negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3 + target_pooled_prompt_embeds = pooled_prompt_embeds + target_negative_pooled_prompt_embeds = negative_pooled_prompt_embeds + + image = self.image_processor.preprocess(image) + + image_latents = self.prepare_image_latents( + image, + batch_size, + num_images_per_prompt, + pooled_prompt_embeds.dtype, + device, + self.do_classifier_free_guidance, + ) + + height, width = image_latents.shape[-2:] + height = height * self.vae_scale_factor + width = width * self.vae_scale_factor + + if self.do_classifier_free_guidance: + prompt_embeds_t5 = torch.cat([negative_prompt_embeds_t5, negative_prompt_embeds_t5, prompt_embeds_t5], dim=0) + prompt_embeds_llama3 = torch.cat([negative_prompt_embeds_llama3, negative_prompt_embeds_llama3, prompt_embeds_llama3], dim=1) + pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) + + target_prompt_embeds_t5 = torch.cat([target_negative_prompt_embeds_t5, target_prompt_embeds_t5], dim=0) + target_prompt_embeds_llama3 = torch.cat([target_negative_prompt_embeds_llama3, target_prompt_embeds_llama3], dim=1) + target_pooled_prompt_embeds = torch.cat([target_negative_pooled_prompt_embeds, target_pooled_prompt_embeds], dim=0) + + # 4. Prepare latent variables + num_channels_latents = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + pooled_prompt_embeds.dtype, + device, + generator, + latents, + ) + + # 5. Prepare timesteps + mu = calculate_shift(self.transformer.max_seq) + scheduler_kwargs = {"mu": mu} + if isinstance(self.scheduler, UniPCMultistepScheduler): + self.scheduler.set_timesteps(num_inference_steps, device=device) # , shift=math.exp(mu)) + timesteps = self.scheduler.timesteps + else: + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + **scheduler_kwargs, + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + # 6. Denoising loop + refine_stage = False + if reload_keys is not None: + load_info = self.transformer.load_state_dict(reload_keys['editing'], strict=False) + assert len(load_info.unexpected_keys) == 0 + self.transformer.enable_adapters() + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if reload_keys is not None and i == int(num_inference_steps * (1.0 - refine_strength)): + self.transformer.disable_adapters() + load_info = self.transformer.load_state_dict(reload_keys['refine'], strict=False) + assert len(load_info.unexpected_keys) == 0 + logger.info(f"Refining start at step {i}") + refine_stage = True + if self.interrupt: + continue + if refine_stage: + latent_model_input_with_condition = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + prompt_embeds_t5 = target_prompt_embeds_t5 + prompt_embeds_llama3 = target_prompt_embeds_llama3 + pooled_prompt_embeds = target_pooled_prompt_embeds + else: + latent_model_input = torch.cat([latents] * 3) if self.do_classifier_free_guidance else latents + latent_model_input_with_condition = torch.cat([latent_model_input, image_latents], dim=-1) + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input_with_condition.shape[0]) + noise_pred = self.transformer( + hidden_states=latent_model_input_with_condition, + timesteps=timestep, + encoder_hidden_states_t5=prompt_embeds_t5, + encoder_hidden_states_llama3=prompt_embeds_llama3, + pooled_embeds=pooled_prompt_embeds, + return_dict=False, + )[0] + # perform guidance + if self.do_classifier_free_guidance: + if refine_stage: + uncond, full_cond = noise_pred.chunk(2) + noise_pred = uncond + self.guidance_scale * (full_cond - uncond) + noise_pred = noise_pred[..., :latents.shape[-1]] + else: + uncond, image_cond, full_cond = noise_pred.chunk(3) + noise_pred = uncond + self.image_guidance_scale * (image_cond - uncond) + self.guidance_scale * ( + full_cond - image_cond) + noise_pred = noise_pred[..., :latents.shape[-1]] + + noise_pred = -noise_pred + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + if latents.dtype != latents_dtype: + if torch.backends.mps.is_available(): + # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272 + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds_t5 = callback_outputs.pop("prompt_embeds_t5", prompt_embeds_t5) + prompt_embeds_llama3 = callback_outputs.pop("prompt_embeds_llama3", prompt_embeds_llama3) + pooled_prompt_embeds = callback_outputs.pop("pooled_prompt_embeds", pooled_prompt_embeds) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + if output_type == "latent": + image = latents + + else: + latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor + + image = self.vae.decode(latents, return_dict=False)[0] + image = self.image_processor.postprocess(image, output_type=output_type) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return HiDreamImagePipelineOutput(images=image) diff --git a/modules/infotext.py b/modules/infotext.py index 497879d31..d655be68e 100644 --- a/modules/infotext.py +++ b/modules/infotext.py @@ -68,7 +68,6 @@ def parse(infotext): debug(f'Raw: {infotext}') remaining = infotext.replace('\nSteps:', ' Steps:') - # TODO infotext: handle using regex instead params = [' steps:', ' seed:', ' width:', ' height:', ' sampler:', ' size:', ' cfg scale:'] # first param is one of those params += ['\nsteps:', '\nseed:', '\nwidth:', '\nheight:', '\nsampler:', '\nsize:', '\ncfg scale:'] params += ['.steps:', '.seed:', '.width:', '.height:', '.sampler:', '.size:', '.cfg scale:'] diff --git a/modules/intel/ipex/diffusers.py b/modules/intel/ipex/diffusers.py index 8a10aa546..fa725f2ad 100644 --- a/modules/intel/ipex/diffusers.py +++ b/modules/intel/ipex/diffusers.py @@ -1,6 +1,7 @@ from functools import wraps import torch import diffusers # pylint: disable=import-error +from diffusers.utils import torch_utils # pylint: disable=import-error, unused-import # noqa: F401 # pylint: disable=protected-access, missing-function-docstring, line-too-long @@ -61,13 +62,11 @@ def hidream_rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor: def ipex_diffusers(device_supports_fp64=False, can_allocate_plus_4gb=False): - # get around lazy imports - from diffusers.utils import torch_utils # pylint: disable=import-error, unused-import diffusers.utils.torch_utils.fourier_filter = fourier_filter if not device_supports_fp64: # get around lazy imports - from diffusers.models import transformers as diffusers_transformers # pylint: disable=import-error, unused-import - from diffusers.models import controlnets as diffusers_controlnets # pylint: disable=import-error, unused-import + from diffusers.models import transformers as diffusers_transformers # pylint: disable=import-error, unused-import # noqa: F401 + from diffusers.models import controlnets as diffusers_controlnets # pylint: disable=import-error, unused-import # noqa: F401 diffusers.models.embeddings.FluxPosEmbed = FluxPosEmbed diffusers.models.transformers.transformer_flux.FluxPosEmbed = FluxPosEmbed diffusers.models.controlnets.controlnet_flux.FluxPosEmbed = FluxPosEmbed diff --git a/modules/interrogate/openclip.py b/modules/interrogate/openclip.py index 792b4df85..554678642 100644 --- a/modules/interrogate/openclip.py +++ b/modules/interrogate/openclip.py @@ -350,7 +350,7 @@ def interrogate_batch(batch_files, batch_folder, batch_str, clip_model, blip_mod files += [f.name for f in batch_folder] if batch_str is not None and len(batch_str) > 0 and os.path.exists(batch_str) and os.path.isdir(batch_str): from modules.files_cache import list_files - files += list(list_files(batch_str, ext_filter=['.png', '.jpg', '.jpeg', '.webp'], recursive=recursive)) + files += list(list_files(batch_str, ext_filter=['.png', '.jpg', '.jpeg', '.webp', '.jxl'], recursive=recursive)) if len(files) == 0: shared.log.warning('Interrogate batch: type=clip no images') return '' diff --git a/modules/interrogate/vqa.py b/modules/interrogate/vqa.py index b7f6aabb3..461ec6c73 100644 --- a/modules/interrogate/vqa.py +++ b/modules/interrogate/vqa.py @@ -27,6 +27,7 @@ vlm_models = { "Google Pix Textcaps": "google/pix2struct-textcaps-base", # 1.1GB "Google PaliGemma 2 3B": "google/paligemma2-3b-pt-224", "Alibaba Qwen VL2 2B": "Qwen/Qwen2-VL-2B-Instruct", + "Alibaba Qwen 2.5 Omni 3B": "Qwen/Qwen2.5-Omni-3B", "Huggingface Smol VL2 0.5B": "HuggingFaceTB/SmolVLM-500M-Instruct", "Huggingface Smol VL2 2B": "HuggingFaceTB/SmolVLM-Instruct", "Salesforce BLIP Base": "Salesforce/blip-vqa-base", # 1.5GB @@ -67,6 +68,8 @@ vlm_prompts = [ def b64(image): + if image is None: + return '' with io.BytesIO() as stream: image.save(stream, 'JPEG') values = stream.getvalue() @@ -125,7 +128,7 @@ def qwen(question: str, image: Image.Image, repo: str = None, system_prompt: str devices.torch_gc() sd_models.move_model(model, devices.device) question = question.replace('<', '').replace('>', '').replace('_', ' ') - system_prompt = system_prompt or shared.opts.vlm_system + system_prompt = system_prompt or shared.opts.interrogate_vlm_system conversation = [ { "role": "system", @@ -169,19 +172,20 @@ def gemma(question: str, image: Image.Image, repo: str = None, system_prompt: st devices.torch_gc() sd_models.move_model(model, devices.device) question = question.replace('<', '').replace('>', '').replace('_', ' ') - system_prompt = system_prompt or shared.opts.vlm_system + system_prompt = system_prompt or shared.opts.interrogate_vlm_system + + system_content = [] + if system_prompt is not None and len(system_prompt) > 4: + system_content.append({"type": "text", "text": system_prompt}) + + user_content = [] + if image is not None: + user_content.append({"type": "image", "image": b64(image)}) + if question is not None and len(question) > 4: + user_content.append({"type": "text", "text": question}) conversation = [ - { - "role": "system", - "content": [{"type": "text", "text": system_prompt}] - }, - { - "role": "user", - "content": [ - {"type": "image", "image": b64(image)}, - {"type": "text", "text": question} - ] - } + { "role": "system", "content": system_content}, + { "role": "user", "content": user_content }, ] inputs = processor.apply_chat_template( conversation, @@ -294,7 +298,7 @@ def smol(question: str, image: Image.Image, repo: str = None, system_prompt: str devices.torch_gc() sd_models.move_model(model, devices.device) question = question.replace('<', '').replace('>', '').replace('_', ' ') - system_prompt = system_prompt or shared.opts.vlm_system + system_prompt = system_prompt or shared.opts.interrogate_vlm_system conversation = [ { "role": "system", @@ -632,7 +636,7 @@ def batch(model_name, system_prompt, batch_files, batch_folder, batch_str, quest files += [f.name for f in batch_folder] if batch_str is not None and len(batch_str) > 0 and os.path.exists(batch_str) and os.path.isdir(batch_str): from modules.files_cache import list_files - files += list(list_files(batch_str, ext_filter=['.png', '.jpg', '.jpeg', '.webp'], recursive=recursive)) + files += list(list_files(batch_str, ext_filter=['.png', '.jpg', '.jpeg', '.webp', '.jxl'], recursive=recursive)) if len(files) == 0: shared.log.warning('Interrogate batch: type=vlm no images') return '' diff --git a/modules/model_flex.py b/modules/model_flex.py index 9ef5f44e1..953a22ec1 100644 --- a/modules/model_flex.py +++ b/modules/model_flex.py @@ -82,8 +82,8 @@ def load_flex(checkpoint_info, diffusers_load_config={}): ) sd_hijack_te.init_hijack(pipe) diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["flex2"] = Flex2Pipeline - diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["fluxcfgzero"] = Flex2Pipeline - diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["fluxcfgzero"] = Flex2Pipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["flex2"] = Flex2Pipeline + diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["flex2"] = Flex2Pipeline del text_encoder_2 del transformer diff --git a/modules/model_hidream.py b/modules/model_hidream.py index 138952680..bd0159130 100644 --- a/modules/model_hidream.py +++ b/modules/model_hidream.py @@ -42,6 +42,8 @@ def load_transformer(repo_id, diffusers_load_config={}): def load_text_encoders(repo_id, diffusers_load_config={}): + if repo_id == 'HiDream-ai/HiDream-E1-Full': + repo_id = 'HiDream-ai/HiDream-I1-Full' # use I1 for t5 and llm load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) shared.log.debug(f'Load model: type=HiDream te3="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') text_encoder_3 = transformers.T5EncoderModel.from_pretrained( @@ -92,7 +94,19 @@ def load_hidream(checkpoint_info, diffusers_load_config={}): load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model') shared.log.debug(f'Load model: type=HiDream model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - pipe = diffusers.HiDreamImagePipeline.from_pretrained( + if 'I1' in repo_id: + cls = diffusers.HiDreamImagePipeline + elif 'E1' in repo_id: + from modules.hidream.pipeline_hidream_image_editing import HiDreamImageEditingPipeline + cls = HiDreamImageEditingPipeline + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["hidream-e1"] = diffusers.HiDreamImagePipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["hidream-e1"] = HiDreamImageEditingPipeline + diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["hidream-e1"] = HiDreamImageEditingPipeline + else: + shared.log.error(f'Load model: type=HiDream model="{checkpoint_info.name}" repo="{repo_id}" not recognized') + return False + + pipe = cls.from_pretrained( repo_id, transformer=transformer, text_encoder_3=text_encoder_3, @@ -101,6 +115,7 @@ def load_hidream(checkpoint_info, diffusers_load_config={}): cache_dir=shared.opts.diffusers_dir, **load_args, ) + sd_hijack_te.init_hijack(pipe) del text_encoder_3 del text_encoder_4 diff --git a/modules/model_quant.py b/modules/model_quant.py index 5538a9dd8..95a106c6f 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -207,7 +207,7 @@ def load_bnb(msg='', silent=False): if not installed('bitsandbytes'): if devices.backend == 'cuda': # forcing a version will uninstall the multi-backend-refactor branch of bnb - install('bitsandbytes==0.45.1', quiet=True) + install('bitsandbytes==0.45.5', quiet=True) log.warning('Quantization: bitsandbytes installed please restart') try: import bitsandbytes @@ -263,6 +263,8 @@ def load_nncf(msg='', silent=False): if not installed('nncf'): install('nncf==2.16.0', quiet=True) log.warning('Quantization: nncf installed please restart') + install('jstyleson', quiet=True) + install('texttable', quiet=True) try: import nncf intel_nncf = nncf @@ -350,6 +352,7 @@ def nncf_compress_model(model, op=None, sd_model=None, send_to_device=True, do_g num_bits = 8 if shared.opts.nncf_compress_weights_mode in {"INT8", "INT8_SYM", "INT8_ASYM"} else 4 is_asym_mode = shared.opts.nncf_compress_weights_mode in {"INT8", "INT4", "INT8_ASYM", "INT4_ASYM"} model = apply_nncf_to_module(model, num_bits, is_asym_mode, quant_conv=shared.opts.nncf_quantize_conv_layers) + model.quantization_method = 'NNCF' if send_to_device: nncf_send_to_device(model, devices.device) diff --git a/modules/modelloader.py b/modules/modelloader.py index 7bbb4f341..15ecb7dbe 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -26,7 +26,7 @@ def hf_login(token=None): global loggedin # pylint: disable=global-statement token = token or shared.opts.huggingface_token install('hf_xet', quiet=True) - if token is None or len(token) <= 2: + if token is None or len(token) <= 4: log.debug('HF login: no token provided') return False if os.environ.get('HUGGING_FACE_HUB_TOKEN', None) is not None: @@ -41,8 +41,9 @@ def hf_login(token=None): hf.logout() hf.login(token=token, add_to_git_credential=False, write_permission=False) text = stdout.getvalue() or '' + obfuscated_token = 'hf_...' + token[-4:] line = [l for l in text.split('\n') if 'Token' in l] - log.info(f'HF login: token="{hf.constants.HF_TOKEN_PATH}" {line[0] if len(line) > 0 else text}') + log.info(f'HF login: token="{obfuscated_token}" fn="{hf.constants.HF_TOKEN_PATH}" {line[0] if len(line) > 0 else text}') loggedin = token return True diff --git a/modules/processing_args.py b/modules/processing_args.py index 016e0bcbe..6ef43667b 100644 --- a/modules/processing_args.py +++ b/modules/processing_args.py @@ -35,7 +35,10 @@ def task_specific_kwargs(p, model): } elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.IMAGE_2_IMAGE or is_img2img_model) and len(getattr(p, 'init_images', [])) > 0: if shared.sd_model_type == 'sdxl' and hasattr(model, 'register_to_config'): - model.register_to_config(requires_aesthetics_score = False) + if model.__class__.__name__ in sd_models.i2i_pipes: + pass + else: + model.register_to_config(requires_aesthetics_score = False) if 'hires' not in p.ops: p.ops.append('img2img') if p.vae_type == 'Remote': @@ -71,7 +74,10 @@ def task_specific_kwargs(p, model): } elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INPAINTING or is_img2img_model) and len(getattr(p, 'init_images', [])) > 0: if shared.sd_model_type == 'sdxl' and hasattr(model, 'register_to_config'): - model.register_to_config(requires_aesthetics_score = False) + if model.__class__.__name__ in [sd_models.i2i_pipes]: + pass + else: + model.register_to_config(requires_aesthetics_score = False) if p.detailer_enabled: p.ops.append('detailer') else: @@ -144,7 +150,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t 'StableDiffusion' in model.__class__.__name__ or 'StableCascade' in model.__class__.__name__ or 'Flux' in model.__class__.__name__ or - 'HiDreamImage' in model.__class__.__name__ + 'HiDreamImagePipeline' in model.__class__.__name__ # hidream-e1 has different embeds ): try: prompt_parser_diffusers.embedder = prompt_parser_diffusers.PromptEmbedder(prompts, negative_prompts, steps, clip_skip, p) @@ -161,7 +167,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t if 'prompt' in possible: if 'OmniGen' in model.__class__.__name__: prompts = [p.replace('|image|', '<|image_1|>') for p in prompts] - if 'HiDreamImage' in model.__class__.__name__: + if 'HiDreamImage' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None: args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds') prompt_embeds = prompt_parser_diffusers.embedder('prompt_embeds') args['prompt_embeds_t5'] = prompt_embeds[0] @@ -180,7 +186,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t else: args['prompt'] = prompts if 'negative_prompt' in possible: - if 'HiDreamImage' in model.__class__.__name__: + if 'HiDreamImage' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None: args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds') negative_prompt_embeds = prompt_parser_diffusers.embedder('negative_prompt_embeds') args['negative_prompt_embeds_t5'] = negative_prompt_embeds[0] diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 2c1f16cb6..03d0b7b78 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -54,7 +54,7 @@ def restore_state(p: processing.StableDiffusionProcessing): def process_base(p: processing.StableDiffusionProcessing): txt2img = is_txt2img() - use_refiner_start = txt2img and is_refiner_enabled(p) and not p.is_hr_pass and p.refiner_start > 0 and p.refiner_start < 1 + use_refiner_start = is_refiner_enabled(p) and (not p.is_hr_pass) use_denoise_start = not txt2img and p.refiner_start > 0 and p.refiner_start < 1 shared.sd_model = update_pipeline(shared.sd_model, p) diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index 7d417525b..2535127ff 100644 --- a/modules/processing_helpers.py +++ b/modules/processing_helpers.py @@ -23,7 +23,7 @@ def is_txt2img(): def is_refiner_enabled(p): - return p.enable_hr and p.refiner_steps > 0 and p.refiner_start > 0 and p.refiner_start < 1 and shared.sd_refiner is not None + return p.enable_hr and (p.refiner_steps > 0) and (p.refiner_start > 0) and (p.refiner_start < 1) and (shared.sd_refiner is not None) def setup_color_correction(image): @@ -454,12 +454,15 @@ def calculate_base_steps(p, use_denoise_start, use_refiner_start): if len(getattr(p, 'timesteps', [])) > 0: return None if not is_txt2img(): - if use_denoise_start and shared.sd_model_type == 'sdxl': + cls = shared.sd_model.__class__.__name__ + if cls in sd_models.i2i_pipes: + steps = p.steps + elif 'Flex' in cls: + steps = p.steps + elif 'HiDreamImageEditingPipeline' in cls: + steps = p.steps + elif use_denoise_start and (shared.sd_model_type == 'sdxl'): steps = p.steps // (1 - p.refiner_start) - elif 'Flex' in shared.sd_model.__class__.__name__: - steps = p.steps - elif shared.sd_model_type == 'omnigen': - steps = p.steps elif p.denoising_strength > 0: steps = (p.steps // p.denoising_strength) + 1 else: @@ -536,19 +539,17 @@ def set_latents(p): return latents -last_circular = False -def apply_circular(enable, model): - global last_circular # pylint: disable=global-statement +def apply_circular(enable: bool, model): if not hasattr(model, 'unet') or not hasattr(model, 'vae'): return - if last_circular == enable: + if getattr(model, 'texture_tiling', False) == enable: return try: for layer in [layer for layer in model.unet.modules() if type(layer) is torch.nn.Conv2d]: layer.padding_mode = 'circular' if enable else 'zeros' for layer in [layer for layer in model.vae.modules() if type(layer) is torch.nn.Conv2d]: layer.padding_mode = 'circular' if enable else 'zeros' - last_circular = enable + model.texture_tiling = enable except Exception as e: debug(f"Diffusers tiling failed: {e}") diff --git a/modules/sd_checkpoint.py b/modules/sd_checkpoint.py index f2d7f4f7b..52d96dedc 100644 --- a/modules/sd_checkpoint.py +++ b/modules/sd_checkpoint.py @@ -158,6 +158,7 @@ def list_models(): shared.log.info(f'Available Models: safetensors="{shared.opts.ckpt_dir}":{len(safetensors_list)} diffusers="{shared.opts.diffusers_dir}":{len(diffusers_list)} items={len(checkpoints_list)} time={time.time()-t0:.2f}') checkpoints_list = dict(sorted(checkpoints_list.items(), key=lambda cp: cp[1].filename)) + def update_model_hashes(): txt = [] lst = [ckpt for ckpt in checkpoints_list.values() if ckpt.hash is None] diff --git a/modules/sd_detect.py b/modules/sd_detect.py index 4a40453cd..8a413757c 100644 --- a/modules/sd_detect.py +++ b/modules/sd_detect.py @@ -104,12 +104,14 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False): index = shared.readfile(index, silent=True) cls = index.get('_class_name', None) if cls is not None: - pipeline = getattr(diffusers, cls) - if 'Flux' in pipeline.__name__ and guess != 'FLEX': + pipeline = getattr(diffusers, cls, None) + if pipeline is None: + pipeline = cls + if callable(pipeline) and 'Flux' in pipeline.__name__ and guess != 'FLEX': guess = 'FLUX' - if 'StableDiffusion3' in pipeline.__name__: + if callable(pipeline) and 'StableDiffusion3' in pipeline.__name__: guess = 'Stable Diffusion 3' - if 'Lumina2' in pipeline.__name__: + if callable(pipeline) and 'Lumina2' in pipeline.__name__: guess = 'Lumina 2' # switch for specific variant if guess == 'Stable Diffusion' and 'inpaint' in f.lower(): diff --git a/modules/sd_models.py b/modules/sd_models.py index 4e7664d6d..10a34ad44 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -31,19 +31,27 @@ debug_process = log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not Non diffusers_version = int(diffusers.__version__.split('.')[1]) checkpoint_tiles = checkpoint_titles # legacy compatibility pipe_switch_task_exclude = [ - 'StableDiffusionReferencePipeline', - 'StableDiffusionAdapterPipeline', - 'AnimateDiffPipeline', - 'AnimateDiffSDXLPipeline', - 'OmniGenPipeline', - 'StableDiffusion3ControlNetPipeline', - 'InstantIRPipeline', - 'FluxFillPipeline', + 'AnimateDiffPipeline', 'AnimateDiffSDXLPipeline', 'FluxControlPipeline', - 'PixelSmithXLPipeline', - 'PhotoMakerStableDiffusionXLPipeline', - 'StableDiffusionXLInstantIDPipeline', + 'FluxFillPipeline', + 'InstantIRPipeline', 'LTXConditionPipeline', + 'OmniGenPipeline', + 'PhotoMakerStableDiffusionXLPipeline', + 'PixelSmithXLPipeline', + 'StableDiffusion3ControlNetPipeline', + 'StableDiffusionAdapterPipeline', + 'StableDiffusionAdapterPipeline', 'StableDiffusionXLAdapterPipeline', + 'StableDiffusionControlNetXSPipeline', 'StableDiffusionXLControlNetXSPipeline', + 'StableDiffusionReferencePipeline', + 'StableDiffusionXLInstantIDPipeline', +] +i2i_pipes = [ + 'LEditsPPPipelineStableDiffusion', + 'LEditsPPPipelineStableDiffusionXL', + 'OmniGenPipeline', + 'StableDiffusionAdapterPipeline', 'StableDiffusionXLAdapterPipeline', + 'StableDiffusionControlNetXSPipeline', 'StableDiffusionXLControlNetXSPipeline', ] @@ -276,7 +284,7 @@ def load_diffuser_initial(diffusers_load_config, op='model'): def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op='model'): sd_model = None - unload_model_weights() + unload_model_weights(op=op) shared.sd_model = None try: if model_type in ['Stable Cascade']: # forced pipeline @@ -655,7 +663,7 @@ class DiffusersTaskType(Enum): def get_diffusers_task(pipe: diffusers.DiffusionPipeline) -> DiffusersTaskType: cls = pipe.__class__.__name__ - if cls in ["LEditsPPPipelineStableDiffusion", "LEditsPPPipelineStableDiffusionXL", "OmniGenPipeline"]: # special case + if cls in i2i_pipes: # special case return DiffusersTaskType.IMAGE_2_IMAGE elif 'ImageToVideo' in cls or cls in ['LTXConditionPipeline', 'StableVideoDiffusionPipeline']: # i2v pipelines return DiffusersTaskType.IMAGE_2_IMAGE @@ -782,6 +790,9 @@ def set_diffuser_pipe(pipe, new_pipe_type): if new_pipe_type == DiffusersTaskType.TEXT_2_IMAGE: clean_diffuser_pipe(pipe) + if hasattr(pipe, 'no_task_switch'): + del pipe.no_task_switch + return pipe if get_diffusers_task(pipe) == new_pipe_type: return pipe diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py index 3d83cb29d..ca5fa21fd 100644 --- a/modules/sd_samplers_diffusers.py +++ b/modules/sd_samplers_diffusers.py @@ -232,6 +232,8 @@ class DiffusionSampler: self.config['beta_schedule'] = 'scaled_linear' elif shared.opts.schedulers_beta_schedule == 'cosine': self.config['beta_schedule'] = 'squaredcos_cap_v2' + elif shared.opts.schedulers_beta_schedule == 'sigmoid': + self.config['beta_schedule'] = 'sigmoid' timesteps = re.split(',| ', shared.opts.schedulers_timesteps) timesteps = [int(x) for x in timesteps if x.isdigit()] diff --git a/modules/sd_vae_remote.py b/modules/sd_vae_remote.py index c751bc083..741d349bc 100644 --- a/modules/sd_vae_remote.py +++ b/modules/sd_vae_remote.py @@ -28,13 +28,6 @@ dtypes = { } -def h1_pack_latents(latents, _batch_size, _num_channels_latents, _height, _width): # TODO hidream: pack latents for remote vae - # latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) - # latents = latents.permute(0, 2, 4, 1, 3, 5) - # latents = latents.reshape(batch_size, (height // 2) * (width // 2) // (num_channels_latents * 4), num_channels_latents * 4) - return latents - - def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_type: str = None) -> Image.Image: from modules import devices, shared, errors, modelloader tensors = [] @@ -57,9 +50,6 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_ latent = latent_copy[i] if model_type != 'f1': latent = latent.unsqueeze(0) - # if model_type == 'h1': - # num_channels_latents = shared.sd_model.transformer.config.in_channels - # latent = h1_pack_latents(latent, 1, num_channels_latents, height, width) # pylint: disable=protected-access params = { "input_tensor_type": "binary", "shape": list(latent.shape), diff --git a/modules/shared.py b/modules/shared.py index 611a96831..07e81d7be 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -814,7 +814,7 @@ options_templates.update(options_section(('sampler-params', "Sampler Settings"), "schedulers_sigma": OptionInfo("default", "Sigma algorithm", gr.Radio, {"choices": ['default', 'karras', 'exponential', 'polyexponential'], "visible": False}), # managed from ui.py for backend diffusers - "schedulers_beta_schedule": OptionInfo("default", "Beta schedule", gr.Dropdown, {"choices": ['default', 'linear', 'scaled_linear', 'squaredcos_cap_v2'], "visible": False}), + "schedulers_beta_schedule": OptionInfo("default", "Beta schedule", gr.Dropdown, {"choices": ['default', 'linear', 'scaled_linear', 'squaredcos_cap_v2', 'sigmoid'], "visible": False}), "schedulers_use_thresholding": OptionInfo(False, "Use dynamic thresholding", gr.Checkbox, {"visible": False}), "schedulers_timestep_spacing": OptionInfo("default", "Timestep spacing", gr.Dropdown, {"choices": ['default', 'linspace', 'leading', 'trailing'], "visible": False}), 'schedulers_timesteps': OptionInfo('', "Timesteps", gr.Textbox, {"visible": False}), @@ -973,7 +973,6 @@ options_templates.update(options_section(('extra_networks', "Networks"), { "extra_networks_styles_sep": OptionInfo("

Styles

", "", gr.HTML), "extra_networks_styles": OptionInfo(True, "Show reference styles"), - "extra_networks_save_unparsed": OptionInfo(True, "Save unparsed prompt"), "extra_networks_apply_unparsed": OptionInfo(True, "Restore unparsed prompt"), "extra_networks_embed_sep": OptionInfo("

Embeddings

", "", gr.HTML), diff --git a/modules/shared_items.py b/modules/shared_items.py index 1bcb47334..d8796daf5 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -17,6 +17,7 @@ pipelines = { 'Stable Diffusion XL': getattr(diffusers, 'StableDiffusionXLPipeline', None), 'Stable Diffusion XL Inpaint': getattr(diffusers, 'StableDiffusionXLInpaintPipeline', None), 'Stable Diffusion XL Instruct': getattr(diffusers, 'StableDiffusionXLInstructPix2PixPipeline', None), + 'Stable Diffusion XL Refiner': getattr(diffusers, 'StableDiffusionXLImg2ImgPipeline', None), 'Stable Cascade': getattr(diffusers, 'StableCascadeCombinedPipeline', None), 'Stable Diffusion 3.x': getattr(diffusers, 'StableDiffusion3Pipeline', None), 'Latent Consistency Model': getattr(diffusers, 'LatentConsistencyModelPipeline', None), diff --git a/modules/styles.py b/modules/styles.py index 5bda3728b..1369315c5 100644 --- a/modules/styles.py +++ b/modules/styles.py @@ -152,11 +152,6 @@ def apply_styles_to_extra(p, style: Style): reference_style = get_reference_style() extra = infotext.parse(reference_style) if shared.opts.extra_network_reference_values else {} - if not hasattr(p, 'original_prompt'): - p.original_prompt = p.prompt - if not hasattr(p, 'original_negative'): - p.original_negative = p.negative_prompt - style_extra = apply_wildcards_to_prompt(style.extra, [style.wildcards], silent=True) style_extra = ' ' + style_extra.lower() extra.update(infotext.parse(style_extra)) @@ -340,6 +335,11 @@ class StyleDatabase: return prompt def apply_styles_to_extra(self, p): + if len(getattr(p, 'original_prompt', '')) == 0: + p.original_prompt = p.prompt + if len(getattr(p, 'original_negative', '')) == 0: + p.original_negative = p.negative_prompt + if p.styles is None: return if p.styles is None or not isinstance(p.styles, list): diff --git a/modules/ui_control.py b/modules/ui_control.py index b10f2838e..9e3de1cf1 100644 --- a/modules/ui_control.py +++ b/modules/ui_control.py @@ -359,7 +359,7 @@ def create_ui(_blocks: gr.Blocks=None): controlnetxs_ui_units.append(unit_ui) units.append(unit.Unit( unit_type = 'xs', - index = 1, + index = i, enabled = enabled, result_txt = result_txt, enabled_cb = enabled_cb, diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index 023eacbad..e8367547a 100644 --- a/modules/ui_extra_networks.py +++ b/modules/ui_extra_networks.py @@ -783,8 +783,10 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): def show_details(text, img, desc, info, meta, description, prompt, negative, parameters, wildcards, params, _dummy1=None, _dummy2=None): page, item = get_item(state, params) - valid = item is not None and hasattr(item, 'name') and hasattr(item, 'filename') - if valid: + is_style = (page is not None) and (page.title == 'Style') + is_valid = (item is not None) and hasattr(item, 'name') and hasattr(item, 'filename') + + if is_valid: stat = os.stat(item.filename) if os.path.exists(item.filename) else None desc = item.description fullinfo = shared.readfile(os.path.splitext(item.filename)[0] + '.json', silent=True) @@ -795,16 +797,14 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): item.filename = None shared.log.warning('Network: show details not supported for compound item') info = None - """ - if prompt is not None: + if prompt is not None and len(prompt) > 0: item.prompt = prompt - if negative is not None: + if negative is not None and len(negative) > 0: item.negative = negative - if description is not None: + if description is not None and len(description) > 0: item.description = description - if wildcards is not None: + if wildcards is not None and len(wildcards) > 0: item.wildcards = wildcards - """ meta = page.metadata.get(item.name, {}) or {} if type(meta) is str: try: @@ -887,7 +887,6 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): {note} ''' - is_style = (page is not None) and (page.title == 'Style') return [ text, # gr.html img, # gr.image @@ -899,7 +898,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): gr.update(value=negative, visible=is_style), # gr.textbox gr.update(value=parameters, visible=is_style), # gr.textbox gr.update(value=wildcards, visible=is_style), # gr.textbox - gr.update(visible=valid), # details ui visible + gr.update(visible=is_valid), # details ui visible gr.update(visible=not is_style), # details ui tabs visible gr.update(visible=is_style), # details ui text visible ] @@ -940,10 +939,9 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): return ui_refresh_click(title) def ui_save_click(): - if shared.opts.extra_networks_save_unparsed: - from modules.processing_info import get_last_args - params, text = get_last_args() - else: + from modules.processing_info import get_last_args + params, text = get_last_args() + if (not params) or (not text) or (len(text) == 0): filename = os.path.join(paths.data_path, "params.txt") if os.path.exists(filename): with open(filename, "r", encoding="utf8") as file: @@ -957,13 +955,12 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): return res def ui_quicksave_click(name): - if shared.opts.extra_networks_save_unparsed: - from modules.processing_info import get_last_args - params, text = get_last_args() - else: - if name is None or len(name) < 1: - shared.log.warning("Network quick save style: no name provided") - return + if name is None or len(name) < 1: + shared.log.warning("Network quick save style: no name provided") + return + from modules.processing_info import get_last_args + params, text = get_last_args() + if (not params) or (not text) or (len(text) == 0): fn = os.path.join(paths.data_path, "params.txt") if os.path.exists(fn): with open(fn, "r", encoding="utf8") as file: @@ -983,9 +980,9 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): } shared.writefile(item, fn, silent=True) if len(prompt) > 0: - shared.log.debug(f'Network quick save style: item="{name}" filename="{fn}" unparsed={shared.opts.extra_networks_unparsed}') + shared.log.debug(f'Networks type=style quicksave style: item="{name}" filename="{fn}" prompt="{prompt}"') else: - shared.log.warning(f'Network quick save model: item="{name}" filename="{fn}" prompt is empty') + shared.log.warning(f'Networks type=style quicksave model: item="{name}" filename="{fn}" prompt is empty') def ui_sort_cards(sort_order): if shared.opts.extra_networks_sort != sort_order: diff --git a/modules/ui_sections.py b/modules/ui_sections.py index 7e464130e..f516147da 100644 --- a/modules/ui_sections.py +++ b/modules/ui_sections.py @@ -307,7 +307,7 @@ def create_sampler_options(tabname): sampler_sigma = gr.Dropdown(label='Sigma method', elem_id=f"{tabname}_sampler_sigma", choices=['default', 'karras', 'betas', 'exponential', 'lambdas', 'flowmatch'], value=shared.opts.schedulers_sigma, type='value') sampler_spacing = gr.Dropdown(label='Timestep spacing', elem_id=f"{tabname}_sampler_spacing", choices=['default', 'linspace', 'leading', 'trailing'], value=shared.opts.schedulers_timestep_spacing, type='value') with gr.Row(elem_classes=['flex-break']): - sampler_beta = gr.Dropdown(label='Beta schedule', elem_id=f"{tabname}_sampler_beta", choices=['default', 'linear', 'scaled', 'cosine'], value=shared.opts.schedulers_beta_schedule, type='value') + sampler_beta = gr.Dropdown(label='Beta schedule', elem_id=f"{tabname}_sampler_beta", choices=['default', 'linear', 'scaled', 'cosine', 'sigmoid'], value=shared.opts.schedulers_beta_schedule, type='value') sampler_prediction = gr.Dropdown(label='Prediction method', elem_id=f"{tabname}_sampler_prediction", choices=['default', 'epsilon', 'sample', 'v_prediction', 'flow_prediction'], value=shared.opts.schedulers_prediction_type, type='value') with gr.Row(elem_classes=['flex-break']): sampler_presets = gr.Dropdown(label='Timesteps presets', elem_id=f"{tabname}_sampler_presets", choices=['None', 'AYS SD15', 'AYS SDXL'], value='None', type='value') diff --git a/modules/ui_video_vlm.py b/modules/ui_video_vlm.py new file mode 100644 index 000000000..b311a4e0b --- /dev/null +++ b/modules/ui_video_vlm.py @@ -0,0 +1,66 @@ +import gradio as gr +from PIL import Image +from modules import shared + + +models = [ + 'Google Gemma 3 4B', +] +system_prompts = { + 'prefix': 'You are a video director.', + 'desc': 'You will provide a prompt that can be used to create short, creative and dynamic video. Describe the subject of the video, add movement that fits the overall tone of the video, add camera movement.', + 'suffix': 'NSFW content and nudity are allowed and must be fully described if present. Movement should be dynamic and creative. Do not specify duration and assume video is short. Avoid slow-motion and prefer faster movements. Output should be a single short paragraph without explanations', + 'example': 'Example: "Short video of beautiful blonde woman in her 20ies wearing a long flowing red dress. She is briskly walking on the beach during sunset and performing a pirouette ending with her hand pointing at the camera as she smiles. Camera is moving around her and zooming to her face. Sun is setting in the background causing changes in colors and shadows to move dynamically."', + + 't2v-prompt': 'You are a given short prompt with basic instructions.', + 't2v-noprompt': '', + 'i2v-prompt': 'You are given an image as a starting point and a short prompt with basic instructions.', + 'i2v-noprompt': 'You are given an image as a starting point.', +} + + +def enhance_prompt(enable:bool, model:str=None, image=None, prompt:str='', system_prompt:str=''): + from modules.interrogate import vqa + if not enable: + return prompt + if model is None or len(model) < 4: + model = models[0] + if image is not None and not isinstance(image, Image.Image): + image = Image.fromarray(image) + if prompt is None or len(prompt) < 4: + prompt = ' ' + if system_prompt is None or len(system_prompt) < 4: + if image is not None: + if prompt is not None and len(prompt) > 4: + core_prompt = system_prompts['i2v-prompt'] + else: + core_prompt = system_prompts['i2v-noprompt'] + else: + if prompt is not None and len(prompt) > 4: + core_prompt = system_prompts['t2v-prompt'] + else: + core_prompt = system_prompts['t2v-noprompt'] + system_prompt = f"{system_prompts['prefix']} {core_prompt} {system_prompts['desc']} {system_prompts['suffix']} {system_prompts['example']}" + shared.log.debug(f'Video prompt enhance: model="{model}" image={image} prompt="{prompt}"') + # shared.log.trace(f'Video prompt enhance: system="{system_prompt}"') + answer = vqa.interrogate(question='', prompt=prompt, system_prompt=system_prompt, image=image, model_name=model, quiet=False) + shared.log.debug(f'Video prompt enhance: answer="{answer}"') + return answer + + +def create_ui(prompt_element:gr.Textbox, image_element:gr.Image): + with gr.Accordion('Prompt enhance', open=False): + with gr.Row(): + enable = gr.Checkbox(label='Enable', value=False) + btn_enhance = gr.Button(value='Enhance now', elem_id='btn_enhance') + with gr.Row(): + model = gr.Dropdown(label='Model', choices=models, value=models[0]) + with gr.Row(): + system_prompt = gr.Textbox(label='System prompt', placeholder='override system prompt with user-provided prompt', lines=3) + btn_enhance.click( + fn=enhance_prompt, + inputs=[enable, model, image_element, prompt_element, system_prompt], + outputs=prompt_element, + show_progress=True, + ) + return enable, model, system_prompt diff --git a/modules/zluda_hijacks.py b/modules/zluda_hijacks.py index 4be4b71d0..f8d831986 100644 --- a/modules/zluda_hijacks.py +++ b/modules/zluda_hijacks.py @@ -33,16 +33,8 @@ MEM_BUS_WIDTH = { } -_topk = torch.topk -def topk(input: torch.Tensor, *args, **kwargs): # pylint: disable=redefined-builtin - device = input.device - values, indices = _topk(input.cpu(), *args, **kwargs) - return torch.return_types.topk((values.to(device), indices.to(device),)) - - class DeviceProperties: PROPERTIES_OVERRIDE = { - "regs_per_multiprocessor": 65535, # sometimes gcnArchName contains device name ("AMD Radeon RX ..."), not architecture name ("gfx...") "gcnArchName": "UNKNOWN ARCHITECTURE", } @@ -68,7 +60,6 @@ def torch__C__cuda_getCurrentRawStream(device): def do_hijack(): - torch.topk = topk if zluda.default_agent is not None: DeviceProperties.PROPERTIES_OVERRIDE["gcnArchName"] = zluda.default_agent.name torch.cuda._get_device_properties = torch_cuda__get_device_properties # pylint: disable=protected-access @@ -104,10 +95,13 @@ def do_hijack(): query = torch.nn.functional.pad(query, [0, 8 - head_size_og % 8]) key = torch.nn.functional.pad(key, [0, 8 - head_size_og % 8]) value = torch.nn.functional.pad(value, [0, 8 - head_size_og % 8]) - out_padded = interface_fa.fwd( - query.transpose(1, 2), + query = query.transpose(1, 2) + out_padded = torch.zeros_like(query) + interface_fa.fwd( + query, key.transpose(1, 2), value.transpose(1, 2), + out_padded, dropout_p, scale, is_causal, diff --git a/modules/zluda_installer.py b/modules/zluda_installer.py index 71ac5b21c..027c4f2ca 100644 --- a/modules/zluda_installer.py +++ b/modules/zluda_installer.py @@ -69,8 +69,8 @@ def set_default_agent(agent: rocm.Agent): default_agent = agent -def is_reinstall_needed() -> bool: # ZLUDA<3.8.7 - return not os.path.exists(os.path.join(path, 'cufftw.dll')) +def is_reinstall_needed() -> bool: # ZLUDA<3.9.4 + return os.path.exists(os.path.join(path, 'cudart.dll')) def install(): @@ -78,7 +78,7 @@ def install(): return platform = "windows" - commit = os.environ.get("ZLUDA_HASH", "dba64c0966df2c71e82255e942c96e2e1cea3a2d") + commit = os.environ.get("ZLUDA_HASH", "8d2128caf460b853b165cab0b4d8826b6b734ae7") if os.environ.get("ZLUDA_NIGHTLY", "0") == "1": log.warning("Environment variable 'ZLUDA_NIGHTLY' will be removed. Please use command-line argument '--use-nightly' instead.") args.use_nightly = True diff --git a/requirements.txt b/requirements.txt index eaa429d27..ce264dcb8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -31,7 +31,7 @@ invisible-watermark pi-heif # versioned -rich==13.9.4 +rich==14.0.0 safetensors==0.5.3 tensordict==0.1.2 peft==0.15.2 @@ -46,9 +46,9 @@ opencv-contrib-python-headless==4.9.0.80 einops==0.4.1 gradio==3.43.2 huggingface_hub==0.30.2 -numexpr==2.8.8 +numexpr==2.10.2 numpy==1.26.4 -numba==0.59.1 +numba==0.61.2 protobuf==4.25.3 pytorch_lightning==1.9.4 tokenizers==0.21.1 diff --git a/scripts/prompt_enhance.py b/scripts/prompt_enhance.py index 5125f4dcc..adc80d094 100644 --- a/scripts/prompt_enhance.py +++ b/scripts/prompt_enhance.py @@ -16,6 +16,12 @@ class Options: models = { 'google/gemma-3-1b-it': {}, 'google/gemma-3-4b-it': {}, + 'Qwen/Qwen3-0.6B-FP8': {}, + 'Qwen/Qwen3-1.7B-FP8': {}, + 'Qwen/Qwen3-4B-FP8': {}, + 'Qwen/Qwen3-0.6B': {}, + 'Qwen/Qwen3-1.7B': {}, + 'Qwen/Qwen3-4B': {}, 'Qwen/Qwen2.5-0.5B-Instruct': {}, 'Qwen/Qwen2.5-1.5B-Instruct': {}, 'Qwen/Qwen2.5-3B-Instruct': {}, @@ -50,6 +56,7 @@ class Options: do_sample: bool = True temperature: float = 0.15 repetition_penalty: float = 1.2 + thinking_mode: bool = False class Script(scripts.Script): @@ -195,7 +202,7 @@ class Script(scripts.Script): filtered = re.sub(pattern, '', prompt) return filtered, matches - def enhance(self, model: str=None, prompt:str=None, system:str=None, prefix:str=None, suffix:str=None, sample:bool=None, tokens:int=None, temperature:float=None, penalty:float=None): + def enhance(self, model: str=None, prompt:str=None, system:str=None, prefix:str=None, suffix:str=None, sample:bool=None, tokens:int=None, temperature:float=None, penalty:float=None, thinking:bool=False): model = model or self.options.default prompt = prompt or self.prompt.value prefix = prefix or '' @@ -204,6 +211,7 @@ class Script(scripts.Script): tokens = tokens or self.options.max_tokens penalty = penalty or self.options.repetition_penalty temperature = temperature or self.options.temperature + thinking = thinking or self.options.thinking_mode sample = sample if sample is not None else self.options.do_sample while self.busy: time.sleep(0.1) @@ -223,6 +231,7 @@ class Script(scripts.Script): inputs = self.tokenizer.apply_chat_template( chat_template, add_generation_prompt=True, + enable_thinking=thinking, tokenize=True, return_dict=True, return_tensors="pt", @@ -270,7 +279,7 @@ class Script(scripts.Script): response = self.post(response, prefix, suffix, networks) shared.log.info(f'Prompt enhance: model="{model}" time={t1-t0:.2f} inputs={input_len} outputs={outputs.shape[-1]} prompt={len(prompt)} response={len(response)}') if debug_enabled: - shared.log.trace(f'Prompt enhance: sample={sample} tokens={tokens} temperature={temperature} penalty={penalty}') + shared.log.trace(f'Prompt enhance: sample={sample} tokens={tokens} temperature={temperature} penalty={penalty} thinking={thinking}') shared.log.trace(f'Prompt enhance: prompt="{prompt}"') shared.log.trace(f'Prompt enhance: response="{response}"') self.busy = False @@ -279,7 +288,7 @@ class Script(scripts.Script): return prompt return response - def apply(self, prompt, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty): + def apply(self, prompt, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode): response = self.enhance( prompt=prompt, prefix=prompt_prefix, @@ -290,6 +299,7 @@ class Script(scripts.Script): tokens=max_tokens, temperature=temperature, penalty=repetition_penalty, + thinking=thinking_mode, ) if apply_prompt: return [response, response] @@ -339,6 +349,8 @@ class Script(scripts.Script): with gr.Row(): temperature = gr.Slider(label='Temperature', value=self.options.temperature, minimum=0.0, maximum=1.0, step=0.01, interactive=True) repetition_penalty = gr.Slider(label='Repetition penalty', value=self.options.repetition_penalty, minimum=0.0, maximum=2.0, step=0.01, interactive=True) + with gr.Row(): + thinking_mode = gr.Checkbox(label='Thinking mode', value=False, interactive=True) gr.HTML('
') with gr.Accordion('Input', open=False, elem_id='prompt_enhance_system_prompt'): with gr.Row(): @@ -355,15 +367,15 @@ class Script(scripts.Script): clear_btn.click(fn=lambda: '', inputs=[], outputs=[prompt_output]) copy_btn = gr.Button(value='Set prompt', elem_id='prompt_enhance_copy', variant='secondary') copy_btn.click(fn=lambda x: x, inputs=[prompt_output], outputs=[self.prompt]) - apply_btn.click(fn=self.apply, inputs=[self.prompt, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty], outputs=[prompt_output, self.prompt]) - return [apply_auto, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty] + apply_btn.click(fn=self.apply, inputs=[self.prompt, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode], outputs=[prompt_output, self.prompt]) + return [apply_auto, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode] def after_component(self, component, **kwargs): # searching for actual ui prompt components if getattr(component, 'elem_id', '') in ['txt2img_prompt', 'img2img_prompt', 'control_prompt', 'video_prompt']: self.prompt = component def before_process(self, p: processing.StableDiffusionProcessing, *args, **kwargs): # pylint: disable=unused-argument - apply_auto, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty = args + apply_auto, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty, thinking_mode = args if not apply_auto and not p.enhance_prompt: return if shared.state.skipped or shared.state.interrupted: @@ -383,6 +395,7 @@ class Script(scripts.Script): tokens=max_tokens, temperature=temperature, penalty=repetition_penalty, + thinking=thinking_mode, ) p.extra_generation_params['LLM'] = llm_model shared.state.end() diff --git a/scripts/pulid_ext.py b/scripts/pulid_ext.py index 5cf9a6778..366b13d32 100644 --- a/scripts/pulid_ext.py +++ b/scripts/pulid_ext.py @@ -137,7 +137,7 @@ class Script(scripts.Script): shared.log.error('PuLID: no images') return None - supported_model_list = ['sdxl'] + supported_model_list = ['sdxl', 'f1'] if shared.sd_model_type not in supported_model_list: shared.log.error(f'PuLID: class={shared.sd_model.__class__.__name__} model={shared.sd_model_type} required={supported_model_list}') return None @@ -168,15 +168,12 @@ class Script(scripts.Script): p.batch_size = 1 sdp = shared.opts.cross_attention_optimization == "Scaled-Dot-Product" - sampler_fn = getattr(self.pulid.sampling, f'sample_{sampler}', None) strength = getattr(p, 'pulid_strength', strength) zero = getattr(p, 'pulid_zero', zero) ortho = getattr(p, 'pulid_ortho', ortho) sampler = getattr(p, 'pulid_sampler', sampler) restore = getattr(p, 'pulid_restore', restore) p.pulid_restore = restore - if sampler_fn is None: - sampler_fn = self.pulid.sampling.sample_dpmpp_2m_sde if shared.sd_model_type == 'sdxl' and not hasattr(shared.sd_model, 'pipe'): try: @@ -203,9 +200,42 @@ class Script(scripts.Script): shared.log.error(f'PuLID: failed to create pipeline: {e}') errors.display(e, 'PuLID') return None + elif shared.sd_model_type == 'f1': + # TODO nunchaku: pulid-f1 + shared.log.error('PuLID: f1 not supported') + return None + if shared.sd_model_type == 'sdxl': + processed = self.run_sdxl(p, images, strength, zero, sampler, ortho, restore, offload, version) + elif shared.sd_model_type == 'f1': + processed = None + else: + shared.log.error(f'PuLID: class={shared.sd_model.__class__.__name__} model={shared.sd_model_type} required={supported_model_list}') + processed = None + return processed + + def after(self, p: processing.StableDiffusionProcessing, processed: processing.Processed, *args): # pylint: disable=unused-argument + _strength, _zero, _sampler, _ortho, _gallery, restore, _offload, _version = args + if shared.sd_model_type == "sdxl" and hasattr(shared.sd_model, 'pipe'): + restore = getattr(p, 'pulid_restore', restore) + if restore: + if hasattr(shared.sd_model, 'app'): + shared.sd_model.app = None + shared.sd_model.ip_adapter = None + shared.sd_model.face_helper = None + shared.sd_model.clip_vision_model = None + shared.sd_model.handler_ante = None + shared.sd_model = shared.sd_model.pipe + devices.torch_gc(force=True) + shared.log.debug(f'PuLID complete: class={shared.sd_model.__class__.__name__} preprocess={self.preprocess:.2f} pipe={"restore" if restore else "cache"}') + return processed + + def run_sdxl(self, p: processing.StableDiffusionProcessing, images: list, strength: float, zero: int, sampler: str, ortho: str, restore: bool, offload: bool, version: str): + sampler_fn = getattr(self.pulid.sampling, f'sample_{sampler}', None) + if sampler_fn is None: + sampler_fn = self.pulid.sampling.sample_dpmpp_2m_sde shared.sd_model.sampler = sampler_fn - shared.log.info(f'PuLID: class={shared.sd_model.__class__.__name__} version="{version}" sdp={sdp} strength={strength} zero={zero} ortho={ortho} sampler={sampler_fn} images={[i.shape for i in images]} offload={offload} restore={restore}') + shared.log.info(f'PuLID: class={shared.sd_model.__class__.__name__} version="{version}" strength={strength} zero={zero} ortho={ortho} sampler={sampler_fn} images={[i.shape for i in images]} offload={offload} restore={restore}') self.pulid.attention.NUM_ZERO = zero self.pulid.attention.ORTHO = ortho == 'v1' self.pulid.attention.ORTHO_v2 = ortho == 'v2' @@ -256,19 +286,3 @@ class Script(scripts.Script): # interim = [Image.fromarray(img) for img in shared.sd_model.debug_img_list] # shared.log.debug(f'PuLID: time={t1-t0:.2f}') return processed - - def after(self, p: processing.StableDiffusionProcessing, processed: processing.Processed, *args): # pylint: disable=unused-argument - _strength, _zero, _sampler, _ortho, _gallery, restore, _offload, _version = args - if hasattr(shared.sd_model, 'pipe') and shared.sd_model_type == "sdxl": - restore = getattr(p, 'pulid_restore', restore) - if restore: - if hasattr(shared.sd_model, 'app'): - shared.sd_model.app = None - shared.sd_model.ip_adapter = None - shared.sd_model.face_helper = None - shared.sd_model.clip_vision_model = None - shared.sd_model.handler_ante = None - shared.sd_model = shared.sd_model.pipe - devices.torch_gc(force=True) - shared.log.debug(f'PuLID complete: class={shared.sd_model.__class__.__name__} preprocess={self.preprocess:.2f} pipe={"restore" if restore else "cache"}') - return processed diff --git a/wiki b/wiki index 1afa48853..7d860a3b4 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 1afa48853740c39d93de07f64876e898dcb1daaa +Subproject commit 7d860a3b461df5e5976a8ecc0b66c2750b3b521c