add training (lora, foundation, data)

qwenimage/datamodels.py

@@ -0,0 +1,40 @@
+
+
+
+from diffusers.image_processor import PipelineImageInput
+from pydantic import BaseModel, ConfigDict, Field
+import torch
+
+from wandml.foundation.datamodels import FluxInputs
+from wandml.trainers.datamodels import ExperimentTrainerParameters
+
+
+class QwenInputs(BaseModel):
+    image: PipelineImageInput | None = None
+    prompt: str| list[str] | None = None
+    height: int|None = None
+    width: int|None = None
+    negative_prompt: str| list[str] | None = None
+    true_cfg_scale: float = 1.0
+    num_inference_steps: int = 50
+    generator: torch.Generator | list[torch.Generator] | None = None
+    max_sequence_length: int = 512
+    vae_image_override: int | None = 512 * 512
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        # extra="allow",
+    )
+
+
+class QwenConfig(ExperimentTrainerParameters):
+    load_multi_view_lora: bool = False
+    train_max_sequence_length: int = 512
+    train_dist: str = "linear" # "logit-normal"
+    train_shift: bool = True
+    inference_dist: str = "linear"
+    inference_shift: bool = True
+    static_mu: float | None = None
+    loss_weight_dist: str | None = None # "scaled_clipped_gaussian", "logit-normal"
+
+    vae_image_size: int = 1024 * 1024

qwenimage/datasets.py

@@ -0,0 +1,56 @@
+
+
+from pathlib import Path
+import random
+
+from PIL import Image
+from wandml.core.datamodels import SourceDataType
+from wandml.core.source import Source
+
+class StyleSource(Source):
+    _data_types = [
+        SourceDataType(name="image", type=Image.Image),
+        SourceDataType(name="text", type=str),
+    ]
+    def __init__(self, data_dir, prompt, set_len=None):
+        data_dir = Path(data_dir)
+        self.images = list(data_dir.iterdir())
+        self.prompt = prompt
+        self.set_len = set_len
+    
+    def __len__(self):
+        if self.set_len is not None:
+            return self.set_len
+        else:
+            return len(self.images)
+    
+    def __getitem__(self, idx):
+        idx = idx % len(self.images)
+        im_pil = Image.open(self.images[idx]).convert("RGB")
+        return im_pil, self.prompt
+
+class StyleSourceWithRandomRef(Source):
+    _data_types = [
+        SourceDataType(name="image", type=Image.Image),
+        SourceDataType(name="text", type=str),
+        SourceDataType(name="reference", type=Image.Image),
+    ]
+    def __init__(self, data_dir, prompt, ref_dir, set_len=None):
+        data_dir = Path(data_dir)
+        self.images = list(data_dir.iterdir())
+        self.ref_images = list(Path(ref_dir).iterdir())
+        self.prompt = prompt
+        self.set_len = set_len
+    
+    def __len__(self):
+        if self.set_len is not None:
+            return self.set_len
+        else:
+            return len(self.images)
+    
+    def __getitem__(self, idx):
+        idx = idx % len(self.images)
+        im_pil = Image.open(self.images[idx]).convert("RGB")
+        rand_ref = random.choice(self.ref_images)
+        ref_pil = Image.open(rand_ref).convert("RGB")
+        return im_pil, self.prompt, ref_pil

qwenimage/finetuner.py

@@ -0,0 +1,60 @@
+import os
+from typing import Optional
+import uuid
+import hashlib
+
+from peft import LoraConfig
+
+from wandml.utils.debug import ftimed
+from wandml.finetune.lora.lora import LoraFinetuner
+
+class QwenLoraFinetuner(LoraFinetuner):
+
+    @ftimed
+    def load(self, load_path, lora_rank=16, lora_config:Optional[LoraConfig]=None):
+        """
+        Loads new lora on flux transformer if not loaded. Loads lora safetensors from load_path. Specify specific lora config using lora_rank or lora_config.
+        """
+        if "transformer" in self.modules:
+            return super().load(load_path)
+        
+        if lora_config:
+            self.foundation.transformer = self.add_module(
+                "transformer",
+                self.foundation.transformer,
+                lora_config=lora_config
+            )
+        else:
+            target_modules = [
+                'to_q',
+                'to_k',
+                'to_v',
+                'to_qkv',
+
+                'add_q_proj',
+                'add_k_proj',
+                'add_v_proj',
+                'to_added_qkv',
+
+                'proj',
+                'txt_in',
+                'img_in',
+                'txt_mod.1',
+                'img_mod.1',
+                'proj_out',
+                'to_add_out',
+                'to_out.0'
+                'net.2',
+                'linear',
+                'linear_2',
+                'linear_1',
+            ]
+            self.foundation.transformer = self.add_module(
+                "transformer",
+                self.foundation.transformer,
+                target_modules=target_modules,
+                lora_rank=lora_rank,
+            )
+            self.foundation.transformer.to(dtype=self.foundation.dtype)
+
+        return super().load(load_path)

qwenimage/foundation.py

