Merge branch 'diffusion-phase11-consumer-int8' into diffusion-phase12-fbcache

studio/backend/core/inference/diffusion.py

@@ -45,6 +45,7 @@ from .diffusion_memory import (
     snapshot_device_memory,
 )
 from .diffusion_speed import (
+    SPEED_DEFAULT,
     SPEED_OFF,
     apply_speed_optims,
     resolve_speed_mode,
@@ -62,6 +63,7 @@ from .diffusion_prequant import (
     resolve_prequant_source,
 )
 from .diffusion_transformer_quant import (
+    DEFAULT_MIN_LINEAR_FEATURES,
     dense_transformer_supported,
     normalize_transformer_quant,
     quantize_transformer,
@@ -243,6 +245,22 @@ class DiffusionBackend:
                 base, rfilename, hf_token, cancel_event = self._cancel_event
             )
 
+    @staticmethod
+    def _detect_family_for_pick(
+        repo_id: str, gguf_filename: Optional[str], family_override: Optional[str]
+    ) -> Optional[DiffusionFamily]:
+        """Detect the family from the repo id, falling back to the combined
+        path/filename for a direct local .gguf pick. The frontend splits such a
+        pick into (parent dir, basename), so the family keyword can live only in
+        the filename (e.g. /models/z-image-turbo-Q4_K_M.gguf) while the parent
+        directory carries none; scan it too when the directory alone is
+        undetectable. Only used as a fallback, so remote 'org/name' picks and
+        explicit overrides behave exactly as before."""
+        fam = detect_family(repo_id, family_override)
+        if fam is None and gguf_filename and not family_override:
+            fam = detect_family(f"{repo_id}/{gguf_filename}", family_override)
+        return fam
+
     def validate_load_request(
         self,
         repo_id: str,
@@ -259,7 +277,7 @@ class DiffusionBackend:
             raise ValueError(
                 "gguf_filename is required: this backend loads single-file GGUF checkpoints only."
             )
-        fam = detect_family(repo_id, family_override)
+        fam = self._detect_family_for_pick(repo_id, gguf_filename, family_override)
         if fam is None:
             raise ValueError(
                 f"Could not infer a diffusion family for '{repo_id}'. Pass family_override (z-image)."
@@ -272,7 +290,12 @@ class DiffusionBackend:
         local_root = Path(repo_id).expanduser()
         if local_root.exists():
             resolve_local_gguf_child(local_root, gguf_filename)
-        elif repo_id.startswith(("/", "~", "./", "../")) or local_root.is_absolute():
+        elif (
+            # POSIX path-shaped, a "."/".." prefix (covers ./ ../ and their Windows .\ ..\
+            # forms), a Windows separator anywhere (never present in a bare "org/name" HF
+            # id), or an absolute path on this OS.
+            repo_id.startswith(("/", "\\", "~", ".")) or "\\" in repo_id or local_root.is_absolute()
+        ):
             raise FileNotFoundError(f"Local model path does not exist: {repo_id}")
         return fam
 
@@ -345,7 +368,9 @@ class DiffusionBackend:
             # Resolve the base repo and estimate sizes on this thread (both network
             # calls) so begin_load returns instantly; the bar shows raw bytes until
             # the total lands. This is the only writer of _loading's fields here.
-            fam = detect_family(kwargs["repo_id"], kwargs.get("family_override"))
+            fam = self._detect_family_for_pick(
+                kwargs["repo_id"], kwargs.get("gguf_filename"), kwargs.get("family_override")
+            )
             base = _resolve_base_repo(
                 kwargs["repo_id"], kwargs.get("base_repo"), fam, kwargs.get("hf_token")
             )
@@ -353,10 +378,12 @@ class DiffusionBackend:
             expected, base_files = self._estimate_download_bytes(
                 kwargs["repo_id"], kwargs.get("gguf_filename"), base, kwargs.get("hf_token")
             )
-            loading = self._loading
-            if loading is not None:
-                loading.base_repo = base
-                loading.expected_bytes = expected
+            with self._lock:
+                # Stamp progress only if this load is still current; a superseding
+                # load (or unload) has its own token and its own _LoadingState.
+                if self._load_token == token and self._loading is not None:
+                    self._loading.base_repo = base
+                    self._loading.expected_bytes = expected
             # Download outside the lock so unload()/an eviction can preempt the
             # multi-GB pull; load_pipeline below then assembles from the cache.
             self._prefetch_files(
@@ -378,9 +405,13 @@ class DiffusionBackend:
             if self._load_token != token:
                 return
             logger.error("diffusion.load_failed: %s", exc)
+            # Redact native paths: this error is surfaced verbatim via the
+            # load-progress poll, and Studio can run as a shared server.
+            from utils.native_path_leases import redact_native_paths
+
             with self._lock:
                 if self._load_token == token and self._loading is not None:
-                    self._loading.error = str(exc)
+                    self._loading.error = redact_native_paths(str(exc))
 
     def load_progress(self) -> dict[str, Any]:
         """Phase + downloaded/total bytes for the in-flight load (cache-scan based)."""
@@ -412,7 +443,11 @@ class DiffusionBackend:
         total = 0
         base_files: list[str] = []
         try:
-            if gguf_filename:
+            # Skip the Hub size lookup for a LOCAL gguf path: model_info(repo_id) would
+            # raise on a filesystem path and (caught below) skip the base-repo lookup too,
+            # so the companion VAE/text-encoder files would never be prefetched and would
+            # instead download synchronously under the load lock.
+            if gguf_filename and not Path(repo_id).expanduser().exists():
                 info = api.model_info(repo_id, files_metadata = True, token = hf_token)
                 total += sum(s.size or 0 for s in info.siblings if s.rfilename == gguf_filename)
             base_info = api.model_info(base_repo, files_metadata = True, token = hf_token)
@@ -563,6 +598,18 @@ class DiffusionBackend:
                 # the quant noise floor), dense models stay bit-identical `off`. An
                 # explicit speed_mode (incl. "off") is honored verbatim.
                 effective_speed = resolve_speed_mode(speed_mode, is_gguf = bool(gguf_filename))
+                # A torchao-quantized dense transformer runs its matmuls through the
+                # regional torch.compile; UNcompiled (eager) it is ~30x slower and would
+                # lose to the GGUF fallback. A dense model otherwise resolves to `off`, so
+                # force at least `default` (regional compile) whenever the quant engaged,
+                # or the opt-in "fast" path silently commits an eager, pathologically slow
+                # pipeline.
+                if transformer_quant_engaged is not None and effective_speed == SPEED_OFF:
+                    logger.info(
+                        "diffusion.transformer_quant: forcing speed_mode=default "
+                        "(quantized transformer must be compiled; eager is ~30x slower)"
+                    )
+                    effective_speed = SPEED_DEFAULT
                 # Opt-in speed optims run BEFORE placement (channels_last / compile
                 # must precede CPU offload). Snapshot the process-wide backend flags
                 # first so unload can restore them: TF32 / cudnn.benchmark are global,
@@ -603,6 +650,16 @@ class DiffusionBackend:
                     cache_active = cache_engaged is not None,
                     logger = logger,
                 )
+                if transformer_quant_engaged is not None and not speed_applied.get("compiled"):
+                    # Promotion above could not engage compile (e.g. the family is not
+                    # compile-friendly, or compile_repeated_blocks failed): the quantized
+                    # transformer is now running eager, which is far slower than the GGUF
+                    # path it replaced. Surface it loudly rather than hiding the regression.
+                    logger.warning(
+                        "diffusion.transformer_quant: %s engaged but the transformer is NOT "
+                        "compiled; eager torchao quant is ~30x slower than GGUF here",
+                        transformer_quant_engaged,
+                    )
                 # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4),
                 # also before placement so the offload hooks move the smaller weights.
                 te_quant = quantize_text_encoders(
@@ -696,6 +753,9 @@ class DiffusionBackend:
                     dtype = dtype,
                     hf_token = hf_token,
                     scheme = scheme,
+                    # Reject a checkpoint built with a different Linear filter than the
+                    # dense path uses, so the prequant and runtime-quant models match.
+                    min_features = DEFAULT_MIN_LINEAR_FEATURES,
                     logger = logger,
                 )
                 if transformer is not None: