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.

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
 
 
 # --------------------------------------------------------------------------- #