demo_sparse_lavida.py

# ADOBE CONFIDENTIAL
# Copyright 2025 Adobe
# All Rights Reserved.
# NOTICE: All information contained herein is, and remains
# the property of Adobe and its suppliers, if any. The intellectual
# and technical concepts contained herein are proprietary to Adobe
# and its suppliers and are protected by all applicable intellectual
# property laws, including trade secret and copyright laws.
# Dissemination of this information or reproduction of this material
# is strictly forbidden unless prior written permission is obtained
# from Adobe.


import os
os.environ['DEBUG_FIX_PADDING'] = '1'
os.environ['NOT_ALWASY_DO_2DPOOL'] = '1'
from llava.eval.predict_t2i_edit import build_model

from llava.eval.predict_t2i_edit import text_to_image
from llava.eval.predict_t2i_edit import create_plan,get_feedback,create_plan_edit
import os
from llava.model.utils import maybe_truncate_last_dim,pad_along_last_dim
from PIL import Image
from llava.mm_utils import resize_and_center_crop
from llava.mm_utils import get_model_name_from_path, process_images, tokenizer_image_token
import yaml
import torch
from llava.eval.predict_grounding import predict_grounding
# from llava.eval.demo_utils import  visualize_boxes,extract_bounding_boxes
import copy
from llava.mm_utils import get_model_name_from_path, process_images, tokenizer_image_token
from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, IGNORE_INDEX,SKIP_DOWN_SAMPLE
from llava.conversation import conv_templates, SeparatorStyle
from matplotlib import pyplot as plt
from tqdm.auto import tqdm

pretrained = 'path-to-your-model'

tokenizer, model, image_processor = build_model(pretrained=pretrained)

import yaml
def load_yaml(fp):
    """
    Load a YAML file from a file path or file-like object.

    Parameters:
        fp (str or file-like): File path or file-like object to the YAML file.

    Returns:
        dict or list: Parsed YAML content.
    """
    if isinstance(fp, str):
        with open(fp, 'r') as f:
            return yaml.safe_load(f)
    else:
        return yaml.safe_load(fp)
model.requires_grad_(False)
model.eval()



#prompt = 'lush garden in the middle of a dimly lit old library, fantasy, realistic, 4k'
#prompt = '4k, ultra-hd, photorealistic A cinematic wide shot of a lone astronaut standing on a desolate, alien planet, bathed in the glow of a binary sunset. Dust swirls around their boots. Highly detailed, 8K, sci-fi art, dramatic lighting.'

import time
# Text to Image
def run_text_to_image():

    prompt = 'pink sky. White clouds. bright stars. mountains. trees. river. anime style. high quality.'
    micro_cond = 'ORIGINAL WIDTH : 1024; ORIGINAL HEIGHT : 1024; TOP : 0; LEFT : 0; SCORE : 7.120; '
    config_t2i = load_yaml('llava/eval/1024_eval.yaml')
    config_t2i['config'].update(temperature=0.86)

    cache_kwargs = dict(use_cache=True,cache_vq=True,truncate_vq=True,cache_prompt=True,cache_rounds=2,add_register=True,cache_buffer_round=3,cache_buffer_round_end=3)
    config_t2i['config'].update(
                                micro_cond=micro_cond,
                                # is_legacy=True,
                                # guidance_scale=5,
                                return_compute_positions=True,
                            block_policy=2)

    with torch.no_grad():
        t0 = time.time()
        img1 = text_to_image(model,prompt,tokenizer=tokenizer,**config_t2i['config'],image_resolution=1024, n_tokens=4096,is_legacy=True)

        t1 = time.time()
    time_tiout_cache = t1-t0
    
    config_t2i['config'].update(**cache_kwargs)
    with torch.no_grad():
        t0 = time.time()
        img2 = text_to_image(model,prompt,tokenizer=tokenizer,**config_t2i['config'],image_resolution=1024, n_tokens=4096,is_legacy=True)

        t1 = time.time()
    time_tiwith_cache = t1-t0
    print("Time without cache:", time_tiout_cache)
    print("Time with cache:", time_tiwith_cache)
    img1.save('logs/t2i_output_no_cache.png')
    img2.save('logs/t2i_output_with_cache.png')

