yujiepan / qwen-image-edit-plus-tiny-random

This tiny model is for debugging. It is randomly initialized with the config adapted from Qwen/Qwen-Image-Edit-2509 .

File size:

• ~10MB text_encoder/model.safetensors
• ~200KB transformer/diffusion_pytorch_model.safetensors
• ~5MB vae/diffusion_pytorch_model.safetensors

Example usage:

import torch
from diffusers import QwenImageEditPlusPipeline
from PIL import Image
import numpy as np

model_id = "yujiepan/qwen-image-edit-plus-tiny-random"
torch_dtype = torch.bfloat16
device = "cuda"
pipe = QwenImageEditPlusPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
pipe = pipe.to(device)

positive_magic = {
    "en": "Ultra HD, 4K, cinematic composition.",
    "zh": "超清，4K，电影级构图"
}
prompt = '''A coffee shop'''
prompt += 'Some dummy random texts to make prompt long enough ' * 10
negative_prompt = " "

# Generate with different aspect ratios
aspect_ratios = {
    "1:1": (1328, 1328),
    "16:9": (1664, 928),
    "9:16": (928, 1664),
    "4:3": (1472, 1140),
    "3:4": (1140, 1472)
}

image1 = Image.fromarray(np.random.randint(0, 255, (1328, 1328, 3), dtype=np.uint8))
image2 = Image.fromarray(np.random.randint(0, 255, (1664, 928, 3), dtype=np.uint8))

for width, height in aspect_ratios.values():
    image = pipe(
        image=[image1, image2],
        prompt=prompt + positive_magic["en"],
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        num_inference_steps=4,
        true_cfg_scale=4.0,
        generator=torch.Generator(device="cuda").manual_seed(42)
    ).images[0]
    print(image)

Codes to create this repo:

import json

import torch
from diffusers import (
    AutoencoderKLQwenImage,
    DiffusionPipeline,
    FlowMatchEulerDiscreteScheduler,
    QwenImageEditPlusPipeline,
    QwenImagePipeline,
    QwenImageTransformer2DModel,
)
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoTokenizer, Qwen2_5_VLForConditionalGeneration, AutoProcessor, Qwen2VLProcessor
from transformers.generation import GenerationConfig

source_model_id = "Qwen/Qwen-Image-Edit-2509"
save_folder = "/tmp/yujiepan/qwen-image-edit-plus-tiny-random"

torch.set_default_dtype(torch.bfloat16)
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
    source_model_id, subfolder='scheduler')
tokenizer = AutoTokenizer.from_pretrained(
    source_model_id, subfolder='tokenizer')

def save_json(path, obj):
    import json
    from pathlib import Path
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    with open(path, 'w', encoding='utf-8') as f:
        json.dump(obj, f, indent=2, ensure_ascii=False)

def init_weights(model):
    import torch
    from transformers import set_seed
    set_seed(42)
    model = model.cpu()
    with torch.no_grad():
        for name, p in sorted(model.named_parameters()):
            torch.nn.init.normal_(p, 0, 0.1)
            print(name, p.shape, p.dtype, p.device)

with open(hf_hub_download(source_model_id, filename='text_encoder/config.json', repo_type='model'), 'r', encoding='utf - 8') as f:
    config = json.load(f)
    config.update({
        'hidden_size': 32,
        'intermediate_size': 64,
        'max_window_layers': 1,
        'num_attention_heads': 2,
        'num_hidden_layers': 2,
        'num_key_value_heads': 1,
        'sliding_window': 64,
        'tie_word_embeddings': True,
        'use_sliding_window': True,
    })
    del config['torch_dtype']
    config['rope_scaling']['mrope_section'] = [4, 2, 2]
    config['text_config'].update({
        'hidden_size': 32,
        'intermediate_size': 64,
        'num_attention_heads': 2,
        'num_hidden_layers': 2,
        'num_key_value_heads': 1,
        'sliding_window': 64,
        'tie_word_embeddings': True,
        'max_window_layers': 1,
        'use_sliding_window': True,
        'layer_types': ['full_attention', 'sliding_attention']
    })
    del config['text_config']['torch_dtype']
    config['text_config']['rope_scaling']['mrope_section'] = [4, 2, 2]
    config['vision_config'].update(
        {
            'depth': 2,
            'fullatt_block_indexes': [0],
            'hidden_size': 32,
            'intermediate_size': 64,
            'num_heads': 2,
            'out_hidden_size': 32,
        }
    )
    del config['vision_config']['torch_dtype']
    save_json(f'{save_folder}/text_encoder/config.json', config)
    text_encoder_config = AutoConfig.from_pretrained(
        f'{save_folder}/text_encoder')
    text_encoder = Qwen2_5_VLForConditionalGeneration(
        text_encoder_config).to(torch.bfloat16)
    generation_config = GenerationConfig.from_pretrained(
        source_model_id, subfolder='text_encoder')
    # text_encoder.config.generation_config = generation_config
    text_encoder.generation_config = generation_config
    init_weights(text_encoder)

with open(hf_hub_download(source_model_id, filename='transformer/config.json', repo_type='model'), 'r', encoding='utf-8') as f:
    config = json.load(f)
    config.update({
        'attention_head_dim': 32,
        'axes_dims_rope': [8, 12, 12],
        'joint_attention_dim': 32,
        'num_attention_heads': 1,
        'num_layers': 2,
    })
    if 'pooled_projection_dim' in config:
        del config['pooled_projection_dim']  # not used
    save_json(f'{save_folder}/transformer/config.json', config)
    transformer_config = QwenImageTransformer2DModel.load_config(
        f'{save_folder}/transformer')
    transformer = QwenImageTransformer2DModel.from_config(
        transformer_config)
    init_weights(transformer)

with open(hf_hub_download(source_model_id, filename='vae/config.json', repo_type='model'), 'r', encoding='utf-8') as f:
    config = json.load(f)
    config.update({
        'num_res_blocks': 1,
        'base_dim': 16,
        'dim_mult': [1, 2, 4, 4],
    })
    del config['latents_mean']  # not used
    del config['latents_std']  # not used
    save_json(f'{save_folder}/vae/config.json', config)
    vae_config = AutoencoderKLQwenImage.load_config(f'{save_folder}/vae')
    vae = AutoencoderKLQwenImage.from_config(vae_config)
    init_weights(vae)

pipeline = QwenImageEditPlusPipeline(
    scheduler=scheduler,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    transformer=transformer,
    vae=vae,
    processor=Qwen2VLProcessor.from_pretrained(
        source_model_id, subfolder='processor'),
)
pipeline = pipeline.to(torch.bfloat16)
pipeline.save_pretrained(save_folder, safe_serialization=True)
print(pipeline)