serve text-to-image and image-to-image from one copy of the weights (v0.4.94)

A checkpoint is a UNet, a VAE and its text encoders; which task it performs is
a property of the pipeline class wrapped around them, not of the weights. Both
tasks already worked after 0.4.93, but each endpoint loaded its OWN pipeline
under its own key -- so a client that creates a character with a model and then
edits it with the same one held an 18 GB SDXL fine-tune twice, and the second
load evicted the first. Alternating between the two thrashes the card
continuously.

Reading the two loaders side by side turned up the other half of it:
_load_img2img_pipeline ignored the model's configuration altogether -- dtype
hardcoded from global args, no offload, no quantization, no VAE override. Only
_load_diffusers_pipeline honours models.json. A 4-bit model was therefore
served at 4 bits for generation and bf16 for editing, out of two copies.

So the model is loaded once, through the configured loader, under the manager's
own key -- its VRAM accounting and eviction now see one model -- and each task
gets a view built with AutoPipelineFor*.from_pipe around the components already
resident. /v1/images/generations gets the same treatment in reverse, so a model
an edit loaded can create a character without reloading.

The views are cached on the base pipeline, not as extra models-dict entries.
Deliberately: they share the base's tensors and so must die exactly when it
does, and a second models-dict key per task is what let a cache key reach
from_pretrained in 0.4.93. The img2img:/inpaint: keys are gone from the module
entirely.

Task detection asks diffusers' own tables rather than a signature, because a
ControlNet text-to-image pipeline requires an `image` -- the control hint, not
a source -- and a signature check would file it as img2img. Tested.

Two behaviours worth knowing: the views share the UNet, so a LoRA or
IP-Adapter applied through one task is visible to the others (same weights, by
design); and SD, SDXL, Flux and Z-Image all have entries for all three tasks,
so this is not an SD-only arrangement. I had assumed Z-Image was unmapped and
wrote a test asserting it; the test caught that and both it and a code comment
are corrected.

28 tests here, 187 across the image files; full suite 4549 passed with the same
15 pre-existing failures as 0.4.91 (diff empty).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01S8LAP3ueNzQJtwzhgVbq66

codai/__init__.py

@@ -16,7 +16,7 @@
 
 # Canonical product version for CoderAI — single source of truth. Both the API
 # metadata and the admin web UI read from here.
-__version__ = "0.4.93"
+__version__ = "0.4.94"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

codai/api/images.py

