diffsynth 2.0.14__tar.gz → 2.0.16__tar.gz

{diffsynth-2.0.14 → diffsynth-2.0.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsynth
-Version: 2.0.14
+Version: 2.0.16
 Summary: Enjoy the magic of Diffusion models!
 Author: ModelScope Team
 License: Apache-2.0

{diffsynth-2.0.14 → diffsynth-2.0.16}/README.md

@@ -34,7 +34,9 @@ We believe that a well-developed open-source code framework can lower the thresh
 
 > Currently, the development personnel of this project are limited, with most of the work handled by [Artiprocher](https://github.com/Artiprocher) and [mi804](https://github.com/mi804). Therefore, the progress of new feature development will be relatively slow, and the speed of responding to and resolving issues is limited. We apologize for this and ask developers to understand.
 
-- **June 16, 2026**: We have added a new Template model for ACE-Step: [vocals2music](https://www.modelscope.cn/models/DiffSynth-Studio/acestep15xlsft-vocals2music). For more details, please refer to the [documentation](/docs/zh/Model_Details/ACE-Step.md) and [example code](/examples/ace_step/).
+- **June 24, 2026** Krea-2 is now open-source, and we have provided full support. For more details, please refer to the [documentation](/docs/en/Model_Details/Krea-2.md) and [example code](/examples/krea2/).
+
+- **June 16, 2026**: We have added a new Template model for ACE-Step: [vocals2music](https://www.modelscope.cn/models/DiffSynth-Studio/acestep15xlsft-vocals2music). For more details, please refer to the [documentation](/docs/en/Model_Details/ACE-Step.md) and [example code](/examples/ace_step/).
 
 - **June 15, 2026** We have open-sourced Image-to-LoRA V2, compressing the hours-long training process for image style LoRAs into a single model inference step, thereby exploring a new paradigm for LoRA model training. The [technical report](https://arxiv.org/abs/2606.13809) has been released. This release includes three models:
     * [DiffSynth-Studio/ZImage-i2L-v2](https://modelscope.cn/models/DiffSynth-Studio/ZImage-i2L-v2): Adapted for the Z-Image model
@@ -1042,6 +1044,59 @@ Example code for Ideogram 4 is available at: [/examples/ideogram4/](/examples/id
 
 </details>
 
+#### Krea-2: [/docs/en/Model_Details/Krea-2.md](/docs/en/Model_Details/Krea-2.md)
+
+<details>
+
+<summary>Quick Start</summary>
+
+Running the following code will quickly load the [krea/Krea-2-Raw](https://www.modelscope.cn/models/krea/Krea-2-Raw) model and perform inference. VRAM management is enabled, and the framework will automatically control the loading of model parameters based on available VRAM. The model can run with a minimum of 24GB VRAM.
+
+```python
+from diffsynth.pipelines.krea2 import Krea2Pipeline, ModelConfig
+import torch
+
+vram_config = {
+    "offload_dtype": "disk",
+    "offload_device": "disk",
+    "onload_dtype": torch.float8_e4m3fn,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.float8_e4m3fn,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
+pipe = Krea2Pipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="krea/Krea-2-Raw", origin_file_pattern="raw.safetensors", **vram_config),
+        ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors", **vram_config),
+        ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config),
+    ],
+    tokenizer_config=ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern=""),
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 1,
+)
+prompt = "A cat standing on a stone."
+image = pipe(prompt, seed=0, num_inference_steps=52, cfg_scale=4.5)
+image.save("image.jpg")
+```
+
+</details>
+
+<details>
+
+<summary>Examples</summary>
+
+Example code for Krea-2 is available at: [/examples/krea2/](/examples/krea2/)
+
+| Model ID | Inference | Low VRAM Inference | Full Training | Full Training Validation | LoRA Training | LoRA Training Validation |
+|-|-|-|-|-|-|-|
+|[krea/Krea-2-Raw](https://www.modelscope.cn/models/krea/Krea-2-Raw)|[code](/examples/krea2/model_inference/Krea-2-Raw.py)|[code](/examples/krea2/model_inference_low_vram/Krea-2-Raw.py)|[code](/examples/krea2/model_training/full/Krea-2-Raw.sh)|[code](/examples/krea2/model_training/validate_full/Krea-2-Raw.py)|[code](/examples/krea2/model_training/lora/Krea-2-Raw.sh)|[code](/examples/krea2/model_training/validate_lora/Krea-2-Raw.py)|
+|[krea/Krea-2-Turbo](https://www.modelscope.cn/models/krea/Krea-2-Turbo)|[code](/examples/krea2/model_inference/Krea-2-Turbo.py)|[code](/examples/krea2/model_inference_low_vram/Krea-2-Turbo.py)|[code](/examples/krea2/model_training/full/Krea-2-Turbo.sh)|[code](/examples/krea2/model_training/validate_full/Krea-2-Turbo.py)|[code](/examples/krea2/model_training/lora/Krea-2-Turbo.sh)|[code](/examples/krea2/model_training/validate_lora/Krea-2-Turbo.py)|
+
+</details>
+
 ### Video Synthesis
 
 https://github.com/user-attachments/assets/1d66ae74-3b02-40a9-acc3-ea95fc039314

{diffsynth-2.0.14 → diffsynth-2.0.16}/diffsynth/configs/model_configs.py

@@ -917,6 +917,14 @@ stable_diffusion_xl_series = [
         "state_dict_converter": "diffsynth.utils.state_dict_converters.stable_diffusion_vae.SDVAEStateDictConverter",
         "extra_kwargs": {"scaling_factor": 0.13025, "sample_size": 1024, "force_upcast": True},
     },
+    {
+        # Example: ModelConfig(model_id="stabilityai/stable-diffusion-xl-base-1.0", origin_file_pattern="sd_xl_base_1.0.safetensors")
+        "model_hash": "4cf64a799d04260df438c6f33c9a047e",
+        "model_name": "stable_diffusion_xl_unet",
+        "model_class": "diffsynth.models.stable_diffusion_xl_unet.SDXLUNet2DConditionModel",
+        "extra_kwargs": {"attention_head_dim": [5, 10, 20], "transformer_layers_per_block": [1, 2, 10], "use_linear_projection": True, "addition_embed_type": "text_time", "addition_time_embed_dim": 256, "projection_class_embeddings_input_dim": 2816},
+        "state_dict_converter": "diffsynth.utils.state_dict_converters.sdxl.SDXLUNetStateDictConverter_Original2Diffusers",
+    }
 ]
 
 stable_diffusion_series = [
@@ -1163,8 +1171,38 @@ ideogram4_series = [
     },
 ]
 
+krea2_series = [
+    {
+        # Example: ModelConfig(model_id="AI-ModelScope/Krea-2", origin_file_pattern="raw.safetensors"),
+        "model_hash": "4be51f0ecc9b9dfc9825abb3f7a0b4ae",
+        "model_name": "krea2_dit",
+        "model_class": "diffsynth.models.krea2_dit.SingleStreamDiT",
+        "state_dict_converter": "diffsynth.utils.state_dict_converters.krea2_dit.Krea2DiTStateDictConverter",
+    },
+    {
+        # Example: ModelConfig(model_id="krea/Krea-2-Raw", origin_file_pattern="raw.safetensors")
+        "model_hash": "2a05e69084a6d608516ee47c8ad03334",
+        "model_name": "krea2_dit",
+        "model_class": "diffsynth.models.krea2_dit.SingleStreamDiT",
+    },
+    {
+        # Example: ModelConfig(model_id="krea/Krea-2-Raw", origin_file_pattern="transformer/*.safetensors")
+        "model_hash": "cddb3118d9f8250bf7b08c331c5c55aa",
+        "model_name": "krea2_dit",
+        "model_class": "diffsynth.models.krea2_dit.SingleStreamDiT",
+        "state_dict_converter": "diffsynth.utils.state_dict_converters.krea2_dit.Krea2DiTStateDictConverter_Diffusers",
+    },
+    {
+        # Example: ModelConfig(model_id="Qwen/Qwen3-VL-4B-Instruct", origin_file_pattern="*.safetensors")
+        "model_hash": "b8750c24f732c87797f551196c4cef78",
+        "model_name": "krea2_text_encoder",
+        "model_class": "diffsynth.models.krea2_text_encoder.Krea2TextEncoder",
+        "state_dict_converter": "diffsynth.utils.state_dict_converters.krea2_text_encoder.Krea2TextEncoderStateDictConverter",
+    },
+]
+
 MODEL_CONFIGS = (
     stable_diffusion_xl_series + stable_diffusion_series + qwen_image_series + wan_series + flux_series + flux2_series + ernie_image_series
     + z_image_series + ltx2_series + anima_series + mova_series + joyai_image_series + ace_step_series + hidream_o1_image_series
-    + image_metrics_series + ideogram4_series
+    + image_metrics_series + ideogram4_series + krea2_series
 )

{diffsynth-2.0.14 → diffsynth-2.0.16}/diffsynth/configs/vram_management_module_maps.py

@@ -393,6 +393,18 @@ VRAM_MANAGEMENT_MODULE_MAPS = {
         "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
         "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
     },
+    "diffsynth.models.krea2_dit.SingleStreamDiT": {
+        "diffsynth.models.krea2_dit.RMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+        "diffsynth.models.krea2_dit.DoubleSharedModulation": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "diffsynth.models.krea2_dit.SimpleModulation": "diffsynth.core.vram.layers.AutoWrappedModule",
+    },
+    "diffsynth.models.krea2_text_encoder.Krea2TextEncoder": {
+        "torch.nn.Linear": "diffsynth.core.vram.layers.AutoWrappedLinear",
+        "torch.nn.Embedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextRotaryEmbedding": "diffsynth.core.vram.layers.AutoWrappedModule",
+        "transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextRMSNorm": "diffsynth.core.vram.layers.AutoWrappedModule",
+    },
 }
 
 def QwenImageTextEncoder_Module_Map_Updater():

{diffsynth-2.0.14 → diffsynth-2.0.16}/diffsynth/core/data/operators.py

@@ -1,10 +1,32 @@
-import math, warnings
+import math, warnings, time
 import torch, torchvision, imageio, os
 import imageio.v3 as iio
+import urllib.request
+from io import BytesIO
 from PIL import Image
 from einops import repeat
 
 
+def _is_url(path):
+    return isinstance(path, str) and path.startswith(("http://", "https://"))
+
+
+def _load_image_from_url(url, timeout=30, max_retries=3):
+    request = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0 (compatible; DiffSynth-Studio)"})
+    last_error = None
+    for attempt in range(max_retries):
+        try:
+            with urllib.request.urlopen(request, timeout=timeout) as response:
+                image = Image.open(BytesIO(response.read()))
+                image.load()
+            return image
+        except Exception as e:
+            last_error = e
+            if attempt < max_retries - 1:
+                time.sleep(2 ** attempt)
+    raise RuntimeError(f"Failed to load image from URL '{url}' after {max_retries} attempts: {last_error}") from last_error
+
+
 class DataProcessingPipeline:
     def __init__(self, operators=None):
         self.operators: list[DataProcessingOperator] = [] if operators is None else operators
@@ -55,12 +77,17 @@ class ToStr(DataProcessingOperator):
 
 
 class LoadImage(DataProcessingOperator):
-    def __init__(self, convert_RGB=True, convert_RGBA=False):
+    def __init__(self, convert_RGB=True, convert_RGBA=False, url_timeout=30, url_max_retries=3):
         self.convert_RGB = convert_RGB
         self.convert_RGBA = convert_RGBA
-    
+        self.url_timeout = url_timeout
+        self.url_max_retries = url_max_retries
+
     def __call__(self, data: str):
-        image = Image.open(data)
+        if _is_url(data):
+            image = _load_image_from_url(data, timeout=self.url_timeout, max_retries=self.url_max_retries)
+        else:
+            image = Image.open(data)
         if self.convert_RGB: image = image.convert("RGB")
         if self.convert_RGBA: image = image.convert("RGBA")
         return image
@@ -242,6 +269,9 @@ class ToAbsolutePath(DataProcessingOperator):
         self.base_path = base_path
         
     def __call__(self, data):
+        # A remote URL is already absolute; joining it with base_path would corrupt it.
+        if _is_url(data):
+            return data
         return os.path.join(self.base_path, data)
 
 

{diffsynth-2.0.14 → diffsynth-2.0.16}/diffsynth/diffusion/ddim_scheduler.py

@@ -87,8 +87,9 @@ class DDIMScheduler():
     
     
     def add_noise(self, original_samples, noise, timestep):
-        sqrt_alpha_prod = math.sqrt(self.alphas_cumprod[int(timestep.flatten().tolist()[0])])
-        sqrt_one_minus_alpha_prod = math.sqrt(1 - self.alphas_cumprod[int(timestep.flatten().tolist()[0])])
+        timestep_id = max(min(int(timestep.flatten().tolist()[0]), len(self.alphas_cumprod)-1), 0)
+        sqrt_alpha_prod = math.sqrt(self.alphas_cumprod[timestep_id])
+        sqrt_one_minus_alpha_prod = math.sqrt(1 - self.alphas_cumprod[timestep_id])
         noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
         return noisy_samples
     
@@ -97,8 +98,9 @@ class DDIMScheduler():
         if self.prediction_type == "epsilon":
             return noise
         else:
-            sqrt_alpha_prod = math.sqrt(self.alphas_cumprod[int(timestep.flatten().tolist()[0])])
-            sqrt_one_minus_alpha_prod = math.sqrt(1 - self.alphas_cumprod[int(timestep.flatten().tolist()[0])])
+            timestep_id = max(min(int(timestep.flatten().tolist()[0]), len(self.alphas_cumprod)-1), 0)
+            sqrt_alpha_prod = math.sqrt(self.alphas_cumprod[timestep_id])
+            sqrt_one_minus_alpha_prod = math.sqrt(1 - self.alphas_cumprod[timestep_id])
             target = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
             return target
         

{diffsynth-2.0.14 → diffsynth-2.0.16}/diffsynth/diffusion/flow_match.py

@@ -4,7 +4,7 @@ from typing_extensions import Literal
 
 class FlowMatchScheduler():
 
-    def __init__(self, template: Literal["FLUX.1", "Wan", "Qwen-Image", "FLUX.2", "Z-Image", "LTX-2", "Qwen-Image-Lightning", "ERNIE-Image", "ACE-Step", "Ideogram4"] = "FLUX.1"):
+    def __init__(self, template: Literal["FLUX.1", "Wan", "Qwen-Image", "FLUX.2", "Z-Image", "LTX-2", "Qwen-Image-Lightning", "ERNIE-Image", "ACE-Step", "Ideogram4", "Krea-2"] = "FLUX.1"):
         self.set_timesteps_fn = {
             "FLUX.1": FlowMatchScheduler.set_timesteps_flux,
             "Wan": FlowMatchScheduler.set_timesteps_wan,
@@ -17,6 +17,7 @@ class FlowMatchScheduler():
             "ACE-Step": FlowMatchScheduler.set_timesteps_ace_step,
             "HiDream-O1-Image": FlowMatchScheduler.set_timesteps_hidream_o1_image,
             "Ideogram4": FlowMatchScheduler.set_timesteps_ideogram4,
+            "Krea-2": FlowMatchScheduler.set_timesteps_krea2,
         }.get(template, FlowMatchScheduler.set_timesteps_flux)
         self.num_train_timesteps = 1000
 
@@ -233,6 +234,23 @@ class FlowMatchScheduler():
         sigmas = (1 - sigmas)[:-1]
         return sigmas, timesteps
 
+    @staticmethod
+    def set_timesteps_krea2(num_inference_steps=28, denoising_strength=1.0, dynamic_shift_len=None, y1=0.5, y2=1.15, mu=None):
+        x1 = 256
+        x2 = 6400
+        sigma = 1
+        ts = torch.linspace(denoising_strength, 0, num_inference_steps + 1)
+        if mu is None and dynamic_shift_len is None:
+            # Training
+            mu = 0.8
+        elif mu is None:
+            # Raw
+            slope = (y2 - y1) / (x2 - x1)
+            mu = slope * dynamic_shift_len + (y1 - slope * x1)
+        ts = math.exp(mu) / (math.exp(mu) + (1.0 / ts - 1.0) ** sigma)
+        sigmas, timesteps = ts[:-1], ts[:-1]
+        return sigmas, timesteps
+
     @staticmethod
     def set_timesteps_ltx2(num_inference_steps=100, denoising_strength=1.0, dynamic_shift_len=None, terminal=0.1, special_case=None):
         num_train_timesteps = 1000

{diffsynth-2.0.14 → diffsynth-2.0.16}/diffsynth/diffusion/runner.py

@@ -1,4 +1,4 @@
-import os, torch, importlib
+import os, json, torch, importlib
 from tqdm import tqdm
 from accelerate import Accelerator
 from .training_module import DiffusionTrainingModule
@@ -16,6 +16,20 @@ def get_optimizer_class(customized_optimizer=None):
         return getattr(module, class_name)
 
 
+def save_training_args(args):
+    output_path = getattr(args, "output_path", None) if args is not None else None
+    if output_path is None:
+        return
+    try:
+        os.makedirs(args.output_path, exist_ok=True)
+        save_path = os.path.join(args.output_path, "training_args.json")
+        with open(save_path, "w", encoding="utf-8") as f:
+            json.dump(vars(args), f, indent=4, ensure_ascii=False, default=str)
+        print(f"Training arguments saved to `{save_path}`.")
+    except Exception as e:
+        print(f"Warning: failed to save training arguments: {e}")
+
+
 def launch_training_task(
     accelerator: Accelerator,
     dataset: torch.utils.data.Dataset,
@@ -44,6 +58,9 @@ def launch_training_task(
         cpu_offload_split_threshold = args.cpu_offload_split_threshold
         customized_optimizer = args.customized_optimizer
 
+    if accelerator.is_main_process:
+        save_training_args(args)
+
     optimizer_class = get_optimizer_class(customized_optimizer)
     optimizer = optimizer_class(model.trainable_modules(), lr=learning_rate, weight_decay=weight_decay)
     scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer)