diffusion skills framework

examples/flux2/model_inference/FLUX.2-klein-base-4B-skills.py

@@ -0,0 +1,56 @@
+from diffsynth.diffusion.skills import SkillsPipeline
+from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig
+import torch
+from PIL import Image
+
+
+pipe = Flux2ImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"),
+        ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"),
+)
+skills = SkillsPipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="DiffSynth-Studio/F2KB4B-Skills-ControlNet"),
+        ModelConfig(model_id="DiffSynth-Studio/F2KB4B-Skills-Brightness"),
+    ],
+)
+skill_cache = skills(
+    positive_inputs = [
+        {
+            "model_id": 0,
+            "image": Image.open("xxx.jpg"),
+            "prompt": "一位长发少女，四周环绕着魔法粒子",
+        },
+        {
+            "model_id": 1,
+            "scale": 0.6,
+        },
+    ],
+    negative_inputs = [
+        {
+            "model_id": 0,
+            "image": Image.open("xxx.jpg"),
+            "prompt": "一位长发少女，四周环绕着魔法粒子",
+        },
+        {
+            "model_id": 1,
+            "scale": 0.5,
+        },
+    ],
+    pipe=pipe,
+)
+image = pipe(
+    prompt="一位长发少女，四周环绕着魔法粒子",
+    seed=0, rand_device="cuda", num_inference_steps=50, cfg_scale=4,
+    height=1024, width=1024,
+    **skill_cache,
+)
+image.save("image.jpg")