UniPic2-Metaquery-GRPO-Flash

🌌 UniPic2-Metaquery-GRPO-Flash

📖 Introduction

UniPic2-Metaquery-GRPO-Flash is a quantized variant of UniPic2-MetaQuery-GRPO, offering end-to-end image understanding, text-to-image (T2I) generation, and image editing. Optimized for efficiency, it runs smoothly on NVIDIA RTX 40-series GPUs with under 16 GB VRAM — without any performance degradation.

📊 Benchmarks

🧠 Usage

1. Clone the Repository

git clone https://github.com/SkyworkAI/UniPic
cd UniPic-2

2. Set Up the Environment

conda create -n unipic python=3.10
conda activate unipic
pip install -r requirements.txt

3.Text-to-Image Generation

import torch
from PIL import Image
from unipicv2.pipeline_stable_diffusion_3_kontext import StableDiffusion3KontextPipeline
from unipicv2.transformer_sd3_kontext import SD3Transformer2DKontextModel
from unipicv2.stable_diffusion_3_conditioner import StableDiffusion3Conditioner
from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor
from diffusers import FlowMatchEulerDiscreteScheduler, AutoencoderKL,BitsAndBytesConfig

# Load model components
pretrained_model_name_or_path = "/path/to/UniPic2-Metaquery-Flash/UniPic2-Metaquery"
vlm_path   = "/path/to/UniPic2-Metaquery-Flash/Qwen2.5-VL-7B-Instruct-AWQ"

quant = "int4"  # {"int4", "fp16"}

bnb4 = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,  # 与 LMM/Cond 对齐
)

if quant == "int4":
    transformer = SD3Transformer2DKontextModel.from_pretrained(
        pretrained_model_name_or_path, subfolder="transformer",
        quantization_config=bnb4, device_map="auto", low_cpu_mem_usage=True
    )
elif quant == "fp16":
    transformer = SD3Transformer2DKontextModel.from_pretrained(
        pretrained_model_name_or_path, subfolder="transformer",
        torch_dtype=torch.float16, device_map="auto", low_cpu_mem_usage=True
    )
else:
    raise ValueError(f"Unsupported quant: {quant}")

vae = AutoencoderKL.from_pretrained(
    pretrained_model_name_or_path, subfolder="vae",
    torch_dtype=torch.float16, device_map="auto", low_cpu_mem_usage=True).cuda()

# Load Qwen2.5-VL model
lmm = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    vlm_path,
    torch_dtype=torch.bfloat16,device_map="auto",
    attn_implementation="flash_attention_2")

processor = Qwen2_5_VLProcessor.from_pretrained(vlm_path)
processor.chat_template = processor.chat_template.replace(
    "{% if loop.first and message['role'] != 'system' %}system\nYou are a helpful assistant.\n{% endif %}",
    "")

# 加上cuda
conditioner = StableDiffusion3Conditioner.from_pretrained(
    pretrained_model_name_or_path, subfolder="conditioner", torch_dtype=torch.float16).cuda()

scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")

# Create pipeline (note: text encoders set to None)
pipeline = StableDiffusion3KontextPipeline(
    transformer=transformer, vae=vae,
    text_encoder=None, tokenizer=None,
    text_encoder_2=None, tokenizer_2=None,
    text_encoder_3=None, tokenizer_3=None,
    scheduler=scheduler)

# Prepare prompts
prompt = 'a pig with wings and a top hat flying over a happy futuristic scifi city'
negative_prompt = ''

messages = [[{"role": "user", "content": [{"type": "text", "text": f'Generate an image: {txt}'}]}]
            for txt in [prompt, negative_prompt]]

texts = [processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True) for msg in messages]
inputs = processor(text=texts, images=None, videos=None, padding=True, return_tensors="pt").to("cuda")

# Process with Qwen2.5-VL
input_ids, attention_mask = inputs.input_ids, inputs.attention_mask
input_ids = torch.cat([input_ids, input_ids.new_zeros(2, conditioner.config.num_queries)], dim=1)
attention_mask = torch.cat([attention_mask, attention_mask.new_ones(2, conditioner.config.num_queries)], dim=1)
inputs_embeds = lmm.get_input_embeddings()(input_ids)
inputs_embeds[:, -conditioner.config.num_queries:] = conditioner.meta_queries[None].expand(2, -1, -1)

outputs = lmm.model(inputs_embeds=inputs_embeds, attention_mask=attention_mask, use_cache=False)
hidden_states = outputs.last_hidden_state[:, -conditioner.config.num_queries:]
prompt_embeds, pooled_prompt_embeds = conditioner(hidden_states)

# Generate image
image = pipeline(
    prompt_embeds=prompt_embeds[:1],
    pooled_prompt_embeds=pooled_prompt_embeds[:1],
    negative_prompt_embeds=prompt_embeds[1:],
    negative_pooled_prompt_embeds=pooled_prompt_embeds[1:],
    height=512, width=384,
    num_inference_steps=50,
    guidance_scale=3.5,
    generator=torch.Generator(device=transformer.device).manual_seed(42)
).images[0]

image.save("text2image.png")
print(f"Image saved to text2image.png (quant={quant})")

4. Image Editing

# Load image for editing
image = Image.open("text2image.png")
image = fix_longer_edge(image, image_size=512)

prompt = "remove the pig's hat"
negative_prompt = "blurry, low quality, low resolution, distorted, deformed, broken content, missing parts, damaged details, artifacts, glitch, noise, pixelated, grainy, compression artifacts, bad composition, wrong proportion, incomplete editing, unfinished, unedited areas."

# Prepare messages with image input
messages = [[{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": txt}]}]
            for txt in [prompt, negative_prompt]]

texts = [processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True) for msg in messages]

min_pixels = max_pixels = int(image.height * 28 / 32 * image.width * 28 / 32)
inputs = processor(
    text=texts, images=[image]*2,
    min_pixels=min_pixels, max_pixels=max_pixels,
    videos=None, padding=True, return_tensors="pt").cuda()

# Process with vision understanding
input_ids, attention_mask, pixel_values, image_grid_thw = \
    inputs.input_ids, inputs.attention