ibyteohdear/Qwen-Rapid-AIO

import os
import torch
from huggingface_hub import HfApi, hf_hub_download
from safetensors.torch import load_file
from diffusers import QwenImageEditPlusPipeline



HF_TOKEN = os.getenv("HF_TOKEN")
api = HfApi()
dtype = torch.bfloat16

DEST_REPO = "ibyteohdear/Qwen-Rapid-AIO-v23"
FILENAME = "Qwen-Rapid-Full-Bundle"



os.makedirs(FILENAME, exist_ok=True)

print("Accessing v rapid checkpoint...")
v_path = hf_hub_download(
    token=HF_TOKEN,
    repo_id="Phr00t/Qwen-Image-Edit-Rapid-AIO",
    filename="v23/Qwen-Rapid-AIO-SFW-v23.safetensors",
    repo_type="model"


)


print("Loading state dict into CPU memory...")
state_dict = load_file(v_path)







transformer_weights = {}
vae_weights = {}
text_encoder_weights = {}




first_key = next(iter(state_dict.keys()))
print(f"Format detection - first key detected: {first_key}")




for k, v in state_dict.items():
    if k.startswith("model.diffusion_model."):
        transformer_weights[k.replace("model.diffusion_model.", "")] = v
    elif k.startswith("transformer."):
        transformer_weights[k.replace("transformer.", "")] = v
    elif k.startswith("first_stage_model."):
        vae_weights[k.replace("first_stage_model.", "")] = v
    elif k.startswith("vae."):
        vae_weights[k.replace("vae.", "")] = v
    elif "text_encoder" in k or "conditioner" in k:
        if "conditioner.embedders.0." in k:
            text_encoder_weights[k.replace("conditioner.embedders.0.", "")] = v
        elif "text_encoder." in k:
            text_encoder_weights[k.replace("text_encoder.", "")] = v

print("Initializing base pipeline...")
pipe = QwenImageEditPlusPipeline.from_pretrained(
    "ibyteohdear/Qwen-Image-Edit-2511",




    torch_dtype=dtype,

)

print("Injecting custom weights into pipeline components...")
if transformer_weights:
    pipe.transformer.load_state_dict(transformer_weights, strict=False)
if vae_weights:
    pipe.vae.load_state_dict(vae_weights, strict=False)
if text_encoder_weights and hasattr(pipe, "text_encoder") and pipe.text_encoder is not None:
    pipe.text_encoder.load_state_dict(text_encoder_weights, strict=False)

# Clean up source memory
del state_dict
torch.cuda.empty_cache()

print(f"Saving full pipeline components locally to {FILENAME}...")
pipe.save_pretrained(FILENAME)

print(f"Uploading complete bundle to {DEST_REPO}...")
api.create_repo(repo_id=DEST_REPO, token=HF_TOKEN, private=False, exist_ok=True)

api.upload_folder(
    folder_path=FILENAME,
    repo_id=DEST_REPO,
    repo_type="model",
    commit_message="Upload full Qwen Rapid AIO pipeline (Transformer, VAE, and Text Encoder)",
    token=HF_TOKEN,
)

print("Upload complete!")