2604

README.md

@@ -79,6 +79,13 @@ image.show()
 upscaled = pipe.image_upscale("media/girl.jpg")
 ```
 
+```markdown
+| | |
+|:---:|:---:|
+| <img src="media/123456789.jpg" height="512"/> | <img src="media/123456789.png" height="512"/> |
+| original | 2xupscale |
+```
+
 ### Prompt refine
 ```
 refined = pipe.refine_prompts("girl")
@@ -214,6 +221,10 @@ nohup accelerate launch train.py &
  - Rubles: [For users from Russia](https://www.tbank.ru/cf/90ensBQqpJj)
  - DOGE: DEw2DR8C7BnF8GgcrfTzUjSnGkuMeJhg83
  - BTC: 3JHv9Hb8kEW8zMAccdgCdZGfrHeMhH1rpN
+   USTD
+ - Ethereum / Polygon / BNB SmartChain: 0xD4388B6698dFaE1460E72099D4F208aaCA4f6E6C
+ - Tron: TD7ey4h9igPGdcrcBcnZaz56R5tNgRZNvV
+ - Solana: MMYFJeYEtYHrSNFHChytJDHbEDniXrnAxPNLhJ1LbkB
 
 ## Contacts
 Please contact with us if you may provide some GPU's or money on training

pipeline_sdxs.py

@@ -248,7 +248,6 @@ class SdxsPipeline(DiffusionPipeline):
             
             # Encode -> Decode (using mean for deterministic upscale)
             latents = self.vae.encode(tensors).latent_dist.mean
-            latents = latents * self.vae_latents_std.to(latents) + self.vae_latents_mean.to(latents)
             decoded = self.vae.decode(latents.to(self.vae.dtype))[0]
             
             # 4. Post-process: Denormalize and Crop

test.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1f99ca4f959a85a5453820eedad6b5c1fde29fcadd30a3eaae22097c96a0a7a0
-size 16433141
+oid sha256:bd0d1f7f30e064d8fb03320c204b7d67fd5c7df87da6919af818d6e2f60ee1a2
+size 12104053

train.py

@@ -10,6 +10,7 @@ import bitsandbytes as bnb
 import torch.nn.functional as F
 import argparse
 
+from datetime import datetime
 from diffusers import UNet2DConditionModel, AsymmetricAutoencoderKL, FlowMatchEulerDiscreteScheduler
 from transformers import Qwen3_5Tokenizer, Qwen3_5ForConditionalGeneration
 from torch.utils.data import DataLoader, Sampler
@@ -21,7 +22,7 @@ from tqdm import tqdm
 from PIL import Image, ImageOps
 from torch.utils.checkpoint import checkpoint
 from diffusers.models.attention_processor import AttnProcessor2_0
-from datetime import datetime
+from contextlib import nullcontext
 
 # Muon not tested! pip install git+https://github.com/recoilme/muon_adamw8bit.git
 from muon_adamw8bit import MuonAdamW8bit
@@ -470,12 +471,14 @@ else:
         return min_learning_rate + 0.5 * (base_learning_rate - min_learning_rate) * \
                (1 + math.cos(math.pi * decay_ratio))
     lr_scheduler = LambdaLR(optimizer, lambda step: lr_schedule(step) / base_learning_rate)
-    unet, optimizer, lr_scheduler = accelerator.prepare(unet, optimizer, lr_scheduler)
 
 if torch_compile:
-    print("compiling")
+    print("Compiling UNet... Это займет несколько минут, не прерывайте!")
     unet = torch.compile(unet)
-    print("compiling - ok")
+    print("Compiling - ok")
+
+if not fbp:
+    unet, optimizer, lr_scheduler = accelerator.prepare(unet, optimizer, lr_scheduler)
 
 # Фиксированные семплы
 fixed_samples = get_fixed_samples_by_resolution(dataset)
@@ -702,105 +705,108 @@ for epoch in range(start_epoch, start_epoch + num_epochs):
     unet.train()
     
     for step, (latents, embeddings, attention_mask) in enumerate(dataloader):
-        with accelerator.accumulate(unet):
-            if save_model == False and epoch == 0 and step == 5 :
-                used_gb = torch.cuda.max_memory_allocated() / 1024**3
-                print(f"Шаг {step}: {used_gb:.2f} GB")
-                
-            # шум
-            noise = torch.randn_like(latents, dtype=latents.dtype)
-
-            # 3. Время t (сэмплим, как и раньше, но чуть сжимаем края)
-            u = torch.rand(latents.shape[0], device=latents.device, dtype=latents.dtype)
-            t = u * (1 - 2 * 1e-5) + 1e-5  # Теперь t строго в (0.00001 ... 0.99999)
-            # интерполяция между x0 и шумом
-            noisy_latents = (1.0 - t.view(-1, 1, 1, 1)) * latents + t.view(-1, 1, 1, 1) * noise
-            # делаем integer timesteps для UNet
-            timesteps = t.to(torch.float32).mul(999.0)
-            timesteps = timesteps.clamp(0, scheduler.config.num_train_timesteps - 1)
-
-            # --- Вызов UNet с маской  ---
-            model_pred = unet(
-                noisy_latents, 
-                timesteps, 
-                encoder_hidden_states=embeddings,
-                encoder_attention_mask=attention_mask,
-            ).sample
-            
-            target = noise - latents
-
-            mse_loss = F.mse_loss(model_pred.float(), target.float())                
-            batch_losses.append(mse_loss.detach().item())
-
-            if (global_step % 100 == 0) or (global_step % sink_interval == 0):
-                accelerator.wait_for_everyone()
-
-            losses_dict = {}
-            losses_dict["mse"] = mse_loss
-
-            if (global_step % 100 == 0) or (global_step % sink_interval == 0):
-                accelerator.wait_for_everyone()
-
-            accelerator.backward(mse_loss)
+        
+        if save_model == False and epoch == 0 and step == 5 :
+            used_gb = torch.cuda.max_memory_allocated() / 1024**3
+            print(f"Шаг {step}: {used_gb:.2f} GB")
 
-            if (global_step % 100 == 0) or (global_step % sink_interval == 0):
-                accelerator.wait_for_everyone()
+        amp_context = accelerator.autocast() if torch_compile else nullcontext()
+        with accelerator.accumulate(unet):
+            with amp_context:
+                # шум
+                noise = torch.randn_like(latents, dtype=latents.dtype)
+    
+                # 3. Время t (сэмплим, как и раньше, но чуть сжимаем края)
+                u = torch.rand(latents.shape[0], device=latents.device, dtype=latents.dtype)
+                t = u * (1 - 2 * 1e-5) + 1e-5  # Теперь t строго в (0.00001 ... 0.99999)
+                # интерполяция между x0 и шумом
+                noisy_latents = (1.0 - t.view(-1, 1, 1, 1)) * latents + t.view(-1, 1, 1, 1) * noise
+                # делаем integer timesteps для UNet
+                timesteps = t.to(torch.float32).mul(999.0)
+                timesteps = timesteps.clamp(0, scheduler.config.num_train_timesteps - 1)
+    
+                # --- Вызов UNet с маской  ---
+                model_pred = unet(
+                    noisy_latents, 
+                    timesteps, 
+                    encoder_hidden_states=embeddings,
+                    encoder_attention_mask=attention_mask,
+                ).sample
                 
-            grad = 0.0
-            if not fbp:
-                if accelerator.sync_gradients:
-                    grad_val = accelerator.clip_grad_norm_(unet.parameters(), clip_grad_norm)
-                    grad = grad_val.float().item() if torch.is_tensor(grad_val) else float(grad_val)
-                    optimizer.step()
-                    lr_scheduler.step()
-                    optimizer.zero_grad(set_to_none=True)
-
-            if accelerator.sync_gradients:
-                global_step += 1
-                progress_bar.update(1)
-                if accelerator.is_main_process:
-                    if fbp:
-                        current_lr = base_learning_rate
-                    else:
-                        current_lr = lr_scheduler.get_last_lr()[0]
-                    batch_grads.append(grad)
+                target = noise - latents
+    
+                mse_loss = F.mse_loss(model_pred.float(), target.float())                
+                batch_losses.append(mse_loss.detach().item())
+    
+                if (global_step % 100 == 0) or (global_step % sink_interval == 0):
+                    accelerator.wait_for_everyone()
+    
+                losses_dict = {}
+                losses_dict["mse"] = mse_loss
     
-                    log_data = {}
-                    log_data["loss_mse"] = mse_loss.detach().item()
-                    log_data["lr"] = current_lr
-                    log_data["grad"] = grad
+                if (global_step % 100 == 0) or (global_step % sink_interval == 0):
+                    accelerator.wait_for_everyone()
+    
+                accelerator.backward(mse_loss)
+    
+                if (global_step % 100 == 0) or (global_step % sink_interval == 0):
+                    accelerator.wait_for_everyone()
+                    
+                grad = 0.0
+                if not fbp:
                     if accelerator.sync_gradients:
-                        if use_wandb:
-                            wandb.log(log_data, step=global_step)
-                        if use_comet_ml:
-                            comet_experiment.log_metrics(log_data, step=global_step)
-
-                    current_time = time.time()
-                    is_time_to_sample = (current_time - last_sample_time) >= sample_interval_seconds
-                    if is_time_to_sample or global_step == 50:
-                        # Передаем tuple (emb, mask) для негатива
-                        if save_model:
-                            generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask), global_step)
-                        elif epoch % 10 == 0:
-                            generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask), global_step)
-                        last_n = sink_interval
-                        
-                        if save_model:
-                            has_losses = len(batch_losses) > 0
-                            avg_sample_loss = np.mean(batch_losses[-sink_interval:]) if has_losses else 0.0
-                            last_loss = batch_losses[-1] if has_losses else 0.0
-                            max_loss = max(avg_sample_loss, last_loss)    
-                            should_save = max_loss < min_loss * save_barrier
-                            print(
-                                f"Saving: {should_save} | Max: {max_loss:.4f} | "
-                                f"Last: {last_loss:.4f} | Avg: {avg_sample_loss:.4f}"
-                            )
-                            # 6. Сохранение и обновление
-                            if should_save:
-                                min_loss = max_loss
-                                save_checkpoint(unet)
-                        last_sample_time = current_time
-                        unet.train()
+                        grad_val = accelerator.clip_grad_norm_(unet.parameters(), clip_grad_norm)
+                        grad = grad_val.float().item() if torch.is_tensor(grad_val) else float(grad_val)
+                        optimizer.step()
+                        lr_scheduler.step()
+                        optimizer.zero_grad(set_to_none=True)
+    
+                if accelerator.sync_gradients:
+                    global_step += 1
+                    progress_bar.update(1)
+                    if accelerator.is_main_process:
+                        if fbp:
+                            current_lr = base_learning_rate
+                        else:
+                            current_lr = lr_scheduler.get_last_lr()[0]
+                        batch_grads.append(grad)
+        
+                        log_data = {}
+                        log_data["loss_mse"] = mse_loss.detach().item()
+                        log_data["lr"] = current_lr
+                        log_data["grad"] = grad
+                        if accelerator.sync_gradients:
+                            if use_wandb:
+                                wandb.log(log_data, step=global_step)
+                            if use_comet_ml:
+                                comet_experiment.log_metrics(log_data, step=global_step)
+    
+                        current_time = time.time()
+                        is_time_to_sample = (current_time - last_sample_time) >= sample_interval_seconds
+                        if is_time_to_sample or global_step == 50:
+                            # Передаем tuple (emb, mask) для негатива
+                            if save_model:
+                                generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask), global_step)
+                            elif epoch % 10 == 0:
+                                generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask), global_step)
+                            last_n = sink_interval
+                            
+                            if save_model:
+                                has_losses = len(batch_losses) > 0
+                                avg_sample_loss = np.mean(batch_losses[-sink_interval:]) if has_losses else 0.0
+                                last_loss = batch_losses[-1] if has_losses else 0.0
+                                max_loss = max(avg_sample_loss, last_loss)    
+                                should_save = max_loss < min_loss * save_barrier
+                                print(
+                                    f"Saving: {should_save} | Max: {max_loss:.4f} | "
+                                    f"Last: {last_loss:.4f} | Avg: {avg_sample_loss:.4f}"
+                                )
+                                # 6. Сохранение и обновление
+                                if should_save:
+                                    min_loss = max_loss
+                                    save_checkpoint(unet)
+                            last_sample_time = current_time
+                            unet.train()
 
     if accelerator.is_main_process:
         avg_epoch_loss = np.mean(batch_losses) if len(batch_losses) > 0 else 0.0

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2003b45e59873c0430ebf4889ee94bae5eca7b7aafb898f378689697f5c30c69
+oid sha256:9b2e01643382012adcffc091502a64e908d85a394504a92cf8bb8aa8cec080b9
 size 3210307232