257 lines
9.1 KiB
Python
257 lines
9.1 KiB
Python
|
||
import numpy as np
|
||
import os
|
||
import torch
|
||
import random
|
||
|
||
from accelerate import infer_auto_device_map, load_checkpoint_and_dispatch, init_empty_weights
|
||
from accelerate.utils import BnbQuantizationConfig, load_and_quantize_model
|
||
from PIL import Image
|
||
|
||
from .data.data_utils import add_special_tokens, pil_img2rgb
|
||
from .data.transforms import ImageTransform
|
||
from .inferencer import InterleaveInferencer
|
||
from .modeling.autoencoder import load_ae
|
||
from .modeling.bagel.qwen2_navit import NaiveCache
|
||
from .modeling.bagel import (
|
||
BagelConfig, Bagel, Qwen2Config, Qwen2ForCausalLM,
|
||
SiglipVisionConfig, SiglipVisionModel
|
||
)
|
||
from .modeling.qwen2 import Qwen2Tokenizer
|
||
|
||
|
||
|
||
def load_bagel_model(model_path,quantize_mode,max_mem_per_gpu = "24GiB"):
|
||
print ("Loading model .... ")
|
||
llm_config = Qwen2Config.from_json_file(os.path.join(model_path, "llm_config.json"))
|
||
llm_config.qk_norm = True
|
||
llm_config.tie_word_embeddings = False
|
||
llm_config.layer_module = "Qwen2MoTDecoderLayer"
|
||
vit_config = SiglipVisionConfig.from_json_file(os.path.join(model_path, "vit_config.json"))
|
||
vit_config.rope = False
|
||
vit_config.num_hidden_layers -= 1
|
||
vae_model, vae_config = load_ae(local_path=os.path.join(model_path, "ae.safetensors"))
|
||
config = BagelConfig(
|
||
visual_gen=True,
|
||
visual_und=True,
|
||
llm_config=llm_config,
|
||
vit_config=vit_config,
|
||
vae_config=vae_config,
|
||
vit_max_num_patch_per_side=70,
|
||
connector_act='gelu_pytorch_tanh',
|
||
latent_patch_size=2,
|
||
max_latent_size=64,
|
||
)
|
||
with init_empty_weights():
|
||
language_model = Qwen2ForCausalLM(llm_config)
|
||
vit_model = SiglipVisionModel(vit_config)
|
||
model = Bagel(language_model, vit_model, config)
|
||
model.vit_model.vision_model.embeddings.convert_conv2d_to_linear(vit_config, meta=True)
|
||
|
||
tokenizer = Qwen2Tokenizer.from_pretrained(model_path)
|
||
tokenizer, new_token_ids, _ = add_special_tokens(tokenizer)
|
||
|
||
vae_transform = ImageTransform(1024, 512, 16)
|
||
vit_transform = ImageTransform(980, 224, 14)
|
||
# Model Loading and Multi GPU Infernece Preparing
|
||
|
||
|
||
#max_mem_per_gpu = "24GiB" # Modify it according to your GPU setting. On an A100, 80 GiB is sufficient to load on a single GPU.
|
||
|
||
device_map = infer_auto_device_map(
|
||
model,
|
||
max_memory={0: max_mem_per_gpu},
|
||
no_split_module_classes=["Bagel", "Qwen2MoTDecoderLayer", "SiglipVisionModel"],
|
||
)
|
||
#print(f"device_map is {device_map}")
|
||
|
||
same_device_modules = [
|
||
'language_model.model.embed_tokens',
|
||
'time_embedder',
|
||
'latent_pos_embed',
|
||
'vae2llm',
|
||
'llm2vae',
|
||
'connector',
|
||
'vit_pos_embed'
|
||
]
|
||
|
||
if torch.cuda.device_count() == 1:
|
||
first_device = device_map.get(same_device_modules[0], "cuda:0")
|
||
for k in same_device_modules:
|
||
if k in device_map:
|
||
device_map[k] = first_device
|
||
else:
|
||
device_map[k] = "cuda:0"
|
||
else:
|
||
first_device = device_map.get(same_device_modules[0])
|
||
for k in same_device_modules:
|
||
if k in device_map:
|
||
device_map[k] = first_device
|
||
|
||
if quantize_mode == "fp16":
|
||
model = load_checkpoint_and_dispatch(
|
||
model,
|
||
checkpoint=os.path.join(model_path, "ema.safetensors"),
|
||
device_map=device_map,
|
||
offload_buffers=True,
|
||
offload_folder="offload",
|
||
dtype=torch.bfloat16,
|
||
force_hooks=True,
|
||
).eval()
|
||
elif quantize_mode == "nf4":
|
||
print("Quantizing to NF4")
|
||
bnb_quantization_config = BnbQuantizationConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=False, bnb_4bit_quant_type="nf4")
|
||
model = load_and_quantize_model(
|
||
model,
|
||
weights_location=os.path.join(model_path, "ema.safetensors"),
|
||
bnb_quantization_config=bnb_quantization_config,
|
||
device_map=device_map,
|
||
offload_folder="offload",
|
||
).eval()
|
||
else:#args.mode == 3: # INT8
|
||
print("Quantizing to INT8/fp8")
|
||
bnb_quantization_config = BnbQuantizationConfig(load_in_8bit=True, torch_dtype=torch.float32)
|
||
model = load_and_quantize_model(
|
||
model,
|
||
weights_location=os.path.join(model_path, "ema.safetensors"),
|
||
bnb_quantization_config=bnb_quantization_config,
|
||
device_map=device_map,
|
||
offload_folder="offload",
|
||
).eval()
|
||
|
||
model = model.eval()
|
||
print('Model loaded is done!')
|
||
inferencer = InterleaveInferencer(
|
||
model=model,
|
||
vae_model=vae_model,
|
||
tokenizer=tokenizer,
|
||
vae_transform=vae_transform,
|
||
vit_transform=vit_transform,
|
||
new_token_ids=new_token_ids
|
||
)
|
||
return inferencer
|
||
|
||
|
||
|
||
def set_seed(seed):
|
||
"""Set random seeds for reproducibility"""
|
||
if seed > 0:
|
||
random.seed(seed)
|
||
np.random.seed(seed)
|
||
torch.manual_seed(seed)
|
||
if torch.cuda.is_available():
|
||
torch.cuda.manual_seed(seed)
|
||
torch.cuda.manual_seed_all(seed)
|
||
torch.backends.cudnn.deterministic = True
|
||
torch.backends.cudnn.benchmark = False
|
||
return seed
|
||
|
||
# Text to Image function with thinking option and hyperparameters
|
||
def text_to_image(inferencer,prompt, show_thinking=False, cfg_text_scale=4.0, cfg_interval=0.4,
|
||
timestep_shift=3.0, num_timesteps=50,
|
||
cfg_renorm_min=0.0, cfg_renorm_type="global",
|
||
max_think_token_n=1024, do_sample=False, text_temperature=0.3,
|
||
seed=0, image_ratio="1:1"):
|
||
# Set seed for reproducibility
|
||
set_seed(seed)
|
||
|
||
if image_ratio == "1:1":
|
||
image_shapes = (1024, 1024)
|
||
elif image_ratio == "4:3":
|
||
image_shapes = (768, 1024)
|
||
elif image_ratio == "3:4":
|
||
image_shapes = (1024, 768)
|
||
elif image_ratio == "16:9":
|
||
image_shapes = (576, 1024)
|
||
elif image_ratio == "9:16":
|
||
image_shapes = (1024, 576)
|
||
|
||
# Set hyperparameters
|
||
inference_hyper = dict(
|
||
max_think_token_n=max_think_token_n if show_thinking else 1024,
|
||
do_sample=do_sample if show_thinking else False,
|
||
text_temperature=text_temperature if show_thinking else 0.3,
|
||
cfg_text_scale=cfg_text_scale,
|
||
cfg_interval=[cfg_interval, 1.0], # End fixed at 1.0
|
||
timestep_shift=timestep_shift,
|
||
num_timesteps=num_timesteps,
|
||
cfg_renorm_min=cfg_renorm_min,
|
||
cfg_renorm_type=cfg_renorm_type,
|
||
image_shapes=image_shapes,
|
||
)
|
||
|
||
# Call inferencer with or without think parameter based on user choice
|
||
result = inferencer(text=prompt, think=show_thinking, **inference_hyper)
|
||
return result["image"], result.get("text", "")
|
||
|
||
|
||
# Image Understanding function with thinking option and hyperparameters
|
||
def image_understanding(inferencer,image: Image.Image, prompt: str, show_thinking=False,
|
||
do_sample=False, text_temperature=0.3, max_new_tokens=512):
|
||
if image is None:
|
||
return "Please upload an image."
|
||
|
||
if isinstance(image, np.ndarray):
|
||
image = Image.fromarray(image)
|
||
|
||
image = pil_img2rgb(image)
|
||
|
||
# Set hyperparameters
|
||
inference_hyper = dict(
|
||
do_sample=do_sample,
|
||
text_temperature=text_temperature,
|
||
max_think_token_n=max_new_tokens, # Set max_length
|
||
)
|
||
|
||
# Use show_thinking parameter to control thinking process
|
||
result = inferencer(image=image, text=prompt, think=show_thinking,
|
||
understanding_output=True, **inference_hyper)
|
||
return result["text"]
|
||
|
||
|
||
# Image Editing function with thinking option and hyperparameters
|
||
def edit_image(inferencer,image: Image.Image, prompt: str, show_thinking=False, cfg_text_scale=4.0,
|
||
cfg_img_scale=2.0, cfg_interval=0.0,
|
||
timestep_shift=3.0, num_timesteps=50, cfg_renorm_min=0.0,
|
||
cfg_renorm_type="text_channel", max_think_token_n=1024,
|
||
do_sample=False, text_temperature=0.3, seed=0):
|
||
# Set seed for reproducibility
|
||
set_seed(seed)
|
||
|
||
if image is None:
|
||
return "Please upload an image.", ""
|
||
|
||
if isinstance(image, np.ndarray):
|
||
image = Image.fromarray(image)
|
||
|
||
image = pil_img2rgb(image)
|
||
|
||
# Set hyperparameters
|
||
inference_hyper = dict(
|
||
max_think_token_n=max_think_token_n if show_thinking else 1024,
|
||
do_sample=do_sample if show_thinking else False,
|
||
text_temperature=text_temperature if show_thinking else 0.3,
|
||
cfg_text_scale=cfg_text_scale,
|
||
cfg_img_scale=cfg_img_scale,
|
||
cfg_interval=[cfg_interval, 1.0], # End fixed at 1.0
|
||
timestep_shift=timestep_shift,
|
||
num_timesteps=num_timesteps,
|
||
cfg_renorm_min=cfg_renorm_min,
|
||
cfg_renorm_type=cfg_renorm_type,
|
||
)
|
||
|
||
# Include thinking parameter based on user choice
|
||
result = inferencer(image=image, text=prompt, think=show_thinking, **inference_hyper)
|
||
return result["image"], result.get("text", "")
|
||
|
||
|
||
# Helper function to load example images
|
||
def load_example_image(image_path):
|
||
try:
|
||
return Image.open(image_path)
|
||
except Exception as e:
|
||
print(f"Error loading example image: {e}")
|
||
return None
|
||
|
||
|