import torch
from models import VQVAE, build_vae_var
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoTokenizer, SiglipTextModel
from peft import LoraConfig, get_peft_model
from torchvision.transforms import ToPILImage
import random
import gradio as gr
class SimpleAdapter(nn.Module):
def __init__(self, input_dim=512, hidden_dim=1024, out_dim=1024):
super(SimpleAdapter, self).__init__()
self.layer1 = nn.Linear(input_dim, hidden_dim)
self.norm0 = nn.LayerNorm(input_dim)
self.activation1 = nn.GELU()
self.layer2 = nn.Linear(hidden_dim, out_dim)
self.norm2 = nn.LayerNorm(out_dim)
self._initialize_weights()
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight, gain=0.001)
nn.init.zeros_(m.bias)
elif isinstance(m, nn.LayerNorm):
nn.init.ones_(m.weight)
nn.init.zeros_(m.bias)
def forward(self, x):
x = self.norm0(x)
x = self.layer1(x)
x = self.activation1(x)
x = self.layer2(x)
x = self.norm2(x)
return x
class InrenceTextVAR(nn.Module):
def __init__(self, pl_checkpoint=None, start_class_id=578, hugging_face_token=None, siglip_model='google/siglip-base-patch16-224', device="cpu", MODEL_DEPTH=16):
super(InrenceTextVAR, self).__init__()
self.device = device
self.class_id = start_class_id
# Define layers
patch_nums = (1, 2, 3, 4, 5, 6, 8, 10, 13, 16)
self.vae, self.var = build_vae_var(
V=4096, Cvae=32, ch=160, share_quant_resi=4,
device=device, patch_nums=patch_nums,
num_classes=1000, depth=MODEL_DEPTH, shared_aln=False,
)
self.text_processor = AutoTokenizer.from_pretrained(siglip_model, token=hugging_face_token)
self.siglip_text_encoder = SiglipTextModel.from_pretrained(siglip_model, token=hugging_face_token).to(device)
self.adapter = SimpleAdapter(
input_dim=self.siglip_text_encoder.config.hidden_size,
out_dim=self.var.C # Ensure dimensional consistency
).to(device)
self.apply_lora_to_var()
if pl_checkpoint is not None:
state_dict = torch.load(pl_checkpoint, map_location="cpu")['state_dict']
var_state_dict = {k[len('var.'):]: v for k, v in state_dict.items() if k.startswith('var.')}
vae_state_dict = {k[len('vae.'):]: v for k, v in state_dict.items() if k.startswith('vae.')}
adapter_state_dict = {k[len('adapter.'):]: v for k, v in state_dict.items() if k.startswith('adapter.')}
self.var.load_state_dict(var_state_dict)
self.vae.load_state_dict(vae_state_dict)
self.adapter.load_state_dict(adapter_state_dict)
del self.vae.encoder
def apply_lora_to_var(self):
"""
Applies LoRA (Low-Rank Adaptation) to the VAR model.
"""
def find_linear_module_names(model):
linear_module_names = []
for name, module in model.named_modules():
if isinstance(module, nn.Linear):
linear_module_names.append(name)
return linear_module_names
linear_module_names = find_linear_module_names(self.var)
lora_config = LoraConfig(
r=8,
lora_alpha=32,
target_modules=linear_module_names,
lora_dropout=0.05,
bias="none",
)
self.var = get_peft_model(self.var, lora_config)
@torch.no_grad()
def generate_image(self, text, beta=1, seed=None, more_smooth=False, top_k=0, top_p=0.9):
if seed is None:
seed = random.randint(0, 2**32 - 1)
inputs = self.text_processor([text], padding="max_length", return_tensors="pt").to(self.device)
outputs = self.siglip_text_encoder(**inputs)
pooled_output = outputs.pooler_output # pooled (EOS token) states
pooled_output = F.normalize(pooled_output, p=2, dim=-1) # Normalize delta condition
cond_delta = F.normalize(pooled_output, p=2, dim=-1).to(self.device) # Use correct device
cond_delta = self.adapter(cond_delta)
cond_delta = F.normalize(cond_delta, p=2, dim=-1) # Normalize delta condition
generated_images = self.var.autoregressive_infer_cfg(
B=1,
label_B=self.class_id,
delta_condition=cond_delta[:1],
beta=beta,
alpha=1,
top_k=top_k,
top_p=top_p,
more_smooth=more_smooth,
g_seed=seed
)
image = ToPILImage()(generated_images[0].cpu())
return image
if __name__ == '__main__':
# Initialize the model
checkpoint = 'VARtext_v1.pth' # Replace with your actual checkpoint path
device = 'cpu' if not torch.cuda.is_available() else 'cuda'
state_dict = torch.load(checkpoint, map_location="cpu")
model = InrenceTextVAR(device=device)
model.load_state_dict(state_dict)
model.to(device)
def generate_image_gradio(text, beta=1.0, seed=None, more_smooth=False, top_k=0, top_p=0.9):
print(f"Generating image for text: {text}\n"
f"beta: {beta}\n"
f"seed: {seed}\n"
f"more_smooth: {more_smooth}\n"
f"top_k: {top_k}\n"
f"top_p: {top_p}\n")
image = model.generate_image(text, beta=beta, seed=seed, more_smooth=more_smooth, top_k=int(top_k), top_p=top_p)
return image
with gr.Blocks() as demo:
gr.Markdown("# Text to Image/Video Generator")
with gr.Tab("Generate Image"):
text_input = gr.Textbox(label="Input Text")
beta_input = gr.Slider(label="Beta", minimum=0.0, maximum=2.5, step=0.05, value=1.0)
seed_input = gr.Number(label="Seed", value=None)
more_smooth_input = gr.Checkbox(label="More Smooth", value=False)
top_k_input = gr.Number(label="Top K", value=0)
top_p_input = gr.Slider(label="Top P", minimum=0.0, maximum=1.0, step=0.01, value=0.9)
generate_button = gr.Button("Generate Image")
image_output = gr.Image(label="Generated Image")
generate_button.click(
generate_image_gradio,
inputs=[text_input, beta_input, seed_input, more_smooth_input, top_k_input, top_p_input],
outputs=image_output
)
demo.launch()