@@ -0,0 +1,277 @@
+
+import os
+from pathlib import Path
+import warnings
+
+from PIL import Image
+from diffusers.pipelines.qwenimage.pipeline_qwenimage import QwenImagePipeline
+import torch
+from safetensors.torch import load_file, save_model
+import torch.nn.functional as F
+from einops import rearrange
+
+from qwenimage.datamodels import QwenConfig, QwenInputs
+from qwenimage.debug import print_gpu_memory
+from qwenimage.models.encode_prompt import encode_prompt
+from qwenimage.models.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
+from qwenimage.models.transformer_qwenimage import QwenImageTransformer2DModel
+from qwenimage.sampling import TimestepDistUtils
+from wandml import WandModel
+
+
+class QwenImageFoundation(WandModel):
+    SOURCE = "Qwen/Qwen-Image-Edit-2509"
+    INPUT_MODEL = QwenInputs
+    CACHE_DIR = "qwen_image_edit_2509"
+    PIPELINE = QwenImageEditPlusPipeline
+
+    serialize_modules = ["transformer"]
+
+    def __init__(self, config:QwenConfig, device=None):
+        super().__init__()
+        self.config:QwenConfig = config
+        self.dtype = torch.bfloat16
+        if device is None:
+            default_device = "cuda" if torch.cuda.is_available() else "cpu"
+            self.device = default_device
+        else:
+            self.device = device
+        print(f"{self.device=}")
+
+        pipe = self.PIPELINE.from_pretrained(
+            "Qwen/Qwen-Image-Edit-2509", 
+            transformer=QwenImageTransformer2DModel.from_pretrained(
+                "Qwen/Qwen-Image-Edit-2509",
+                subfolder='transformer',
+                torch_dtype=self.dtype,
+                device_map=self.device
+            ),
+            torch_dtype=self.dtype,
+        )
+        pipe = pipe.to(device=self.device, dtype=self.dtype)
+
+
+        if config.load_multi_view_lora:
+            pipe.load_lora_weights(
+                "dx8152/Qwen-Edit-2509-Multiple-angles", 
+                weight_name="镜头转换.safetensors", adapter_name="angles"
+            )
+            pipe.set_adapters(["angles"], adapter_weights=[1.])
+            pipe.fuse_lora(adapter_names=["angles"], lora_scale=1.25)
+            pipe.unload_lora_weights()
+        
+        self.pipe = pipe
+        self.vae = self.pipe.vae
+        self.transformer = self.pipe.transformer
+        self.text_encoder = self.pipe.text_encoder
+        self.scheduler = self.pipe.scheduler
+
+        self.vae.to(self.dtype)
+        self.vae.eval()
+        self.vae.requires_grad_(False)
+        self.text_encoder.eval()
+        self.text_encoder.requires_grad_(False)
+        self.text_encoder_device = None
+        self.transformer.eval()
+        self.transformer.requires_grad_(False)
+
+        self.timestep_dist_utils = TimestepDistUtils(
+            min_seq_len=self.scheduler.config.base_image_seq_len,
+            max_seq_len=self.scheduler.config.max_image_seq_len,
+            min_mu=self.scheduler.config.base_shift,
+            max_mu=self.scheduler.config.max_shift,
+            train_dist=self.config.train_dist,
+            train_shift=self.config.train_shift,
+            inference_dist=self.config.inference_dist,
+            inference_shift=self.config.inference_shift,
+            static_mu=self.config.static_mu,
+            loss_weight_dist=self.config.loss_weight_dist,
+        )
+        self.static_prompt_embeds = None
+
+    def load(self, load_path):
+        if not isinstance(load_path, Path): 
+            load_path = Path(load_path)
+        if not load_path.is_dir():
+            raise ValueError(f"Expected {load_path=} to be a directory")
+        for module_name in self.serialize_modules:
+            model_state_dict = load_file(load_path / f"{module_name}.safetensors")
+            missing, unexpected = getattr(self, module_name).load_state_dict(model_state_dict, strict=False, assign=True)
+            if missing: 
+                warnings.warn(f"{module_name} missing {missing}")
+            if unexpected: 
+                warnings.warn(f"{module_name} unexpected {unexpected}")
+
+    def save(self, save_path, skip=False):
+        if skip: return
+
+        if not isinstance(save_path, Path): 
+            save_path = Path(save_path)
+        if not save_path.is_dir():
+            raise ValueError(f"Expected {save_path=} to be a directory")
+
+        save_path.mkdir(parents=True, exist_ok=True)
+
+        for module_name in self.serialize_modules:
+            save_model(getattr(self, module_name), save_path / f"{module_name}.safetensors")
+            print(f"Saved {module_name} to {save_path}")
+    
+    def get_train_params(self):
+        return [{"params": [p for p in self.transformer.parameters() if p.requires_grad]}]
+    
+    def pil_to_latents(self, images):
+        image = self.pipe.image_processor.preprocess(images)
+        image = image.unsqueeze(2) # N, C, F=1, H, W
+        image = image.to(device=self.device, dtype=self.dtype)
+        latents = self.pipe.vae.encode(image).latent_dist.mode() # argmax
+
+        latents_mean = (
+            torch.tensor(self.pipe.vae.config.latents_mean)
+            .view(1, self.pipe.vae.config.z_dim, 1, 1, 1)
+            .to(latents.device, latents.dtype)
+        )
+        latents_std = (
+            torch.tensor(self.pipe.vae.config.latents_std)
+            .view(1, self.pipe.vae.config.z_dim, 1, 1, 1)
+            .to(latents.device, latents.dtype)
+        )
+        latents = (latents - latents_mean) / latents_std
+        latents = latents.squeeze(2)
+        return latents.to(dtype=self.dtype)
+
+    def latents_to_pil(self, latents):
+        latents = latents.clone().detach()
+        latents = latents.unsqueeze(2)
+
+        latents = latents.to(self.dtype)
+        latents_mean = (
+            torch.tensor(self.pipe.vae.config.latents_mean)
+            .view(1, self.pipe.vae.config.z_dim, 1, 1, 1)
+            .to(latents.device, latents.dtype)
+        )
+        latents_std = (
+            torch.tensor(self.pipe.vae.config.latents_std)
+            .view(1, self.pipe.vae.config.z_dim, 1, 1, 1)
+            .to(latents.device, latents.dtype)
+        )
+        latents = latents * latents_std + latents_mean
+
+        latents = latents.to(device=self.device, dtype=self.dtype)
+        image = self.pipe.vae.decode(latents, return_dict=False)[0][:, :, 0] # F = 1
+        image = self.pipe.image_processor.postprocess(image)
+        return image
+
+    @staticmethod
+    def pack_latents(latents):
+        packed = rearrange(latents, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+        return packed
+    
+    @staticmethod
+    def unpack_latents(packed, h, w):
+        latents = rearrange(packed, "b (h w) (c ph pw) -> b c (h ph) (w pw)", ph=2, pw=2, h=h, w=w)
+        return latents
+
+    def set_static_prompt(self, prompt:str):
+        self.text_encoder.to(device=self.device)
+        if self.text_encoder_device != "cuda":
+            self.text_encoder_device = "cuda"
+        with torch.no_grad():
+            prompt_embeds, prompt_embeds_mask = encode_prompt(
+                self.text_encoder,
+                self.pipe.tokenizer,
+                prompt,
+                device=self.device,
+                dtype=self.dtype,
+                max_sequence_length = self.config.train_max_sequence_length,
+            )
+        prompt_embeds = prompt_embeds.cpu().clone().detach()
+        prompt_embeds_mask = prompt_embeds_mask.cpu().clone().detach()
+        self.static_prompt_embeds = (prompt_embeds, prompt_embeds_mask)
+
+
+    def preprocess_batch(self, batch):
+        prompts = batch["text"]
+
+        if self.static_prompt_embeds is not None:
+            prompt_embeds, prompt_embeds_mask = self.static_prompt_embeds
+
+        self.text_encoder.to(device=self.device)
+        if self.text_encoder_device != "cuda":
+            self.text_encoder_device = "cuda"
+
+        with torch.no_grad():
+            prompt_embeds, prompt_embeds_mask = encode_prompt(
+                self.text_encoder,
+                self.pipe.tokenizer,
+                prompts,
+                device=self.device,
+                dtype=self.dtype,
+                max_sequence_length = self.config.train_max_sequence_length,
+            )
+        prompt_embeds = prompt_embeds.cpu().clone().detach()
+        prompt_embeds_mask = prompt_embeds_mask.cpu().clone().detach()
+
+
+        batch["prompt_embeds"] = (prompt_embeds, prompt_embeds_mask)
+
+        return batch
+
+    def single_step(self, batch) -> torch.Tensor:
+        self.text_encoder.to(device="cpu") # offload
+        if self.text_encoder_device != "cpu":
+            self.text_encoder_device = "cpu"
+            print_gpu_memory()
+
+        if "prompt_embeds" not in batch:
+            batch = self.preprocess_batch(batch)
+        prompt_embeds, prompt_embeds_mask = batch["prompt_embeds"]
+        prompt_embeds = prompt_embeds.to(device=self.device)
+        prompt_embeds_mask = prompt_embeds_mask.to(device=self.device)
+
+        images = batch["image"]
+        x_0 = self.pil_to_latents(images).to(device=self.device, dtype=self.dtype)
+        x_1 = torch.randn_like(x_0).to(device=self.device, dtype=self.dtype)
+        seq_len = self.timestep_dist_utils.get_seq_len(x_0)
+        batch_size = x_0.shape[0]
+        t = self.timestep_dist_utils.get_train_t([batch_size], seq_len=seq_len).to(device=self.device, dtype=self.dtype)
+        x_t = (1.0 - t) * x_0 + t * x_1
+
+        x_t_1d = self.pack_latents(x_t)
+
+        l_height, l_width = x_0.shape[-2:]
+        img_shapes = [
+            [(1, l_height // 2, l_width // 2), ]
+        ] * batch_size
+        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist()
+        image_rotary_emb = self.transformer.pos_embed(img_shapes, txt_seq_lens, device=x_0.device)
+
+        v_pred_1d = self.transformer(
+            hidden_states=x_t_1d,
+            encoder_hidden_states=prompt_embeds,
+            encoder_hidden_states_mask=prompt_embeds_mask,
+            timestep=t,
+            image_rotary_emb=image_rotary_emb,
+            return_dict=False,
+        )[0]
+
+        v_pred_2d = self.unpack_latents(v_pred_1d, h=l_height//2, w=l_width//2)
+        v_gt_2d = x_1 - x_0
+
+        if self.config.loss_weight_dist is not None:
+            loss = F.mse_loss(v_pred_2d, v_gt_2d, reduction="none").mean(dim=[1,2,3])
+            weights = self.timestep_dist_utils.get_loss_weighting(t)
+            loss = torch.mean(loss * weights)
+        else:
+            loss = F.mse_loss(v_pred_2d, v_gt_2d, reduction="mean")
+        
+        return loss
+
+
+    def base_pipe(self, inputs: QwenInputs) -> list[Image]:
+        self.text_encoder.to(device=self.device)
+        if self.text_encoder_device != "cuda":
+            self.text_encoder_device = "cuda"
+        return self.pipe(**inputs.model_dump()).images
+    
+
+

qwenimage/models/encode_prompt.py

@@ -0,0 +1,59 @@
+import torch
+from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer
+
+tokenizer_max_length = 1024
+prompt_template_encode = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+prompt_template_encode_start_idx = 34
+default_sample_size = 128
+
+def _extract_masked_hidden(hidden_states: torch.Tensor, mask: torch.Tensor):
+    bool_mask = mask.bool()
+    valid_lengths = bool_mask.sum(dim=1)
+    selected = hidden_states[bool_mask]
+    split_result = torch.split(selected, valid_lengths.tolist(), dim=0)
+
+    return split_result
+
+
+def encode_prompt(
+    text_encoder: Qwen2_5_VLForConditionalGeneration,
+    tokenizer: Qwen2Tokenizer,
+    prompt: str | list[str],
+    device: torch.device | None = None,
+    dtype: torch.dtype | None = None,
+    max_sequence_length: int|None = 1024,
+):
+
+
+    prompt = [prompt] if isinstance(prompt, str) else prompt
+
+    template = prompt_template_encode
+    drop_idx = prompt_template_encode_start_idx
+    txt = [template.format(e) for e in prompt]
+    txt_tokens = tokenizer(
+        txt, max_length=tokenizer_max_length + drop_idx, padding=True, truncation=True, return_tensors="pt"
+    ).to(device)
+    encoder_hidden_states = text_encoder(
+        input_ids=txt_tokens.input_ids,
+        attention_mask=txt_tokens.attention_mask,
+        output_hidden_states=True,
+    )
+    hidden_states = encoder_hidden_states.hidden_states[-1]
+    split_hidden_states = _extract_masked_hidden(hidden_states, txt_tokens.attention_mask)
+    split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
+    attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states]
+    max_seq_len = max([e.size(0) for e in split_hidden_states])
+    prompt_embeds = torch.stack(
+        [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]
+    )
+    prompt_embeds_mask = torch.stack(
+        [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]
+    )
+
+    prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+    if max_sequence_length is not None:
+        prompt_embeds = prompt_embeds[:, :max_sequence_length]
+        prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
+
+    return prompt_embeds, prompt_embeds_mask

qwenimage/models/pipeline_qwenimage_edit_plus.py

@@ -763,6 +763,7 @@ class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
             # 5. Prepare timesteps
             sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
             image_seq_len = latents.shape[1]
+            print(f"{image_seq_len=}")
             mu = calculate_shift(
                 image_seq_len,
                 self.scheduler.config.get("base_image_seq_len", 256),
@@ -770,6 +771,7 @@ class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
                 self.scheduler.config.get("base_shift", 0.5),
                 self.scheduler.config.get("max_shift", 1.15),
             )
+            print(f"{mu=}")
             timesteps, num_inference_steps = retrieve_timesteps(
                 self.scheduler,
                 num_inference_steps,
@@ -777,6 +779,7 @@ class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
                 sigmas=sigmas,
                 mu=mu,
             )
+            print(f"{timesteps=}")
             num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
             self._num_timesteps = len(timesteps)
 

qwenimage/sampling.py

@@ -0,0 +1,146 @@
+
+
+
+# math functions for sampling schedule
+import math
+from typing import Callable, Literal
+
+import torch
+
+
+
+class TimestepDistUtils:
+    
+    @staticmethod
+    def t_shift(mu: float, sigma: float, t: torch.Tensor):
+        """
+            see eq.(12) of https://arxiv.org/abs/2506.15742 Black Forest Labs (2025)
+            t' = \frac{e^{\mu}}{e^{\mu} + (1/t - 1)^{\sigma}}
+        """
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+    @staticmethod
+    def lerp_mu( # qwen params
+        seq_len,
+        min_seq_len: int = 256,
+        max_seq_len: int = 8192,
+        min_mu: float = 0.5, 
+        max_mu: float = 0.9,
+        train_dist: str = "linear",
+    ):
+        """
+        Resolution-dependent shifting of timestep schedules
+        from Esser et al. https://arxiv.org/abs/2403.03206 
+        updated with default params for Qwen
+        """
+        m = (max_mu - min_mu) / (max_seq_len - min_seq_len)
+        b = min_mu - m * min_seq_len
+        mu = seq_len * m + b
+        return mu
+
+    @staticmethod
+    def logit_normal(t, mu=0.0, sigma=1.0):
+        """
+            Logit normal PDF, as in logistic(randn(mu, sigma))
+        """
+        pdf = torch.zeros_like(t)
+        z = (torch.logit(t) - mu) / sigma
+        coef = 1.0 / (sigma * torch.sqrt(torch.tensor(2.0 * torch.pi)))
+        pdf = coef * torch.exp(-0.5 * z**2) / (t * (1.0 - t))
+        return pdf
+    
+    @staticmethod
+    def scaled_clipped_gaussian(t):
+        """
+            Heuristic distribution for gaussian wuth mu = 0.5 and sigma=0.5, 
+            clipped to [0,1], with int_0^1dt =1.0
+        """
+        y = torch.exp(-2 * (t - 0.5) ** 2)
+        y = (y - 0.606) * 4.02
+        return y
+    
+    @staticmethod
+    def get_seq_len(latents):
+        if latents.dim() == 4 or latents.dim() == 5:
+            h,w = latents.shape[-2:]
+            seq_len = (h//2)*(w//2)
+        elif latents.dim() == 3:
+            seq_len = latents.shape[1] # [B, L=h*w, C]
+        else:
+            raise ValueError(f"{latents.dim()=} not in 3,4,5")
+        return seq_len
+
+    def __init__(
+        self,
+        min_seq_len=256,
+        max_seq_len=8192,
+        min_mu=0.5,
+        max_mu=0.9,
+        train_dist:Literal["logit-normal", "linear"]="linear",
+        train_shift:bool=True,
+        inference_dist:Literal["logit-normal", "linear"]="linear",
+        inference_shift:bool=True,
+        static_mu:float|None=None,
+        loss_weight_dist: Literal["scaled_clipped_gaussian", "logit-normal"] | None = None,
+    ):
+        self.min_seq_len = min_seq_len
+        self.max_seq_len = max_seq_len
+        self.min_mu = min_mu
+        self.max_mu = max_mu
+        self.train_dist = train_dist
+        self.train_shift = train_shift
+        self.inference_dist = inference_dist
+        self.inference_shift = inference_shift
+        self.static_mu = static_mu
+        self.loss_weight_dist = loss_weight_dist
+
+    def lin_t_to_dist(self, t, seq_len=None):
+        if self.train_dist == "logit-normal":
+            t = self.logit_normal_pdf(t)
+        elif self.train_dist == "linear":
+            pass
+        else:
+            raise ValueError()
+    
+        if self.train_shift:
+            if self.static_mu:
+                mu = self.static_mu
+            elif seq_len:
+                mu = self.lerp_mu(seq_len, self.min_seq_len, self.max_seq_len, self.min_mu, self.max_mu)
+            else:
+                raise ValueError()
+            t = self.t_shift(mu, 1.0, t)
+        return t
+
+    def get_train_t(self, size, seq_len=None):
+        t = torch.rand(size)
+        t = self.lin_t_to_dist(t, seq_len=seq_len)
+        return t
+
+    def get_loss_weighting(self, t):
+        if self.loss_weight_dist == "scaled_clipped_gaussian":
+            w = self.scaled_clipped_gaussian(t)
+        elif self.loss_weight_dist == "logit-normal":
+            w = self.logit_normal_pdf(t)
+        elif self.loss_weight_dist is None:
+            w = torch.ones_like(t)
+        else:
+            raise ValueError()
+        return w
+
+
+    def get_inference_t(self, steps, strength=1.0, seq_len=None, clip_by_strength=True):
+        if clip_by_strength:
+            true_steps = max(1, int(strength * steps)) + 1
+        else:
+            true_steps = max(1, steps) + 1
+        t = torch.linspace(strength, 0.0, true_steps)
+        t = self.lin_t_to_dist(t, seq_len=seq_len)
+        return t
+
+    def inference_ode_step(self, noise_pred: torch.Tensor, latents: torch.Tensor, index: int, t_schedule: torch.Tensor):
+        t = t_schedule[index]
+        t_next = t_schedule[index + 1]
+        d_t = t_next - t
+        latents = latents + d_t * noise_pred
+        return latents

qwenimage/task.py

@@ -0,0 +1,46 @@
+from PIL import Image
+import torch
+from torchvision.transforms import v2 as T
+
+from wandml.core.datamodels import SourceDataType
+from wandml.core.task import Task
+from wandml.data.transforms import RemoveAlphaTransform, RandomDownsize
+
+
+image_transforms = T.Compose([
+    RemoveAlphaTransform(bg_color_rgb=(34, 34, 34)),
+    T.ToImage(),
+    T.RGB(),
+    RandomDownsize(sizes=(384, 512, 768)),
+    T.ToDtype(torch.float, scale=True),
+])
+
+class TextToImageTask(Task):
+    data_types = [
+        SourceDataType(name="text", type=str),
+        SourceDataType(name="image", type=Image.Image),
+    ]
+    type_transforms = [
+        Task.identity,
+        image_transforms,
+    ]
+    sample_input_dict = {
+        "prompt": SourceDataType(name="text", type=str),
+    }
+
+
+class TextToImageWithRefTask(Task):
+    data_types = [
+        SourceDataType(name="text", type=str),
+        SourceDataType(name="image", type=Image.Image),
+        SourceDataType(name="reference", type=Image.Image),
+    ]
+    type_transforms = [
+        Task.identity,
+        image_transforms,
+        image_transforms,
+    ]
+    sample_input_dict = {
+        "prompt": SourceDataType(name="text", type=str),
+        "image": SourceDataType(name="reference", type=Image.Image),
+    }

requirements.txt

@@ -10,7 +10,7 @@ dashscope
 kernels
 torchvision
 peft
-torchao==0.11.0
+torchao==0.14.1 
 pydantic
 pandas
 modal

scripts/train.ipynb

@@ -0,0 +1,529 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "faf9556d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/home/ubuntu/Qwen-Image-Edit-Angles\n"
+     ]
+    }
+   ],
+   "source": [
+    "%cd /home/ubuntu/Qwen-Image-Edit-Angles"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "d74b1b7e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/usr/lib/python3/dist-packages/sklearn/utils/fixes.py:25: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.\n",
+      "  from pkg_resources import parse_version  # type: ignore\n",
+      "2025-11-22 18:13:10.673389: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
+      "2025-11-22 18:13:10.687858: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
+      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
+      "E0000 00:00:1763835190.705243 2236633 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
+      "E0000 00:00:1763835190.710795 2236633 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
+      "W0000 00:00:1763835190.724588 2236633 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
+      "W0000 00:00:1763835190.724603 2236633 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
+      "W0000 00:00:1763835190.724605 2236633 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
+      "W0000 00:00:1763835190.724607 2236633 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
+      "2025-11-22 18:13:10.729261: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
+      "To enable the following instructions: AVX512F AVX512_VNNI AVX512_BF16 AVX512_FP16 AVX_VNNI, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
+     ]
+    },
+    {
+     "ename": "AttributeError",
+     "evalue": "'MessageFactory' object has no attribute 'GetPrototype'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
+      "\u001b[0;31mAttributeError\u001b[0m: 'MessageFactory' object has no attribute 'GetPrototype'"
+     ]
+    },
+    {
+     "ename": "AttributeError",
+     "evalue": "'MessageFactory' object has no attribute 'GetPrototype'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
+      "\u001b[0;31mAttributeError\u001b[0m: 'MessageFactory' object has no attribute 'GetPrototype'"
+     ]
+    },
+    {
+     "ename": "AttributeError",
+     "evalue": "'MessageFactory' object has no attribute 'GetPrototype'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
+      "\u001b[0;31mAttributeError\u001b[0m: 'MessageFactory' object has no attribute 'GetPrototype'"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/ubuntu/.local/lib/python3.10/site-packages/google/api_core/_python_version_support.py:266: FutureWarning: You are using a Python version (3.10.12) which Google will stop supporting in new releases of google.api_core once it reaches its end of life (2026-10-04). Please upgrade to the latest Python version, or at least Python 3.11, to continue receiving updates for google.api_core past that date.\n",
+      "  warnings.warn(message, FutureWarning)\n"
+     ]
+    },
+    {
+     "ename": "AttributeError",
+     "evalue": "'MessageFactory' object has no attribute 'GetPrototype'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
+      "\u001b[0;31mAttributeError\u001b[0m: 'MessageFactory' object has no attribute 'GetPrototype'"
+     ]
+    },
+    {
+     "ename": "AttributeError",
+     "evalue": "'MessageFactory' object has no attribute 'GetPrototype'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
+      "\u001b[0;31mAttributeError\u001b[0m: 'MessageFactory' object has no attribute 'GetPrototype'"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Skipping import of cpp extensions due to incompatible torch version 2.9.1+cu128 for torchao version 0.14.1             Please see https://github.com/pytorch/ao/issues/2919 for more info\n",
+      "TMA benchmarks will be running without grid constant TMA descriptor.\n",
+      "WARNING:bitsandbytes.cextension:Could not find the bitsandbytes CUDA binary at PosixPath('/usr/local/lib/python3.10/dist-packages/bitsandbytes/libbitsandbytes_cuda128.so')\n",
+      "ERROR:bitsandbytes.cextension:Could not load bitsandbytes native library: /lib/x86_64-linux-gnu/libstdc++.so.6: version `GLIBCXX_3.4.32' not found (required by /usr/local/lib/python3.10/dist-packages/bitsandbytes/libbitsandbytes_cpu.so)\n",
+      "Traceback (most recent call last):\n",
+      "  File \"/usr/local/lib/python3.10/dist-packages/bitsandbytes/cextension.py\", line 85, in <module>\n",
+      "    lib = get_native_library()\n",
+      "  File \"/usr/local/lib/python3.10/dist-packages/bitsandbytes/cextension.py\", line 72, in get_native_library\n",
+      "    dll = ct.cdll.LoadLibrary(str(binary_path))\n",
+      "  File \"/usr/lib/python3.10/ctypes/__init__.py\", line 452, in LoadLibrary\n",
+      "    return self._dlltype(name)\n",
+      "  File \"/usr/lib/python3.10/ctypes/__init__.py\", line 374, in __init__\n",
+      "    self._handle = _dlopen(self._name, mode)\n",
+      "OSError: /lib/x86_64-linux-gnu/libstdc++.so.6: version `GLIBCXX_3.4.32' not found (required by /usr/local/lib/python3.10/dist-packages/bitsandbytes/libbitsandbytes_cpu.so)\n",
+      "WARNING:bitsandbytes.cextension:\n",
+      "CUDA Setup failed despite CUDA being available. Please run the following command to get more information:\n",
+      "\n",
+      "python -m bitsandbytes\n",
+      "\n",
+      "Inspect the output of the command and see if you can locate CUDA libraries. You might need to add them\n",
+      "to your LD_LIBRARY_PATH. If you suspect a bug, please take the information from python -m bitsandbytes\n",
+      "and open an issue at: https://github.com/bitsandbytes-foundation/bitsandbytes/issues\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import subprocess\n",
+    "from pathlib import Path\n",
+    "import argparse\n",
+    "\n",
+    "from ruamel.yaml import YAML\n",
+    "import diffusers\n",
+    "\n",
+    "\n",
+    "from wandml.trainers.experiment_trainer import ExperimentTrainer\n",
+    "from wandml import WandDataPipe\n",
+    "import wandml\n",
+    "\n",
+    "from qwenimage.finetuner import QwenLoraFinetuner\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b7b70d58",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "ba2e8778",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "\n",
+    "from qwenimage.datasets import StyleSource\n",
+    "\n",
+    "\n",
+    "src = StyleSource(\"/data/styles-finetune-data-artistic/tarot\", \"<0001>\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "eda50bdf",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from wandml.data.tasks.text_to_image import TextToImageTask\n",
+    "\n",
+    "task = TextToImageTask()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/ubuntu/wand-ml/wandml/core/source.py:14: UserWarning: Deprecated: Use data_types instead of _data_types\n",
+      "  warnings.warn(\"Deprecated: Use data_types instead of _data_types\")\n"
+     ]
+    }
+   ],
+   "source": [
+    "dp = WandDataPipe()\n",
+    "dp.add_source(src)\n",
+    "dp.set_task(task)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "b98b9368",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c4aecbbe70e8441c8f7f7a15ff5a95f6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Fetching 7 files:   0%|          | 0/7 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "self.device='cuda'\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "94340bbd1c674a88ba287f7154815d6b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Fetching 5 files:   0%|          | 0/5 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "44febcb836804a969fd7c2571f7830be",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "379bfa34f1cb46ef9aa6c06b94a5437f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading pipeline components...:   0%|          | 0/6 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "`torch_dtype` is deprecated! Use `dtype` instead!\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "963f6022ee8c46ce9d36a02aad4bc739",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "from qwenimage.datamodels import QwenConfig\n",
+    "from qwenimage.foundation import QwenImageFoundation\n",
+    "\n",
+    "config = QwenConfig()\n",
+    "foundation = QwenImageFoundation(config=config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "7646e8ce",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Loading Lora from None\n"
+     ]
+    }
+   ],
+   "source": [
+    "finetuner = QwenLoraFinetuner(foundation, config)\n",
+    "finetuner.load(None)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "47bcba68",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "foundation=<FoundationEnum.FLUX: 'flux'> instance_data_dir=None class_data_dir=None instance_prompt=None class_prompt=None num_class_images=10 output_dir='output' seed=None size=1024 center_crop=False train_batch_size=1 num_train_epochs=1 max_train_steps=None save_steps=1000 save_path=None gradient_accumulation_steps=1 learning_rate=0.001 learning_rate_1d=1e-06 scale_lr=False lr_scheduler='constant' lr_warmup_steps=0 base_lr=1e-05 max_lr=0.001 step_size_up=2000 cyclic_lr_mode=<CyclicLRMode.TRIANGULAR2: 'triangular2'> cyclic_lr_cycle_momentum=False optim=<OptimizerType.ADAMW: 'adamw'> adam_beta1=0.9 adam_beta2=0.999 adam_weight_decay=0.01 adam_epsilon=1e-08 max_grad_norm=1.0 mixed_precision='bf16' concepts_list=None modifier_tokens=None initializer_tokens=None checkpointing_steps=9999 resume_from_checkpoint=None train_text_encoder=True gcs_bucket=None topic_id='finetune-complete' concepts=None global_step=0 prior_loss_weight=1.0 wand_user_id='test' wand_model_id='testing' wand_model_bucket='wand-finetune' wand_project_name='wand-finetune' num_sample_images=5 prodigy_beta3=None prodigy_decouple=True prodigy_use_bias_correction=False prodigy_safeguard_warmup=False base_cache_dir=PosixPath('/data/wand_cache') num_validation_images=30 log_batch_steps=100 run_name=None record_training=True validation_steps=500 train_sigma_distribution='linear' inference_sigma_distribution='shift' quantize=False gradient_checkpointing=False compile=False lora_map_save_params=False log_model_steps=None resume_optimizer=False sample_steps=500 upload_optimizer=False early_stop=False preprocessing_epoch_len=128 train_regional=False preprocessing_epoch_repetitions=1 lora_rank=16 ema=False composite_reference=False train_color_fix=False num_workers=None wandb_entity='wand-tech' warmup_start_lr=0.0 lr_T_mult=1 lr_T_0=None logger_service='wandb' clearml_task_type='training' load_multi_view_lora=False train_max_sequence_length=512 train_dist='linear' train_shift=True inference_dist='linear' inference_shift=True static_mu=None loss_weight_dist=None\n"
+     ]
+    }
+   ],
+   "source": [
+    "trainer = ExperimentTrainer(foundation,dp,config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "d92855c1",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33meleazhong\u001b[0m to \u001b[32mhttps://api.wandb.ai\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "wandb.init called with:\n",
+      "  project: wand-finetune\n",
+      "  entity: wand-tech\n",
+      "  name: None\n",
+      "  config: {'foundation': <FoundationEnum.FLUX: 'flux'>, 'instance_data_dir': None, 'class_data_dir': None, 'instance_prompt': None, 'class_prompt': None, 'num_class_images': 10, 'output_dir': 'output', 'seed': None, 'size': 1024, 'center_crop': False, 'train_batch_size': 1, 'num_train_epochs': 1, 'max_train_steps': None, 'save_steps': 1000, 'save_path': None, 'gradient_accumulation_steps': 1, 'learning_rate': 0.001, 'learning_rate_1d': 1e-06, 'scale_lr': False, 'lr_scheduler': 'constant', 'lr_warmup_steps': 0, 'base_lr': 1e-05, 'max_lr': 0.001, 'step_size_up': 2000, 'cyclic_lr_mode': <CyclicLRMode.TRIANGULAR2: 'triangular2'>, 'cyclic_lr_cycle_momentum': False, 'optim': <OptimizerType.ADAMW: 'adamw'>, 'adam_beta1': 0.9, 'adam_beta2': 0.999, 'adam_weight_decay': 0.01, 'adam_epsilon': 1e-08, 'max_grad_norm': 1.0, 'mixed_precision': 'bf16', 'concepts_list': None, 'modifier_tokens': None, 'initializer_tokens': None, 'checkpointing_steps': 9999, 'resume_from_checkpoint': None, 'train_text_encoder': True, 'gcs_bucket': None, 'topic_id': 'finetune-complete', 'concepts': None, 'global_step': 0, 'prior_loss_weight': 1.0, 'wand_user_id': 'test', 'wand_model_id': 'testing', 'wand_model_bucket': 'wand-finetune', 'wand_project_name': 'wand-finetune', 'num_sample_images': 5, 'prodigy_beta3': None, 'prodigy_decouple': True, 'prodigy_use_bias_correction': False, 'prodigy_safeguard_warmup': False, 'base_cache_dir': PosixPath('/data/wand_cache'), 'num_validation_images': 30, 'log_batch_steps': 100, 'run_name': None, 'record_training': True, 'validation_steps': 500, 'train_sigma_distribution': 'linear', 'inference_sigma_distribution': 'shift', 'quantize': False, 'gradient_checkpointing': False, 'compile': False, 'lora_map_save_params': False, 'log_model_steps': None, 'resume_optimizer': False, 'sample_steps': 500, 'upload_optimizer': False, 'early_stop': False, 'preprocessing_epoch_len': 128, 'train_regional': False, 'preprocessing_epoch_repetitions': 1, 'lora_rank': 16, 'ema': False, 'composite_reference': False, 'train_color_fix': False, 'num_workers': None, 'wandb_entity': 'wand-tech', 'warmup_start_lr': 0.0, 'lr_T_mult': 1, 'lr_T_0': None, 'logger_service': 'wandb', 'clearml_task_type': 'training', 'load_multi_view_lora': False, 'train_max_sequence_length': 512, 'train_dist': 'linear', 'train_shift': True, 'inference_dist': 'linear', 'inference_shift': True, 'static_mu': None, 'loss_weight_dist': None}\n",
+      "  tags: None\n",
+      "  kwargs: {'save_code': True}\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Tracking run with wandb version 0.23.0"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Run data is saved locally in <code>/home/ubuntu/Qwen-Image-Edit-Angles/wandb/run-20251122_181330-lg6f3z2h</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Syncing run <strong><a href='https://wandb.ai/wand-tech/wand-finetune/runs/lg6f3z2h' target=\"_blank\">graceful-galaxy-731</a></strong> to <a href='https://wandb.ai/wand-tech/wand-finetune' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/developer-guide' target=\"_blank\">docs</a>)<br>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View project at <a href='https://wandb.ai/wand-tech/wand-finetune' target=\"_blank\">https://wandb.ai/wand-tech/wand-finetune</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View run at <a href='https://wandb.ai/wand-tech/wand-finetune/runs/lg6f3z2h' target=\"_blank\">https://wandb.ai/wand-tech/wand-finetune/runs/lg6f3z2h</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using suggested max workers 26\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Train:   0%|          | 0/6 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Memory allocated: 55297.20 MB\n",
+      "Memory reserved: 55616.00 MB\n",
+      "Total memory: 81089.88 MB\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Preprocess Batch: 100%|██████████| 128/128 [00:16<00:00,  7.73it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Memory allocated: 56371.81 MB\n",
+      "Memory reserved: 56676.00 MB\n",
+      "Total memory: 81089.88 MB\n",
+      "Repetition: 0\n"
+     ]
+    },
+    {
+     "ename": "RuntimeError",
+     "evalue": "Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
+      "\u001b[0;32m/tmp/ipykernel_2236633/4032920361.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
+      "\u001b[0;32m~/wand-ml/wandml/utils/debug.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     20\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mDEBUG\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 22\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     23\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m                 \u001b[0mstart_time\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtime\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mperf_counter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/wand-ml/wandml/trainers/experiment_trainer.py\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    295\u001b[0m                             \u001b[0mbatch\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"epoch\"\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mepoch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    296\u001b[0m                             \u001b[0mbatch\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"split\"\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"train\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 297\u001b[0;31m                             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msingle_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    298\u001b[0m                             \u001b[0mbatch_num\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    299\u001b[0m                             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mglobal_step\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/wand-ml/wandml/trainers/experiment_trainer.py\u001b[0m in \u001b[0;36msingle_step\u001b[0;34m(self, batch)\u001b[0m\n\u001b[1;32m    334\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0maccelerator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maccumulate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    335\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0maccelerator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautocast\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 336\u001b[0;31m                 \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msingle_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    337\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mctimed\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"accelerator.backward(loss)\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    338\u001b[0m                 \u001b[0maccelerator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/wand-ml/wandml/core/hooks.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     12\u001b[0m                 \u001b[0;32mreturn\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m             \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmanager\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_pre_hooks\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmethod\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     15\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmanager\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_post_hooks\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/Qwen-Image-Edit-Angles/qwenimage/foundation.py\u001b[0m in \u001b[0;36msingle_step\u001b[0;34m(self, batch)\u001b[0m\n\u001b[1;32m    190\u001b[0m         \u001b[0mbatch_size\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx_0\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    191\u001b[0m         \u001b[0mt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtimestep_dist_utils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_train_t\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mbatch_size\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mseq_len\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mseq_len\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 192\u001b[0;31m         \u001b[0mx_t\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;36m1.0\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mx_0\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mt\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mx_1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    193\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    194\u001b[0m         \u001b[0ml_channels\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtransformer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0min_channels\u001b[0m \u001b[0;34m//\u001b[0m \u001b[0;36m4\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;31mRuntimeError\u001b[0m: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!"
+     ]
+    }
+   ],
+   "source": [
+    "trainer.train()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0eea7b23",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

scripts/train.py

@@ -0,0 +1,64 @@
+# # %%
+# %cd /home/ubuntu/Qwen-Image-Edit-Angles
+
+# %%
+import os
+import subprocess
+from pathlib import Path
+import argparse
+
+from ruamel.yaml import YAML
+import diffusers
+
+
+from wandml.trainers.experiment_trainer import ExperimentTrainer
+from wandml import WandDataPipe
+import wandml
+
+from qwenimage.finetuner import QwenLoraFinetuner
+
+
+# %%
+
+
+# %%
+
+
+from qwenimage.datasets import StyleSourceWithRandomRef
+
+
+src = StyleSourceWithRandomRef("/data/styles-finetune-data-artistic/tarot", "<0001>", "/data/image", set_len=1000)
+
+# %%
+from qwenimage.task import TextToImageWithRefTask
+
+task = TextToImageWithRefTask()
+
+# %%
+dp = WandDataPipe()
+dp.add_source(src)
+dp.set_task(task)
+
+# %%
+from qwenimage.datamodels import QwenConfig
+from qwenimage.foundation import QwenImageFoundation
+
+config = QwenConfig(
+    # preprocessing_epoch_len=0,
+)
+foundation = QwenImageFoundation(config=config)
+
+# %%
+finetuner = QwenLoraFinetuner(foundation, config)
+finetuner.load(None)
+
+# %%
+trainer = ExperimentTrainer(foundation,dp,config)
+
+# %%
+trainer.train()
+
+# %%
+
+
+