def run_image_editing():
    prompt = "Turn the tree branche in the image into a witch's magic wand."
    edit_image = Image.open('assets/witch.jpg')
    edit_image = resize_and_center_crop(edit_image,1024)
    config_edit = load_yaml('llava/eval/1024_eval_edit.yaml')
    cache_kwargs = dict(use_cache=True,cache_vq=True,truncate_vq=True,cache_prompt=True,cache_rounds=2,add_register=True,cache_buffer_round=3,cache_buffer_round_end=3)
    config_edit['config'].update(
                                return_compute_positions=True,
                            block_policy=2)
    with torch.no_grad():
        t0 = time.time()
        img1 = text_to_image(model,prompt, edit_image=edit_image,image_processor=image_processor,tokenizer=tokenizer,**config_edit['config'],image_resolution=1024, n_tokens=4096,is_legacy=False)

        t1 = time.time()
    time_tiout_cache = t1-t0
    print("Time without cache:",)
    config_edit['config'].update(**cache_kwargs)
    with torch.no_grad():
        t0 = time.time()
        img2 = text_to_image(model,prompt, edit_image=edit_image,image_processor=image_processor,tokenizer=tokenizer,**config_edit['config'],image_resolution=1024, n_tokens=4096,is_legacy=False)
        t1 = time.time()
    time_tiwith_cache = t1-t0
    print("Time without cache:", time_tiout_cache)
    print("Time with cache:", time_tiwith_cache)
    img1.save('logs/edit_output_no_cache.png')
    img2.save('logs/edit_output_with_cache.png')


def und():
    device = model.device
    input_image = Image.open('assets/detection.jpg')
    input_image = resize_and_center_crop(input_image,1024)
    image_tensor = process_images([input_image], image_processor, model.config)
    image_tensor = [_image.to(dtype=torch.bfloat16, device=device) for _image in image_tensor]
    conv_template = "llada" 
    prompt = "Are there any person in the image? If so, what is he/she doing? Please also describe the look of the person. Additionlly, write a long story based on what you see in the image."
    question = f"<image>\n {prompt}."
    conv = copy.deepcopy(conv_templates[conv_template])
    conv.append_message(conv.roles[0], question)
    conv.append_message(conv.roles[1], None)
    prompt_question = conv.get_prompt()
    input_ids = tokenizer_image_token(prompt_question, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to(model.device)
    image_sizes = [input_image.size]
    t0 = time.time()
    res = model.generate(
        input_ids,
        images=image_tensor,
        image_sizes=image_sizes,
        do_sample=False,
        temperature=0,
        max_new_tokens=512,
        block_length=64,
        step_ratio=1, # 32 steps
        tokenizer=tokenizer,
        prefix_lm=True,
        verbose=False,
        use_fast_dlm=False,
        truncate=False
    )
    t1 = time.time()
    time_no_truncate = t1 - t0
    t0 = time.time()
    res = model.generate(
        input_ids,
        images=image_tensor,
        image_sizes=image_sizes,
        do_sample=False,
        temperature=0,
        max_new_tokens=512,
        block_length=64,
        step_ratio=1, # 32 steps
        tokenizer=tokenizer,
        prefix_lm=True,
        verbose=False,
        use_fast_dlm=True, # cache only
        truncate=False
    )
    t1 = time.time()
    time_cache_only = t1 - t0
    t0 = time.time()
    res = model.generate(
        input_ids,
        images=image_tensor,
        image_sizes=image_sizes,
        do_sample=False,
        temperature=0,
        max_new_tokens=512,
        block_length=64,
        step_ratio=1, # 32 steps
        tokenizer=tokenizer,
        prefix_lm=True,
        verbose=False,
        use_fast_dlm=True,
        truncate=True
    )
    t1 = time.time()
    time_truncate = t1 - t0
    print("Time prefix cache only:", time_no_truncate)
    print("Time cache only:", time_cache_only)
    print("Time cache + truncate:", time_truncate)
    res = tokenizer.batch_decode(res[0],skip_special_tokens=True)[0].replace('!','')
    print(res)
if __name__ == '__main__':
    run_text_to_image()
    # run_image_editing()
    # und()