Merge remote-tracking branch 'origin/diffusion-train-datasets' into diffusion-train-tab

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,
+    },
 ]
 
 

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