diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index ff026b3814..a5b4c3b06e 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1687,6 +1687,38 @@ _DATASET_EXAMPLES: list[dict] = [ "caption_column": "text", "no_checks": True, }, + { + "id": "smithsonian-butterflies", + "label": "Smithsonian Butterflies", + "repo": "huggan/smithsonian_butterflies_subset", + "description": ( + "100 butterfly specimen photos. The classic diffusers-docs training set. No " + "captions, so pair it with the trigger prompt to teach a butterfly subject." + ), + "license": "CC0 (Smithsonian Open Access)", + "image_cap": 100, + # The metadata columns are species names / boilerplate alt-text, not text-to-image + # captions, so train it as a subject set with the trigger prompt instead. + "suggested_trigger": "a photo of a sks butterfly", + "loader": "hf_dataset", + "caption_column": None, + "no_checks": False, + }, + { + "id": "pixel-nouns", + "label": "Nouns (pixel avatars)", + "repo": "m1guelpf/nouns", + "description": ( + "100 captioned Nouns pixel-art avatars. A captioned style set: each image ships " + "a caption, so the adapter learns the pixel look without a trigger word." + ), + "license": "cc0-1.0", + "image_cap": 100, + "suggested_trigger": None, + "loader": "hf_dataset", + "caption_column": "text", + "no_checks": False, + }, ] diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index 7a07e8eb47..b7a895281a 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -197,12 +197,33 @@ def test_list_dataset_examples(client, ds_root): r = client.get("/api/train/diffusion/dataset-examples") assert r.status_code == 200, r.text ids = {e["id"] for e in r.json()["examples"]} - assert {"dreambooth-dog", "tuxemon", "tarot-1920"} <= ids + assert { + "dreambooth-dog", + "tuxemon", + "tarot-1920", + "smithsonian-butterflies", + "pixel-nouns", + } <= ids dog = next(e for e in r.json()["examples"] if e["id"] == "dreambooth-dog") assert dog["suggested_trigger"] == "a photo of sks dog" assert dog["license"] +def test_list_dataset_examples_large_sets(client, ds_root): + # The two ~100-image sets: butterflies is a subject set (trigger, no caption column), + # nouns is a captioned style set (caption column, no trigger). Both cap at 100. + r = client.get("/api/train/diffusion/dataset-examples") + examples = {e["id"]: e for e in r.json()["examples"]} + butterflies = examples["smithsonian-butterflies"] + assert butterflies["image_cap"] == 100 + assert butterflies["suggested_trigger"] == "a photo of a sks butterfly" + assert "CC0" in butterflies["license"] + nouns = examples["pixel-nouns"] + assert nouns["image_cap"] == 100 + assert nouns["suggested_trigger"] is None + assert nouns["license"] == "cc0-1.0" + + class _FakeImageFeature: # Mimics datasets.Image so _detect_image_column matches by class name. pass