train.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Optional, Union, Dict, Any, Tuple
import json
from sentence_transformers import SentenceTransformer
from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler
from diffusers.utils import randn_tensor
import safetensors.torch


class QwenEmbeddingAdapter(nn.Module):
    """
    Adapter layer to project Qwen3 embeddings (1024) to SDXL-compatible dimensions (2048)
    """
    def __init__(self, qwen_dim=1024, sdxl_dim=2048):
        super().__init__()
        self.projection = nn.Linear(qwen_dim, sdxl_dim)
        self.layer_norm = nn.LayerNorm(sdxl_dim)
        
    def forward(self, qwen_embeddings):
        """
        Args:
            qwen_embeddings: tensor of shape [batch_size, seq_len, 1024]
        Returns:
            projected_embeddings: tensor of shape [batch_size, seq_len, 2048]
        """
        projected = self.projection(qwen_embeddings)
        return self.layer_norm(projected)


class QwenSDXLPipeline:
    """
    SDXL Pipeline with Qwen3 embedding model replacing CLIP text encoders
    """
    def __init__(
        self,
        qwen_model_path: str = "models/Qwen3-Embedding-0.6B",
        unet_path: str = "models/extracted_components/waiNSFWIllustrious_v140_unet.safetensors",
        unet_config_path: str = "models/extracted_components/waiNSFWIllustrious_v140_unet_config.json",
        vae_path: str = "models/extracted_components/waiNSFWIllustrious_v140_vae.safetensors",
        vae_config_path: str = "models/extracted_components/waiNSFWIllustrious_v140_vae_config.json",
        device: str = "cuda",
        dtype: torch.dtype = torch.bfloat16
    ):
        self.device = device
        self.dtype = dtype
        
        # Load Qwen3 embedding model
        print("Loading Qwen3 embedding model...")
        self.qwen_model = SentenceTransformer(qwen_model_path)
        self.qwen_model.to(device)
        
        # Initialize adapter layer
        self.adapter = QwenEmbeddingAdapter()
        self.adapter.to(device, dtype)
        
        # Load UNet
        print("Loading UNet...")
        with open(unet_config_path, 'r') as f:
            unet_config = json.load(f)
        self.unet = UNet2DConditionModel.from_config(unet_config)
        unet_state_dict = safetensors.torch.load_file(unet_path)
        self.unet.load_state_dict(unet_state_dict)
        self.unet.to(device, dtype)
        
        # Load VAE
        print("Loading VAE...")
        with open(vae_config_path, 'r') as f:
            vae_config = json.load(f)
        self.vae = AutoencoderKL.from_config(vae_config)
        vae_state_dict = safetensors.torch.load_file(vae_path)
        self.vae.load_state_dict(vae_state_dict)
        self.vae.to(device, dtype)
        
        # Initialize scheduler
        self.scheduler = DDPMScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")
        
        # Set pipeline attributes
        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
        self.default_sample_size = self.unet.config.sample_size
        
        print("Pipeline initialization complete!")
    
    def encode_prompt_with_qwen(
        self,
        prompt: Union[str, List[str]],
        device: torch.device,
        num_images_per_prompt: int = 1,
        do_classifier_free_guidance: bool = True,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        prompt_embeds: Optional[torch.Tensor] = None,
        negative_prompt_embeds: Optional[torch.Tensor] = None,
        pooled_prompt_embeds: Optional[torch.Tensor] = None,
        negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
    ):
        """
        Encode prompts using Qwen3 embedding model instead of CLIP
        """
        if prompt_embeds is not None:
            return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
        
        # Ensure prompt is a list
        if isinstance(prompt, str):
            prompt = [prompt]
        
        batch_size = len(prompt)
        
        # Encode prompts with Qwen3
        with torch.no_grad():
            # Use query prompt for better text understanding
            qwen_embeddings = self.qwen_model.encode(
                prompt, 
                prompt_name="query",
                convert_to_tensor=True,
                device=device
            )  # Shape: [batch_size, 1024]
        
        # Add sequence dimension and project to SDXL dimensions
        # Expand to sequence length 77 (CLIP's default)
        seq_len = 77
        qwen_embeddings = qwen_embeddings.unsqueeze(1).expand(-1, seq_len, -1)  # [batch_size, 77, 1024]
        
        # Project to SDXL dimensions using adapter
        prompt_embeds = self.adapter(qwen_embeddings.to(self.dtype))  # [batch_size, 77, 2048]
        
        # For SDXL, we need pooled embeddings (global representation)
        pooled_prompt_embeds = prompt_embeds.mean(dim=1)  # [batch_size, 2048]
        
        # Handle negative prompts
        if do_classifier_free_guidance:
            if negative_prompt is None:
                negative_prompt = [""] * batch_size
            elif isinstance(negative_prompt, str):
                negative_prompt = [negative_prompt] * batch_size
            
            # Encode negative prompts
            with torch.no_grad():
                negative_qwen_embeddings = self.qwen_model.encode(
                    negative_prompt,
                    prompt_name="query", 
                    convert_to_tensor=True,
                    device=device
                )
            
            negative_qwen_embeddings = negative_qwen_embeddings.unsqueeze(1).expand(-1, seq_len, -1)
            negative_prompt_embeds = self.adapter(negative_qwen_embeddings.to(self.dtype))
            negative_pooled_prompt_embeds = negative_prompt_embeds.mean(dim=1)
        else:
            negative_prompt_embeds = None
            negative_pooled_prompt_embeds = None
        
        # Duplicate embeddings for each generation per prompt
        if num_images_per_prompt > 1:
            prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
            prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
            
            pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt)
            pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
            
            if negative_prompt_embeds is not None:
                negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
                negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
                
                negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt)
                negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1)
        
        return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
    
    def prepare_latents(
        self, 
        batch_size: int, 
        num_channels_latents: int, 
        height: int, 
        width: int, 
        dtype: torch.dtype, 
        device: torch.device, 
        generator: Optional[torch.Generator] = None, 
        latents: Optional[torch.Tensor] = None
    ):
        """Prepare latent variables for diffusion process"""
        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
        
        if isinstance(generator, list) and len(generator) != batch_size:
            raise ValueError(
                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
            )

        if latents is None:
            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
        else:
            latents = latents.to(device)

        # scale the initial noise by the standard deviation required by the scheduler
        latents = latents * self.scheduler.init_noise_sigma
        return latents
    
    def _get_add_time_ids(
        self, 
        original_size: Tuple[int, int], 
        crops_coords_top_left: Tuple[int, int], 
        target_size: Tuple[int, int], 
        dtype: torch.dtype,
        text_encoder_projection_dim: int = 2048
    ):
        """Get additional time IDs for SDXL micro-conditioning"""
        add_time_ids = list(original_size + crops_coords_top_left + target_size)
        
        passed_add_embed_dim = (
            self.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim
        )
        expected_add_embed_dim = self.unet.config.addition_embed_type_num_heads

        if expected_add_embed_dim != passed_add_embed_dim:
            raise ValueError(
                f"Model expects an added time embedding vector of length {expected_add_embed_dim}, "
                f"but a vector of {passed_add_embed_dim} was created. The model has an incorrect config."
            )

        add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
        return add_time_ids
    
    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]] = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: Optional[int] = 1,
        eta: float = 0.0,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.Tensor] = None,
        prompt_embeds: Optional[torch.Tensor] = None,
        negative_prompt_embeds: Optional[torch.Tensor] = None,
        pooled_prompt_embeds: Optional[torch.Tensor] = None,
        negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        original_size: Optional[Tuple[int, int]] = None,
        crops_coords_top_left: Tuple[int, int] = (0, 0),
        target_size: Optional[Tuple[int, int]] = None,
    ):
        """
        Modified SDXL inference pipeline using Qwen3 embeddings
        """
        # 0. Default height and width to unet
        height = height or self.default_sample_size * self.vae_scale_factor
        width = width or self.default_sample_size * self.vae_scale_factor
        
        original_size = original_size or (height, width)
        target_size = target_size or (height, width)
        
        # 1. Define call parameters
        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]
        
        device = self.device
        do_classifier_free_guidance = guidance_scale > 1.0
        
        # 2. Encode input prompt with Qwen3
        (
            prompt_embeds,
            negative_prompt_embeds,
            pooled_prompt_embeds,
            negative_pooled_prompt_embeds,
        ) = self.encode_prompt_with_qwen(
            prompt=prompt,
            device=device,
            num_images_per_prompt=num_images_per_prompt,
            do_classifier_free_guidance=do_classifier_free_guidance,
            negative_prompt=negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            pooled_prompt_embeds=pooled_prompt_embeds,
            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
        )
        
        # 3. Prepare timesteps
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps
        
        # 4. Prepare latent variables
        num_channels_latents = self.unet.config.in_channels
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
            height,
            width,
            prompt_embeds.dtype,
            device,
            generator,
            latents,
        )
        
        # 5. Prepare added time ids & embeddings (SDXL micro-conditioning)
        add_text_embeds = pooled_prompt_embeds
        text_encoder_projection_dim = pooled_prompt_embeds.shape[-1]  # 2048
        
        add_time_ids = self._get_add_time_ids(
            original_size,
            crops_coords_top_left,
            target_size,
            dtype=prompt_embeds.dtype,
            text_encoder_projection_dim=text_encoder_projection_dim,
        )
        
        if do_classifier_free_guidance:
            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
            add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
            add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0)
        
        prompt_embeds = prompt_embeds.to(device)
        add_text_embeds = add_text_embeds.to(device)
        add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
        
        # 6. Denoising loop
        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
        with torch.cuda.amp.autocast(enabled=(self.dtype == torch.float16)):
            for i, t in enumerate(timesteps):
                # expand the latents if we are doing classifier free guidance
                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
                
                # predict the noise residual
                added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
                noise_pred = self.unet(
                    latent_model_input,
                    t,
                    encoder_hidden_states=prompt_embeds,
                    added_cond_kwargs=added_cond_kwargs,
                    return_dict=False,
                )[0]
                
                # perform guidance
                if do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                
                # compute the previous noisy sample x_t -> x_t-1
                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
        
        # 7. Decode latents to images
        if output_type != "latent":
            # make sure the VAE is in float32 mode, as it overflows in float16
            needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
            
            if needs_upcasting:
                self.vae.to(dtype=torch.float32)
                latents = latents.to(torch.float32)
            
            latents = latents / self.vae.config.scaling_factor
            image = self.vae.decode(latents, return_dict=False)[0]
            
            if needs_upcasting:
                self.vae.to(dtype=torch.float16)
        else:
            image = latents
        
        # 8. Post-process images
        if output_type == "pil":
            image = (image / 2 + 0.5).clamp(0, 1)
            image = image.cpu().permute(0, 2, 3, 1).float().numpy()
            # Convert to PIL
            from PIL import Image
            image = [Image.fromarray((img * 255).astype("uint8")) for img in image]
        
        if not return_dict:
            return (image,)
        
        return {"images": image}


def test_inference():
    """Test the Qwen-SDXL pipeline"""
    print("Initializing Qwen-SDXL Pipeline...")
    
    pipeline = QwenSDXLPipeline(
        device="cuda" if torch.cuda.is_available() else "cpu",
        dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
    )
    
    # Test prompts
    prompts = [
        "A beautiful landscape with mountains and rivers, oil painting style",
        "A cute cat wearing a hat, anime style",
    ]
    
    print("Generating images...")
    for i, prompt in enumerate(prompts):
        print(f"Generating image {i+1}: {prompt}")
        
        result = pipeline(
            prompt=prompt,
            negative_prompt="low quality, blurry, distorted",
            num_inference_steps=20,
            guidance_scale=7.5,
            height=1024,
            width=1024,
        )
        
        # Save image
        if "images" in result:
            image = result["images"][0]
            image.save(f"output_qwen_sdxl_{i+1}.png")
            print(f"Saved: output_qwen_sdxl_{i+1}.png")


if __name__ == "__main__":
    test_inference()