@@ -1098,6 +1098,16 @@ async def _generate_with_diffusers(pipeline, request, global_args, http_request=
     import numpy as np
     import time as time_module
 
+    # The same weights serve both directions: a model an edit request loaded as
+    # img2img can create a character here without a second copy. Unknown
+    # families are left exactly as they are.
+    _view = _task_view(pipeline, "text2img")
+    if _view is None:
+        raise HTTPException(status_code=501, detail=(
+            f"{type(pipeline).__name__} cannot generate from a prompt alone: "
+            f"diffusers offers no text-to-image pipeline for this family."))
+    pipeline = _view
+
     if getattr(request, 'disable_safety_checker', False):
         _disable_safety_checker(pipeline)
 
@@ -2077,24 +2087,147 @@ def _accepts_image(pipeline) -> bool:
     return not named
 
 
-def _as_image_task(pipeline, task: str):
-    """``pipeline`` as an img2img/inpaint pipeline, reusing its weights.
+#: The three things one set of image weights can be asked to do. A checkpoint
+#: is a UNet, a VAE and its text encoders; which of these tasks it performs is
+#: a property of the pipeline class wrapped around them, not of the weights.
+TASKS = ("text2img", "img2img", "inpaint")
+
+#: Where the per-task pipelines hang off the loaded model. Deliberately an
+#: attribute of the base pipeline rather than another entry in the manager's
+#: models dict: the views share the base's components, so they must die exactly
+#: when it does, and a second models-dict key is what let a cache key end up in
+#: a from_pretrained() call (see _TASK_NAMESPACES in codai/models/manager.py).
+_VIEWS_ATTR = "_coderai_task_views"
 
-    ``AutoPipelineFor*.from_pipe`` swaps the task class around the components
-    already in VRAM, so recovering from a text-to-image pipeline costs no
-    second copy of the weights. Returns None when it cannot be converted.
+
+def _task_mappings():
+    from diffusers.pipelines.auto_pipeline import (
+        AUTO_IMAGE2IMAGE_PIPELINES_MAPPING,
+        AUTO_INPAINT_PIPELINES_MAPPING,
+        AUTO_TEXT2IMAGE_PIPELINES_MAPPING,
+    )
+    return {
+        "text2img": AUTO_TEXT2IMAGE_PIPELINES_MAPPING,
+        "img2img": AUTO_IMAGE2IMAGE_PIPELINES_MAPPING,
+        "inpaint": AUTO_INPAINT_PIPELINES_MAPPING,
+    }
+
+
+def _task_of(pipeline) -> Optional[str]:
+    """Which of :data:`TASKS` this pipeline class performs, or None.
+
+    Asked of diffusers' own task tables rather than guessed: a signature check
+    would call a ControlNet text-to-image pipeline "img2img", since its
+    required ``image`` is the control hint. None means a family diffusers does
+    not map, which is left alone -- a smaller set than it sounds, since SD,
+    SDXL, Flux and Z-Image all have entries for all three tasks.
     """
-    from diffusers import AutoPipelineForImage2Image, AutoPipelineForInpainting
-    Auto = AutoPipelineForInpainting if task == "inpaint" else AutoPipelineForImage2Image
+    name = type(pipeline).__name__
+    for task, mapping in _task_mappings().items():
+        if any(getattr(c, "__name__", None) == name for c in mapping.values()):
+            return task
+    return None
+
+
+def _task_view(base, task: str):
+    """``base`` as a pipeline for ``task``, sharing one copy of the weights.
+
+    This is what lets one model serve both character creation and editing.
+    ``AutoPipelineFor*.from_pipe`` builds the task's pipeline class around the
+    components already in VRAM, so a model used for text-to-image and
+    image-to-image costs what it costs once -- 18 GB for an SDXL fine-tune,
+    which is the difference between both tasks being resident and an eviction
+    on every alternation.
+
+    Views are cached on the base and shared, so a LoRA or IP-Adapter applied
+    through one task is visible to the others: same UNet, by design.
+
+    Returns the base unchanged when it already performs ``task`` or when its
+    family is not one diffusers maps, and None when a conversion was needed
+    and failed.
+    """
+    current = _task_of(base)
+    if current == task or current is None:
+        return base
+
+    views = getattr(base, _VIEWS_ATTR, None)
+    if views is None:
+        views = {}
+        try:
+            setattr(base, _VIEWS_ATTR, views)
+        except Exception:
+            pass        # a pipeline that refuses attributes still converts
+    if task in views:
+        return views[task]
+
+    from diffusers import (AutoPipelineForImage2Image, AutoPipelineForInpainting,
+                           AutoPipelineForText2Image)
+    Auto = {"text2img": AutoPipelineForText2Image,
+            "img2img": AutoPipelineForImage2Image,
+            "inpaint": AutoPipelineForInpainting}[task]
     try:
-        converted = Auto.from_pipe(pipeline)
+        view = Auto.from_pipe(base)
     except Exception as exc:
-        _log.warning("cannot use %s as a %s pipeline: %s",
-                     type(pipeline).__name__, task, exc)
+        _log.warning("cannot use %s for %s: %s",
+                     type(base).__name__, task, exc)
         return None
-    _log.info("converted %s to %s (same weights, no second copy)",
-              type(pipeline).__name__, type(converted).__name__)
-    return converted
+    views[task] = view
+    _log.info("%s also serving %s as %s (same weights, no second copy)",
+              type(base).__name__, task, type(view).__name__)
+    return view
+
+
+async def _pipeline_for_task(model_info: dict, task: str):
+    """The pipeline for ``task``, from ONE loaded copy of the model's weights.
+
+    Motus creates characters (text-to-image) and edits them (image-to-image)
+    with the same checkpoint, so loading one per task would hold the weights
+    twice -- 18 GB each for an SDXL fine-tune, and an eviction on every
+    alternation between the two. The model is loaded once, under the manager's
+    own key so its VRAM accounting and eviction see one model, and each task
+    gets a view onto those components.
+
+    The load goes through ``_load_diffusers_pipeline``, which is the only
+    loader that honours the model's config -- precision, offload, quantization,
+    a VAE override. The editing endpoints used to load their own copy with a
+    hardcoded dtype and none of that, so a 4-bit model served text-to-image at
+    4 bits and editing at bf16, out of two copies.
+    """
+    model_name = model_info.get('model_name')
+    model_key = model_info.get('model_key') or f"image:{model_name}"
+    base = model_info.get('model_object') or multi_model_manager.models.get(model_key)
+
+    if base is None:
+        cfg = (multi_model_manager.config.get(model_key)
+               or multi_model_manager.config.get(model_name) or {})
+        try:
+            base = await asyncio.to_thread(
+                _load_diffusers_pipeline, model_name, global_args, cfg)
+        except Exception as e:
+            base = None
+            _log.warning("the configured loader could not open %r for %s (%s); "
+                         "falling back to a task-only load", model_name, task, e)
+        if base is not None:
+            multi_model_manager.add_model(model_key, base)
+            multi_model_manager.current_model_key = model_key
+
+    if base is None:
+        # Last resort: a task-specific load, which is all the editing
+        # endpoints ever did. Not registered under the shared key, because it
+        # carries none of the model's configuration.
+        loader = _load_inpaint_pipeline if task == "inpaint" else _load_img2img_pipeline
+        try:
+            return await asyncio.to_thread(loader, model_name, global_args)
+        except Exception as e:
+            raise HTTPException(status_code=500,
+                                detail=f"Failed to load {task} model: {e}")
+
+    view = _task_view(base, task)
+    if view is None or (task in ("img2img", "inpaint") and not _accepts_image(view)):
+        raise HTTPException(status_code=501, detail=(
+            f"'{model_name}' loaded as {type(base).__name__}, which cannot do "
+            f"{task}: diffusers offers no {task} pipeline for this family."))
+    return view
 
 
 def _load_img2img_pipeline(model_name: str, global_args):
@@ -2163,30 +2296,7 @@ async def create_image_edit(request: ImageEditRequest, http_request: Request = N
         err = model_info.get('error', f"Model '{request.model}' not found or not registered")
         raise HTTPException(status_code=404, detail=err)
 
-    model_key = f"img2img:{model_name}"
-    pipe = multi_model_manager.models.get(model_key)
-
-    if pipe is None:
-        try:
-            pipe = await asyncio.get_event_loop().run_in_executor(
-                None, _load_img2img_pipeline, model_name, global_args
-            )
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Failed to load img2img model: {e}")
-        multi_model_manager.models[model_key] = pipe
-
-    # A pipeline that cannot take an image must never serve an edit: it would
-    # accept `image` into **kwargs, drop it, and answer 200 with a picture
-    # generated from the prompt alone. Convert it (same weights) or say so.
-    if not _accepts_image(pipe):
-        converted = _as_image_task(pipe, "img2img")
-        if converted is None:
-            raise HTTPException(status_code=501, detail=(
-                f"'{model_name}' loaded as {type(pipe).__name__}, which cannot "
-                f"edit an image: it has no `image` input. Editing needs a "
-                f"checkpoint diffusers can open for image-to-image."))
-        pipe = converted
-        multi_model_manager.models[model_key] = pipe
+    pipe = await _pipeline_for_task(model_info, "img2img")
 
     try:
         import torch
@@ -2326,23 +2436,8 @@ async def create_image_inpaint(request: ImageInpaintRequest, http_request: Reque
     model_name = model_info.get('model_name')
     if not model_name:
         raise HTTPException(status_code=404, detail=model_info.get('error', 'Model not found'))
-    model_key = f"inpaint:{model_name}"
-    pipe = multi_model_manager.models.get(model_key)
-    if pipe is None:
-        try:
-            pipe = await asyncio.get_event_loop().run_in_executor(
-                None, _load_inpaint_pipeline, model_name, global_args)
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Failed to load inpaint model: {e}")
-        multi_model_manager.models[model_key] = pipe
-    if not _accepts_image(pipe):
-        converted = _as_image_task(pipe, "inpaint")
-        if converted is None:
-            raise HTTPException(status_code=501, detail=(
-                f"'{model_name}' loaded as {type(pipe).__name__}, which cannot "
-                f"inpaint: it has no `image` input."))
-        pipe = converted
-        multi_model_manager.models[model_key] = pipe
+    pipe = await _pipeline_for_task(model_info, "inpaint")
+    model_key = model_info.get('model_key') or f"image:{model_name}"
     try:
         import torch
         source_img = _decode_b64_image(request.image)

docs/image-prompts.md

@@ -103,6 +103,43 @@ source photo's own aspect at about 1 MP on a 64-pixel grid. Both used to
 default to a square, which squashed every non-square photo — and on inpaint it
 pulled the mask out of alignment with the image as well.
 
+## One model, both directions
+
+A checkpoint is a UNet, a VAE and its text encoders. Whether it generates from
+a prompt, edits a picture or inpaints a region is a property of the *pipeline
+class* wrapped around those weights, not of the weights themselves — so the
+same model serves all three, from one copy in VRAM.
+
+That matters when a client creates a character with a model and then edits it
+with the same one. Loading a pipeline per task would hold an 18 GB SDXL
+fine-tune twice and evict on every alternation between the two. Instead the
+model is loaded once, under the manager's own key so its VRAM accounting and
+eviction see one model, and each task gets a view built with
+`AutoPipelineFor*.from_pipe` around the components already resident.
+
+Two consequences worth knowing:
+
+* The views **share the UNet**, so a LoRA or IP-Adapter applied through one
+  task is visible to the others. That is by design, not a leak.
+* The load goes through the configured loader, so precision, offload,
+  quantization and a VAE override apply to editing as well. They did not
+  before: the edit endpoints loaded their own copy with a hardcoded dtype, so
+  a 4-bit model was served at 4 bits for generation and bf16 for editing, out
+  of two copies.
+
+The task class comes from the checkpoint's `model_index.json`, never from
+substrings in its name. Guessing is what broke editing on the booru-tag
+fine-tunes: `nesaorg/PonyRealism-Uncensored` contains no `xl`, so it was loaded
+with the generic pipeline, which honoured the declared
+`StableDiffusionXLPipeline` — the *text-to-image* class. Every diffusers
+`__call__` ends in `**kwargs`, so `image` and `strength` were accepted and
+dropped, and every "edit" was a fresh 1024×1024 render that never saw the
+photo: `200 OK`, a plausible-looking image, and `size` apparently ignored. A
+pipeline with no `image` input is now converted or refused with `501`.
+
+SD, SDXL, Flux and Z-Image all have entries for all three tasks. A family
+diffusers does not map is left on whatever it was loaded as.
+
 ## Standing prefixes: the ladder a model always wants
 
 CoderAI never adds tags to a prompt by itself. The `score_*` ladder is yours to

tests/test_image_edit_uses_an_image_pipeline.py

@@ -132,31 +132,139 @@ def test_conversion_reuses_the_weights_rather_than_loading_them_twice():
     assert "from_pipe" in src
 
 
-def test_an_unconvertible_pipeline_is_refused_not_served():
-    """501, not a 200 carrying a picture of something else."""
+def _named(name, **attrs):
+    """An object whose class is called ``name``, which is how task detection
+    identifies a pipeline (diffusers' own tables are keyed by class)."""
+    body = {"__call__": lambda self, **kw: None}
+    body.update(attrs)
+    return type(name, (), body)()
+
+
+def test_task_detection_reads_diffusers_tables():
+    assert images_mod._task_of(_named("StableDiffusionXLPipeline")) == "text2img"
+    assert images_mod._task_of(_named("StableDiffusionXLImg2ImgPipeline")) == "img2img"
+    assert images_mod._task_of(_named("StableDiffusionXLInpaintPipeline")) == "inpaint"
+
+
+def test_a_controlnet_text2img_pipeline_is_not_mistaken_for_img2img():
+    """Why detection asks the tables instead of the signature: ControlNet
+    text-to-image requires an `image` -- the control hint, not a source."""
+    pipe = _named("StableDiffusionXLControlNetPipeline")
+    assert images_mod._task_of(pipe) == "text2img"
+
+
+def test_an_unmapped_family_is_left_alone():
+    """A pipeline diffusers does not map is not converted in either direction.
+
+    Note which families ARE mapped, because the list is wider than it looks:
+    Z-Image and Flux both have text2img, img2img and inpaint entries, so they
+    share one copy of the weights across the tasks too.
+    """
+    odd = _named("SomeHouseTrainedPipeline")
+    assert images_mod._task_of(odd) is None
+    assert images_mod._task_view(odd, "text2img") is odd
+    assert images_mod._task_view(odd, "img2img") is odd
+
+
+@pytest.mark.parametrize("cls_name,task", [
+    ("ZImagePipeline", "text2img"),
+    ("ZImageImg2ImgPipeline", "img2img"),
+    ("FluxPipeline", "text2img"),
+    ("FluxImg2ImgPipeline", "img2img"),
+    ("FluxInpaintPipeline", "inpaint"),
+])
+def test_the_modern_families_are_mapped_too(cls_name, task):
+    """So the sharing is not an SD/SDXL-only arrangement."""
+    assert images_mod._task_of(_named(cls_name)) == task
+
+
+def test_asking_for_the_task_it_already_is_costs_nothing():
+    base = _named("StableDiffusionXLImg2ImgPipeline")
+    assert images_mod._task_view(base, "img2img") is base
+    assert not hasattr(base, images_mod._VIEWS_ATTR)
+
+
+def test_one_copy_of_the_weights_serves_every_task(monkeypatch):
+    """The point of the whole arrangement: Motus creates a character with a
+    model and then edits it with the same one, and that must not hold two
+    copies of an 18 GB checkpoint or evict on every alternation."""
+    import diffusers
+    conversions = []
+
+    def fake_from_pipe(pipe, **kw):
+        conversions.append(pipe)
+        return _named("StableDiffusionXLImg2ImgPipeline", image=None)
+
+    monkeypatch.setattr(diffusers.AutoPipelineForImage2Image, "from_pipe",
+                        staticmethod(fake_from_pipe))
+    base = _named("StableDiffusionXLPipeline")
+    first = images_mod._task_view(base, "img2img")
+    second = images_mod._task_view(base, "img2img")
+    assert first is second, "converted twice instead of reusing the view"
+    assert len(conversions) == 1
+    assert conversions[0] is base, "converted from something other than the base"
+    # and the view hangs off the base, so it dies when the weights do
+    assert images_mod._VIEWS_ATTR in vars(base)
+    assert vars(base)[images_mod._VIEWS_ATTR]["img2img"] is first
+
+
+def test_a_failed_conversion_is_reported_not_swallowed(monkeypatch):
+    import diffusers
+
+    def boom(pipe, **kw):
+        raise RuntimeError("no img2img for this family")
+
+    monkeypatch.setattr(diffusers.AutoPipelineForImage2Image, "from_pipe",
+                        staticmethod(boom))
+    assert images_mod._task_view(_named("StableDiffusionXLPipeline"), "img2img") is None
+
+
+def test_a_pipeline_that_cannot_do_the_task_is_refused_with_501():
+    """Not served: a 200 carrying a picture generated from the prompt alone is
+    the outcome this whole change exists to rule out."""
+    import asyncio
+
+    from fastapi import HTTPException
+
+    info = {"model_name": "some/model", "model_key": "image:some/model",
+            "model_object": _named("ZImagePipeline")}   # unmapped, no `image`
+    with pytest.raises(HTTPException) as exc:
+        asyncio.run(images_mod._pipeline_for_task(info, "img2img"))
+    assert exc.value.status_code == 501
+    assert "cannot do img2img" in exc.value.detail
+
+
+def test_the_task_pipelines_share_the_manager_entry():
+    """One models-dict key per model. A second key per task is both a double
+    load and how a cache key reached from_pretrained in the first place."""
+    src = open("codai/api/images.py").read()
+    head = src.index("async def _pipeline_for_task")
+    body = src[head:head + 2600]
+    assert "model_info.get('model_key')" in body
+    assert "add_model(model_key" in body
+    # the task-prefixed keys are gone from this module entirely
+    import re
+    assert not re.search(r'f"img2img:\{model_name\}"', src)
+    assert not re.search(r'f"inpaint:\{model_name\}"', src)
+
+
+def test_the_configured_loader_is_the_one_used():
+    """The editing endpoints used to load their own copy with a hardcoded
+    dtype, so a 4-bit model was served at 4 bits for text-to-image and bf16
+    for editing -- out of two copies of the weights."""
     src = open("codai/api/images.py").read()
-    assert "cannot \nedit an image" in src or "cannot " in src
-    assert src.count("status_code=501") >= 2, \
-        "both edits and inpaint must refuse rather than silently generate"
-    assert "has no `image` input" in src
+    head = src.index("async def _pipeline_for_task")
+    body = src[head:head + 2600]
+    assert "_load_diffusers_pipeline" in body
+    assert body.index("_load_diffusers_pipeline") < body.index("_load_img2img_pipeline")
 
 
-@pytest.mark.parametrize("endpoint", ["create_image_edit", "create_image_inpaint"])
-def test_the_guard_runs_on_a_cached_pipeline_too(endpoint):
-    """A pipeline loaded by an earlier version is already in the manager's
-    models dict, so checking only after a fresh load would leave every running
-    install serving text-to-image until it restarted."""
+def test_text2img_also_gets_a_view_so_the_sharing_works_both_ways():
+    """A model an edit loaded can create a character without a second copy."""
     src = open("codai/api/images.py").read()
-    head = src.index(f"async def {endpoint}")
-    body = src[head:head + 4000]
-    guard = body.index("_accepts_image(pipe)")
-    cached_lookup = body.index("multi_model_manager.models.get(")
-    assert guard > cached_lookup
-    # and it is not nested inside the `if pipe is None:` load branch
-    line_start = body.rindex("\n", 0, guard) + 1
-    indent = len(body[line_start:guard]) - len(body[line_start:guard].lstrip())
-    assert indent <= 4, "the guard is inside the load branch, so a cached " \
-                        "text-to-image pipeline would still be used"
+    head = src.index("async def _generate_with_diffusers")
+    body = src[head:head + 1500]
+    assert '_task_view(pipeline, "text2img")' in body
 
 
 # --------------------------------------------------------------------------- #