diff --git a/.gitignore b/.gitignore index 163aa14..8035a56 100644 --- a/.gitignore +++ b/.gitignore @@ -4,22 +4,52 @@ __pycache__/ *$py.class *.so .Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST # Virtual environments venv/ ENV/ env/ +.venv -# IDE -.vscode/ +# PyCharm .idea/ -# Temporary files -*.sh -=* -GITHUB_ISSUE_RESPONSE.md +# VSCode +.vscode/ -# Distribution / packaging -build/ -dist/ -*.egg-info/ +# Jupyter Notebook +.ipynb_checkpoints + +# Model weights +*.bin +*.safetensors +*.ckpt +*.pth + +# Output files +output/ +output_*/ +*.json.tmp +temp_*/ + +# OS +.DS_Store +Thumbs.db diff --git a/README.md b/README.md index 7ddb973..ce8a3fb 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,7 @@ Install remaining dependencies: pip install -r requirements.txt ``` + ## 4. Inference | | GPU Memory Usage | Time per 256/512/1024/2048/4096 tokens | @@ -95,6 +96,22 @@ pip install -r requirements.txt **Note: The inference time shown here is measured per OmniLottie Lottie tokens, while the inference time reported in our paper is measured per JSON code tokens for fair comparison with baseline methods.** +### Model Format Support + +OmniLottie supports **two model formats**: + +1. **Original Format** (`inference.py` / `app.py`): + - Model file: `pytorch_model.bin` + - For users who downloaded the model before HuggingFace format support + +2. **šŸ¤— HuggingFace Format** (`inference_hf.py` / `app_hf.py`): + - Model files: `model-*.safetensors` + `config.json` + - Supports `from_pretrained()` API for automatic downloading + - Compatible with HuggingFace Hub ecosystem + - **Recommended for new users** + +Both formats produce identical results. Choose based on your model format. + ### Quick Start **Download Model Weights** @@ -106,11 +123,58 @@ pip install huggingface-hub **Download the model from Hugging Face:** ```bash -# Download OmniLottie model +# Download OmniLottie model (HuggingFace format with safetensors) huggingface-cli download OmniLottie/OmniLottie --local-dir /PATH/TO/OmniLottie ``` -**Try with Example Data** +### šŸ¤— Using HuggingFace Format (Recommended) + +If you downloaded the model in HuggingFace format (with `config.json` and `.safetensors` files), use `inference_hf.py` and `app_hf.py`: + +**Using from_pretrained() API (automatic download from HF Hub):** +```bash +# Text-to-Lottie +python inference_hf.py \ + --model_path OmniLottie/OmniLottie \ + --text "A bouncing ball" \ + --output output.json + +# Image-to-Lottie +python inference_hf.py \ + --model_path OmniLottie/OmniLottie \ + --image image.png \ + --text "rotating animation" \ + --output output.json + +# Video-to-Lottie +python inference_hf.py \ + --model_path OmniLottie/OmniLottie \ + --video video.mp4 \ + --output output.json +``` + +**Using local HuggingFace format model:** +```bash +python inference_hf.py \ + --model_path /PATH/TO/OmniLottie \ + --text "A spinning star" \ + --output output.json +``` + +**Launch Gradio demo (HuggingFace format):** +```bash +# Using local model +MODEL_PATH=/PATH/TO/OmniLottie python app_hf.py + +# Or using HF Hub (automatic download) +MODEL_PATH=OmniLottie/OmniLottie python app_hf.py +``` + +### Using Original Format + +If you have the original `pytorch_model.bin` format, use `inference.py` and `app.py`: + +**Try with Example Data (Original Format)** We provide example prompts, images, and videos in the `example/` directory: @@ -245,9 +309,18 @@ python inference.py \ ### Interactive Demo -We provide an interactive generation interface using Gradio: +We provide interactive generation interfaces using Gradio: -- **Local Deployment** +- **Local Deployment (HuggingFace Format - Recommended)** + ```bash + # Using local model + MODEL_PATH=/PATH/TO/OmniLottie python app_hf.py + + # Or using HF Hub (automatic download) + MODEL_PATH=OmniLottie/OmniLottie python app_hf.py + ``` + +- **Local Deployment (Original Format)** ```bash python app.py ``` diff --git a/app_hf.py b/app_hf.py new file mode 100644 index 0000000..8a5efbb --- /dev/null +++ b/app_hf.py @@ -0,0 +1,1017 @@ +""" +OmniLottie Gradio Demo - Hugging Face Compatible Version + +This version uses decoder_hf.py with from_pretrained() to load models, +supporting automatic downloading from Hugging Face Hub. + +Environment Variables: + MODEL_PATH: Model path or HF Hub ID (default: "OmniLottie/OmniLottie") + PROCESSOR_PATH: Processor path (default: "/mnt/jfs-test/Qwen2.5-VL-3B-Instruct") + +Usage: + # Using HF Hub model (automatic download) + python app_hf.py + + # Using local model + MODEL_PATH=/path/to/local/model python app_hf.py + + # Custom configuration + MODEL_PATH=your-org/your-model PROCESSOR_PATH=/path/to/processor python app_hf.py +""" + +import gradio as gr +import json +import torch +import os +import numpy as np +import random +import re +import tempfile +import base64 +import threading +import time +from PIL import Image as PILImage +from decord import VideoReader, cpu + +from decoder_hf import LottieDecoder # Use HF decoder with from_pretrained +from transformers import AutoProcessor +from qwen_vl_utils import process_vision_info +from lottie.objects.lottie_tokenize import LottieTensor +from lottie.objects.lottie_param import ( + from_sequence, ShapeLayer, NullLayer, PreCompLayer, TextLayer, + SolidColorLayer, Font, Chars, + shape_layer_to_json, null_layer_to_json, precomp_layer_to_json, + text_layer_to_json, solid_layer_to_json, font_to_json, char_to_json +) + +SYSTEM_PROMPT = "You are a Lottie animation expert." +VIDEO_PROMPT = "Turn this video into Lottie code." +LOTTIE_BOS = 192398 +LOTTIE_EOS = 192399 +PAD_TOKEN = 151643 + +model = None +processor = None +device = None + +generation_lock = threading.Lock() + +def load_model_once(): + """ + Load OmniLottie model using from_pretrained() for HuggingFace compatibility + Supports automatic downloading from HF Hub or loading from local path + """ + global model, processor, device + + if model is not None: + return model, processor, device + + # Model path - can be HF Hub ID or local path + # Examples: "OmniLottie/OmniLottie" or "/path/to/local/model" + model_path = os.environ.get("MODEL_PATH", "OmniLottie/OmniLottie") + + device = torch.device("cuda:0" if torch.cuda.is_available() else "xpu:0" if torch.xpu.is_available() else "cpu") + + print(f"Loading model from {model_path}...") + + # Load model using from_pretrained (supports HF Hub and local paths) + model = LottieDecoder.from_pretrained( + model_path, + torch_dtype=torch.bfloat16, + trust_remote_code=True + ) + + model = model.to(device).eval() + + # Load processor + processor_path = os.environ.get("PROCESSOR_PATH", "/mnt/jfs-test/Qwen2.5-VL-3B-Instruct") + processor = AutoProcessor.from_pretrained( + processor_path, + padding_side="left", + trust_remote_code=True + ) + + print(f"āœ… Model loaded on {device}") + + return model, processor, device + +def simplify_to_animation_description(text): + if not text or not isinstance(text, str): + return text + + prefixes = [ + r'^The video features?\s+', r'^The scene shows?\s+', + r'^An animation of\s+', r'^There is\s+', r'^It shows?\s+' + ] + + for pattern in prefixes: + text = re.sub(pattern, '', text, flags=re.IGNORECASE) + + if text: + text = text[0].upper() + text[1:] + + return text.strip() + +def add_random_background(img): + if img.mode != 'RGBA': + return img.convert('RGB') + + light_colors = [(255, 255, 255), (245, 245, 245), (250, 250, 250)] + bg_color = random.choice(light_colors) + background = PILImage.new('RGB', img.size, bg_color) + background.paste(img, (0, 0), img) + return background + +def load_frames_from_video(video_path, num_frames=8, max_size=336): + import os + ext = os.path.splitext(video_path)[1].lower() + + frames = [] + + if ext in ('.gif', '.webp'): + try: + img = PILImage.open(video_path) + total_frames = getattr(img, 'n_frames', 1) + + if total_frames < 1: + raise ValueError(f"No frames in {ext.upper()}: {video_path}") + + indices = np.linspace(0, total_frames - 1, min(num_frames, total_frames)).astype(int) + + for idx in indices: + img.seek(idx) + frame = img.convert('RGB') + if max(frame.size) > max_size: + frame.thumbnail((max_size, max_size), PILImage.LANCZOS) + frames.append(frame) + + img.close() + + except Exception as e: + raise ValueError(f"Failed to load {ext.upper()}: {str(e)}") + + else: + try: + vr = VideoReader(video_path, ctx=cpu(0)) + total_frames = len(vr) + + if total_frames < 1: + raise ValueError(f"Video has no frames: {video_path}") + + indices = np.linspace(0, total_frames - 1, num_frames).astype(int) + frames_np = vr.get_batch(indices).asnumpy() + + for f in frames_np: + img = PILImage.fromarray(f) + if max(img.size) > max_size: + img.thumbnail((max_size, max_size), PILImage.LANCZOS) + frames.append(img) + + except Exception as e: + raise ValueError(f"Failed to load video: {str(e)}") + + while len(frames) < num_frames: + frames.append(frames[-1].copy()) + + return frames + +def build_messages(task_type, text_prompt=None, image=None, video_frames=None): + messages = [{"role": "system", "content": SYSTEM_PROMPT}] + + if task_type == "text": + text = simplify_to_animation_description(text_prompt) + messages.append({ + "role": "user", + "content": [{"type": "text", "text": f"Generate Lottie code: {text}"}] + }) + + elif task_type == "image": + text = simplify_to_animation_description(text_prompt) + messages.append({ + "role": "user", + "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": f"Animate this image: {text}"} + ] + }) + + elif task_type == "video": + messages.append({ + "role": "user", + "content": [ + {"type": "video", "video": video_frames, "fps": 8.0}, + {"type": "text", "text": VIDEO_PROMPT} + ] + }) + + return messages + +def prepare_inference_input(processor, messages, device): + text_input = processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + + image_inputs, video_inputs = process_vision_info(messages) + + inputs = processor( + text=[text_input], + images=image_inputs if image_inputs else None, + videos=video_inputs if video_inputs else None, + padding=False, + return_tensors="pt" + ) + + input_ids = inputs['input_ids'] + attention_mask = inputs['attention_mask'] + target_len = 1500 + + if input_ids.shape[1] < target_len: + pad_len = target_len - input_ids.shape[1] + input_ids = torch.cat([ + torch.full((1, pad_len), PAD_TOKEN, dtype=torch.long), + input_ids + ], dim=1) + attention_mask = torch.cat([ + torch.zeros((1, pad_len), dtype=torch.long), + attention_mask + ], dim=1) + + result = { + 'input_ids': input_ids.to(device), + 'attention_mask': attention_mask.to(device), + 'pixel_values': inputs.get('pixel_values').to(device) if inputs.get('pixel_values') is not None else None, + 'image_grid_thw': inputs.get('image_grid_thw').to(device) if inputs.get('image_grid_thw') is not None else None, + 'pixel_values_videos': inputs.get('pixel_values_videos').to(device) if inputs.get('pixel_values_videos') is not None else None, + 'video_grid_thw': inputs.get('video_grid_thw').to(device) if inputs.get('video_grid_thw') is not None else None, + } + + return result + +def generate_lottie(model, inputs, max_tokens, device, use_sampling=False, temperature=0.95, top_p=0.25, top_k=5): + """ē”Ÿęˆ Lottie tokens""" + model.transformer.rope_deltas = None + position_ids, _ = model.transformer.get_rope_index( + input_ids=inputs['input_ids'], + attention_mask=inputs['attention_mask'], + image_grid_thw=inputs.get('image_grid_thw'), + video_grid_thw=inputs.get('video_grid_thw'), + ) + position_ids = position_ids * inputs['attention_mask'][None, ] + + kwargs = { + 'input_ids': inputs['input_ids'], + 'attention_mask': inputs['attention_mask'], + 'pixel_values': inputs.get('pixel_values'), + 'image_grid_thw': inputs.get('image_grid_thw'), + 'pixel_values_videos': inputs.get('pixel_values_videos'), + 'video_grid_thw': inputs.get('video_grid_thw'), + 'position_ids': position_ids, + 'max_new_tokens': max_tokens, + 'eos_token_id': LOTTIE_EOS, + 'pad_token_id': PAD_TOKEN, + 'use_cache': True, + } + + if use_sampling: + kwargs.update({'do_sample': True, 'temperature': temperature, 'top_p': top_p, 'top_k': top_k}) + else: + kwargs.update({'do_sample': False, 'num_beams': 1}) + + with torch.no_grad(): + outputs = model.transformer.generate(**kwargs) + + input_len = inputs['input_ids'].shape[1] + generated_ids = outputs[0][input_len:].tolist() + + del outputs, kwargs, position_ids + + if generated_ids and generated_ids[0] == LOTTIE_BOS: + generated_ids = generated_ids[1:] + if LOTTIE_EOS in generated_ids: + generated_ids = generated_ids[:generated_ids.index(LOTTIE_EOS)] + + return generated_ids + +def fix_lottie_json(anim): + + anim_ip = int(round(anim.get("ip", 0))) + anim_op = int(round(anim.get("op", 16))) + anim["ip"] = anim_ip + anim["op"] = anim_op + anim["fr"] = int(round(anim.get("fr", 8))) + anim["ddd"] = int(anim.get("ddd", 0)) + + def fix_t_recursive(obj): + if isinstance(obj, dict): + if obj.get("a") == 1 and isinstance(obj.get("k"), list): + for kf in obj["k"]: + if isinstance(kf, dict) and "t" in kf: + kf["t"] = int(round(kf["t"])) + for v in obj.values(): + fix_t_recursive(v) + elif isinstance(obj, list): + for item in obj: + fix_t_recursive(item) + + fix_t_recursive(anim) + + max_x = float(anim.get("w", 512)) + max_y = float(anim.get("h", 512)) + + def collect_pos(layer): + nonlocal max_x, max_y + p = layer.get("ks", {}).get("p", {}) + if isinstance(p, dict): + if p.get("a", 0) == 0: + pv = p.get("k", [0, 0]) + if isinstance(pv, list) and len(pv) >= 2: + max_x = max(max_x, float(pv[0])) + max_y = max(max_y, float(pv[1])) + else: + for kf in p.get("k", []): + if isinstance(kf, dict): + for sv in (kf.get("s", []), kf.get("e", [])): + if isinstance(sv, list) and len(sv) >= 2: + max_x = max(max_x, float(sv[0])) + max_y = max(max_y, float(sv[1])) + for sub in layer.get("layers", []): + collect_pos(sub) + + for layer in anim.get("layers", []): + collect_pos(layer) + + anim["w"] = max(512, int((max_x * 1.1 + 15) // 16 * 16)) + anim["h"] = max(512, int((max_y * 1.1 + 15) // 16 * 16)) + + valid_inds = set() + for layer in anim.get("layers", []): + if "ind" in layer: + valid_inds.add(int(layer["ind"])) + + def clean_shapes(shapes): + if not isinstance(shapes, list): + return shapes + cleaned = [] + for sh in shapes: + if not isinstance(sh, dict): + continue + if sh.get("ty") == "gr": + sh["it"] = clean_shapes(sh.get("it", [])) + if not sh["it"]: + continue + has_tr = any(item.get("ty") == "tr" for item in sh["it"] if isinstance(item, dict)) + if not has_tr: + sh["it"].append({ + "ty": "tr", "nm": "", + "a": {"a": 0, "k": [0, 0], "ix": 1}, + "p": {"a": 0, "k": [0, 0], "ix": 2}, + "s": {"a": 0, "k": [100, 100], "ix": 3}, + "r": {"a": 0, "k": 0, "ix": 6}, + "o": {"a": 0, "k": 100, "ix": 7}, + "sk": {"a": 0, "k": 0, "ix": 4}, + "sa": {"a": 0, "k": 0, "ix": 5}, + "hd": False + }) + cleaned.append(sh) + return cleaned + + def fix_layer(layer): + ip = int(round(layer.get("ip", anim_ip))) + op = int(round(layer.get("op", anim_op))) + layer["ip"] = max(anim_ip, ip) + layer["op"] = min(anim_op, max(layer["ip"] + 1, op)) + layer["st"] = int(round(layer.get("st", anim_ip))) + if "ind" in layer: + layer["ind"] = int(layer["ind"]) + if "parent" in layer: + p = int(layer["parent"]) + if p in valid_inds: + layer["parent"] = p + else: + del layer["parent"] + layer.pop("ct", None) + if "shapes" in layer: + layer["shapes"] = clean_shapes(layer["shapes"]) + for sub in layer.get("layers", []): + fix_layer(sub) + return layer + + fixed_layers = [] + for l in anim.get("layers", []): + fix_layer(l) + shapes = l.get("shapes", []) + if l.get("ty") == 4 and not shapes: + continue + fixed_layers.append(l) + anim["layers"] = fixed_layers + + for asset in anim.get("assets", []): + if "layers" in asset: + fixed = [] + for l in asset["layers"]: + fix_layer(l) + if l.get("ty") == 4 and not l.get("shapes"): + continue + fixed.append(l) + asset["layers"] = fixed + + return anim + +def tokens_to_lottie_json(generated_ids): + reconstructed_tensor = LottieTensor.from_list(generated_ids) + reconstructed_sequence = reconstructed_tensor.to_sequence() + reconstructed = from_sequence(reconstructed_sequence) + + json_animation = { + "v": reconstructed.get("v", "5.5.2"), + "fr": reconstructed.get("fr", 8), + "ip": reconstructed.get("ip", 0), + "op": reconstructed.get("op", 16), + "w": reconstructed.get("w", 512), + "h": reconstructed.get("h", 512), + "nm": reconstructed.get("nm", "Animation"), + "ddd": reconstructed.get("ddd", 0), + "assets": [], + "layers": [], + } + + if "fonts" in reconstructed and reconstructed["fonts"]: + fonts_data = reconstructed["fonts"] + if isinstance(fonts_data, dict) and "list" in fonts_data: + fonts_json = {"list": []} + for font in fonts_data["list"]: + if isinstance(font, Font): + fonts_json["list"].append(font_to_json(font)) + else: + fonts_json["list"].append(font) + json_animation["fonts"] = fonts_json + + if "chars" in reconstructed and reconstructed["chars"]: + chars_json = [] + for char in reconstructed["chars"]: + if isinstance(char, Chars): + chars_json.append(char_to_json(char)) + else: + chars_json.append(char) + json_animation["chars"] = chars_json + + for asset in reconstructed.get("assets", []): + asset_json = dict(asset) + if "layers" in asset: + asset_json["layers"] = [] + for layer in asset["layers"]: + if isinstance(layer, ShapeLayer): + asset_json["layers"].append(shape_layer_to_json(layer)) + elif isinstance(layer, NullLayer): + asset_json["layers"].append(null_layer_to_json(layer)) + elif isinstance(layer, PreCompLayer): + asset_json["layers"].append(precomp_layer_to_json(layer)) + elif isinstance(layer, TextLayer): + asset_json["layers"].append(text_layer_to_json(layer)) + elif isinstance(layer, SolidColorLayer): + asset_json["layers"].append(solid_layer_to_json(layer)) + else: + asset_json["layers"].append(layer) + json_animation["assets"].append(asset_json) + + for layer in reconstructed.get("layers", []): + if isinstance(layer, ShapeLayer): + json_animation["layers"].append(shape_layer_to_json(layer)) + elif isinstance(layer, NullLayer): + json_animation["layers"].append(null_layer_to_json(layer)) + elif isinstance(layer, PreCompLayer): + json_animation["layers"].append(precomp_layer_to_json(layer)) + elif isinstance(layer, TextLayer): + json_animation["layers"].append(text_layer_to_json(layer)) + elif isinstance(layer, SolidColorLayer): + json_animation["layers"].append(solid_layer_to_json(layer)) + else: + json_animation["layers"].append(layer) + + json_animation = fix_lottie_json(json_animation) + + return json_animation + +_lottie_js_cache = None + +def get_lottie_js(): + global _lottie_js_cache + if _lottie_js_cache is not None: + return _lottie_js_cache + + local_path = "lottie.min.js" + if os.path.exists(local_path): + with open(local_path, 'r', encoding='utf-8') as f: + _lottie_js_cache = f.read() + print(f"āœ… Loaded local lottie.min.js ({len(_lottie_js_cache)} bytes)") + else: + _lottie_js_cache = '' + print("āš ļø Using CDN lottie.min.js (local file not found)") + + return _lottie_js_cache + +def create_lottie_html(animation_data, height=600): + bg_style = """ + background-image: + linear-gradient(45deg, #666666 25%, transparent 25%), + linear-gradient(-45deg, #666666 25%, transparent 25%), + linear-gradient(45deg, transparent 75%, #666666 75%), + linear-gradient(-45deg, transparent 75%, #666666 75%); + background-size: 16px 16px; + background-position: 0 0, 0 8px, 8px -8px, -8px 0px; + background-color: #444444; + """ + + lottie_js = get_lottie_js() + + if lottie_js.startswith('{lottie_js}" + + animation_json_escaped = json.dumps(animation_data).replace('\\', '\\\\').replace("'", "\\'") + + anim_width = animation_data.get('w', 512) + anim_height = animation_data.get('h', 512) + + inner_html = f""" + + + + {lottie_script} + + + +
+
+
+ + +""" + + inner_html_b64 = base64.b64encode(inner_html.encode('utf-8')).decode('utf-8') + iframe_html = f'' + + return iframe_html + +def save_json_to_temp(lottie_json): + if lottie_json is None: + return None + + temp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False, prefix='lottie_') + json.dump(lottie_json, temp_file, indent=2) + temp_file.close() + return temp_file.name + +def process_text_to_lottie(text_prompt, max_tokens, use_sampling, temperature, top_p, top_k): + with generation_lock: + try: + start_time = time.time() + + if not text_prompt or not text_prompt.strip(): + return None, "āŒ Please enter a text description", None + model, processor, device = load_model_once() + + messages = build_messages("text", text_prompt=text_prompt) + inputs = prepare_inference_input(processor, messages, device) + + generated_ids = generate_lottie( + model, inputs, max_tokens, device, + use_sampling, temperature, top_p, top_k + ) + + del inputs + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if torch.xpu.is_available(): + torch.xpu.empty_cache() + + lottie_json = tokens_to_lottie_json(generated_ids) + + html = create_lottie_html(lottie_json, height=600) + + elapsed_time = time.time() - start_time + + status = f"āœ… Generated {len(generated_ids)} tokens | Layers: {len(lottie_json.get('layers', []))} | {lottie_json.get('fr', 8)} fps | Time: {elapsed_time:.1f}s" + + temp_path = save_json_to_temp(lottie_json) + + return html, status, temp_path + + except Exception as e: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if torch.xpu.is_available(): + torch.xpu.empty_cache() + return None, f"āŒ Error: {str(e)}", None + +def load_image_from_file(file_path): + + ext = os.path.splitext(file_path)[1].lower() + + if ext == '.svg': + import cairosvg + import io + png_bytes = cairosvg.svg2png(url=file_path, output_width=448, output_height=448) + image = PILImage.open(io.BytesIO(png_bytes)) + else: + image = PILImage.open(file_path) + + if image.mode == 'RGBA': + image = add_random_background(image) + else: + image = image.convert('RGB') + + return image + + +def process_image_to_lottie(image_file, text_description, max_tokens, use_sampling, temperature, top_p, top_k): + with generation_lock: + try: + start_time = time.time() + + if image_file is None: + return None, "āŒ Please upload an image", None + + model, processor, device = load_model_once() + + image = load_image_from_file(image_file) + + image = image.resize((448, 448), PILImage.LANCZOS) + + desc = text_description if text_description else "A simple animation" + messages = build_messages("image", text_prompt=desc, image=image) + inputs = prepare_inference_input(processor, messages, device) + + generated_ids = generate_lottie( + model, inputs, max_tokens, device, + use_sampling, temperature, top_p, top_k + ) + + del inputs + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if torch.xpu.is_available(): + torch.xpu.empty_cache() + + lottie_json = tokens_to_lottie_json(generated_ids) + + html = create_lottie_html(lottie_json, height=600) + + elapsed_time = time.time() - start_time + + status = f"āœ… Generated {len(generated_ids)} tokens | Layers: {len(lottie_json.get('layers', []))} | {lottie_json.get('fr', 8)} fps | Time: {elapsed_time:.1f}s" + + temp_path = save_json_to_temp(lottie_json) + + return html, status, temp_path + + except Exception as e: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if torch.xpu.is_available(): + torch.xpu.empty_cache() + return None, f"āŒ Error: {str(e)}", None + +def process_video_to_lottie(video, max_tokens, use_sampling, temperature, top_p, top_k): + with generation_lock: + try: + start_time = time.time() + + if video is None: + return None, "āŒ Please upload a video/GIF/WebP file", None + + import os + ext = os.path.splitext(video)[1].lower() if isinstance(video, str) else '' + if ext not in ('.mp4', '.avi', '.mov', '.gif', '.webp'): + return None, f"āŒ Unsupported format: {ext}. Please upload MP4/AVI/MOV/GIF/WebP", None + + model, processor, device = load_model_once() + + frames = load_frames_from_video(video, num_frames=8) + + messages = build_messages("video", video_frames=frames) + inputs = prepare_inference_input(processor, messages, device) + + generated_ids = generate_lottie( + model, inputs, max_tokens, device, + use_sampling, temperature, top_p, top_k + ) + + del inputs + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if torch.xpu.is_available(): + torch.xpu.empty_cache() + + lottie_json = tokens_to_lottie_json(generated_ids) + + html = create_lottie_html(lottie_json, height=600) + + elapsed_time = time.time() - start_time + + status = f"āœ… Generated {len(generated_ids)} tokens (from {len(frames)} frames) | Layers: {len(lottie_json.get('layers', []))} | {lottie_json.get('fr', 8)} fps | Time: {elapsed_time:.1f}s" + + temp_path = save_json_to_temp(lottie_json) + + return html, status, temp_path + + except Exception as e: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if torch.xpu.is_available(): + torch.xpu.empty_cache() + return None, f"āŒ Error: {str(e)}", None + +def create_gradio_interface(): + + with gr.Blocks(title="OmniLottie Demo Page", theme=gr.themes.Soft()) as demo: + gr.Markdown("# šŸŽØ OmniLottie Demo Page") + gr.Markdown("Offical Demo Page of OmniLottie") + gr.Markdown("Generate Lottie animations from text, images, or videos") + + with gr.Tabs() as tabs: + with gr.Tab("šŸ“ Text-to-Lottie"): + gr.Markdown(""" + ### šŸ’” Prompt Tips for Better Results + + **Good prompts should describe:** + 1. **Main Object**: What is being animated (e.g., "a blue bird", "a yellow folder icon", "an orange piggy bank") + 2. **Motion Pattern**: How it moves (e.g., "appearing, pulsing while sliding", "fading in, floating toward", "bouncing up and down") + 3. **Direction**: Where it moves (e.g., "downward", "toward the top-left", "back to its start") + 4. **Loop Behavior**: How it repeats (e.g., "repeating seamlessly", "looping smoothly", "repeating continuously") + + **Example Patterns:** + - šŸ”„ **Simple Loop**: "a red ball appearing, bouncing up and down, then fading out, repeating seamlessly" + - šŸŽÆ **Movement**: "a blue arrow sliding from left to right, then quickly returning to start, looping continuously" + - šŸ’« **Transformation**: "a yellow star fading in while rotating 360 degrees, holds briefly, then fading out, repeating smoothly" + - šŸŽØ **Static Icon**: "static illustration of a cartoon cat's face with a cute expression, light orange body, red inner ears" + - šŸ‘¤ **Character**: "animated cartoon figure dressed in a beige suit with a white shirt, holding a gray tablet" + + **Pro Tips:** + - Be specific about colors, shapes, and movements + - Describe motion phases clearly (appear ā†’ move ā†’ hold ā†’ return) + - Use descriptive motion verbs: sliding, pulsing, drifting, bouncing, rotating, fading + - For icons: include style details (outline, colors, decorations) + """) + + with gr.Row(): + with gr.Column(scale=1): + text_input = gr.Textbox( + label="Text Description", + placeholder="Example: a blue bird appearing, pulsing while sliding downward, lingers briefly, then growing back while sliding upward to reset, repeating seamlessly", + lines=5 + ) + + with gr.Accordion("āš™ļø Generation Settings", open=False): + gr.Markdown(""" + **Parameter Guide:** + - **Max Tokens**: Higher = more complex animations (slower), Lower = simpler animations (faster) + - **Top-p & Top-k**: Higher = more random/creative, Lower = more stable/consistent + - **Temperature**: Higher = more diverse, Lower = more deterministic + + šŸ’” **Quick Tips:** + - For complex animations: increase max tokens to 5856 + - For faster generation: reduce max tokens to 2048-3072 + - For more creative results: increase top-p (0.5-0.8) and top-k (20-50) + - For consistent results: decrease top-p (0.1-0.25) and top-k (5-10) + """) + text_max_tokens = gr.Slider(512, 5856, value=5556, step=256, label="Max Tokens") + text_use_sampling = gr.Checkbox(label="Use Sampling", value=True) + text_temperature = gr.Slider(0.1, 2.0, value=0.9, step=0.1, label="Temperature") + text_top_p = gr.Slider(0.1, 1.0, value=0.25, step=0.1, label="Top-p") + text_top_k = gr.Slider(1, 100, value=5, step=1, label="Top-k") + + text_generate_btn = gr.Button("šŸš€ Generate", variant="primary", size="lg") + + # Generation time tips + gr.Markdown(""" + ā±ļø **Generation Time:** + - Simple icons/shapes: ~30-60 seconds (1000-2000 tokens) + - Medium animations: ~1-2 minutes (2000-3500 tokens) + - Complex characters: ~4-5 minutes (4500-6000 tokens) + + Please be patient! Complex animations take time to generate. ā˜• + """) + + text_status = gr.Markdown() + + with gr.Column(scale=1): + text_output = gr.HTML(label="Animation Preview") + text_json_file = gr.File(label="šŸ“„ Download JSON", visible=True) + + def get_text_examples(): + examples = [] + demo_txt_path = "./example/demo.txt" + if os.path.exists(demo_txt_path): + with open(demo_txt_path, 'r', encoding='utf-8') as f: + lines = [line.strip() for line in f.readlines() if line.strip()] + examples = [[line] for line in lines[:50]] + return examples + + gr.Examples( + examples=get_text_examples(), + inputs=text_input, + label="šŸ“‚ Example Prompts (Click to Load)", + examples_per_page=10, + cache_examples=False + ) + + text_generate_btn.click( + fn=process_text_to_lottie, + inputs=[text_input, text_max_tokens, text_use_sampling, text_temperature, text_top_p, text_top_k], + outputs=[text_output, text_status, text_json_file] + ) + + with gr.Tab("šŸ–¼ļø Text+Image-to-Lottie"): + with gr.Row(): + with gr.Column(scale=1): + image_input = gr.Image( + label="Upload Image", + type="filepath", + sources=["upload"] + ) + image_text_input = gr.Textbox( + label="Animation Description", + placeholder="Example: The object rotates 360 degrees", + lines=3 + ) + + with gr.Accordion("āš™ļø Generation Settings", open=False): + gr.Markdown(""" + **Parameter Guide:** + - **Max Tokens**: Higher = more complex animations (slower), Lower = simpler (faster) + - **Top-p & Top-k**: Higher = more creative/random, Lower = more stable + - **Temperature**: Controls output diversity + """) + image_max_tokens = gr.Slider(512, 5556, value=5556, step=256, label="Max Tokens") + image_use_sampling = gr.Checkbox(label="Use Sampling", value=True) + image_temperature = gr.Slider(0.1, 2.0, value=0.9, step=0.1, label="Temperature") + image_top_p = gr.Slider(0.1, 1.0, value=0.25, step=0.05, label="Top-p") + image_top_k = gr.Slider(1, 100, value=5, step=1, label="Top-k") + + image_generate_btn = gr.Button("šŸš€ Generate", variant="primary", size="lg") + + # Generation time tips + gr.Markdown(""" + ā±ļø **Generation Time:** ~1-5 minutes depending on complexity + """) + + image_status = gr.Markdown() + + with gr.Column(scale=1): + image_output = gr.HTML(label="Animation Preview") + image_json_file = gr.File(label="šŸ“„ Download JSON", visible=True) + + def get_image_text_examples(): + examples = [] + demo_images_dir = "./example/demo_images" + if os.path.exists(demo_images_dir): + png_files = sorted([f for f in os.listdir(demo_images_dir) if f.endswith('.png')]) + for png_file in png_files[:50]: + base_name = os.path.splitext(png_file)[0] + txt_file = os.path.join(demo_images_dir, f"{base_name}.txt") + png_path = os.path.join(demo_images_dir, png_file) + if os.path.exists(txt_file): + with open(txt_file, 'r', encoding='utf-8') as f: + text_desc = f.read().strip() + examples.append([png_path, text_desc]) + return examples + + gr.Examples( + examples=get_image_text_examples(), + inputs=[image_input, image_text_input], + label="šŸ“‚ Example Images (Click to Load)", + examples_per_page=5, + cache_examples=False + ) + + image_generate_btn.click( + fn=process_image_to_lottie, + inputs=[image_input, image_text_input, image_max_tokens, image_use_sampling, + image_temperature, image_top_p, image_top_k], + outputs=[image_output, image_status, image_json_file] + ) + + # Tab 3: Video-to-Lottie + with gr.Tab("šŸŽ„ Video-to-Lottie"): + with gr.Row(): + with gr.Column(scale=1): + video_input = gr.Video( + label="Upload Video / GIF / WebP", + sources=["upload"] + ) + + with gr.Accordion("āš™ļø Generation Settings", open=False): + gr.Markdown(""" + **Parameter Guide:** + - **Max Tokens**: Higher = more complex animations (slower), Lower = simpler (faster) + - **Top-p & Top-k**: Higher = more creative/random, Lower = more stable + - **Temperature**: Controls output diversity + """) + video_max_tokens = gr.Slider(512, 5556, value=5556, step=256, label="Max Tokens") + video_use_sampling = gr.Checkbox(label="Use Sampling", value=True) + video_temperature = gr.Slider(0.1, 2.0, value=0.9, step=0.1, label="Temperature") + video_top_p = gr.Slider(0.1, 1.0, value=0.25, step=0.05, label="Top-p") + video_top_k = gr.Slider(1, 100, value=5, step=1, label="Top-k") + + video_generate_btn = gr.Button("šŸš€ Generate", variant="primary", size="lg") + + # Generation time tips + gr.Markdown(""" + ā±ļø **Generation Time:** ~2-5 minutes depending on video complexity + """) + + video_status = gr.Markdown() + + with gr.Column(scale=1): + video_output = gr.HTML(label="Animation Preview") + video_json_file = gr.File(label="šŸ“„ Download JSON", visible=True) + + def get_video_examples(): + examples = [] + demo_video_dir = "./example/demo_video" + if os.path.exists(demo_video_dir): + video_files = sorted([os.path.join(demo_video_dir, f) + for f in os.listdir(demo_video_dir) + if f.endswith('.mp4')]) + examples = [[vf] for vf in video_files[:50]] + return examples + + gr.Examples( + examples=get_video_examples(), + inputs=video_input, + label="šŸ“‚ Example Videos (Click to Load)", + examples_per_page=5, + cache_examples=False + ) + + video_generate_btn.click( + fn=process_video_to_lottie, + inputs=[video_input, video_max_tokens, video_use_sampling, + video_temperature, video_top_p, video_top_k], + outputs=[video_output, video_status, video_json_file] + ) + + gr.Markdown("---") + gr.Markdown(""" + āš ļø **Important Note:** This demo processes one request at a time. + If another user is generating, your request will wait in queue until the current one completes. + """) + + return demo + +if __name__ == "__main__": + demo = create_gradio_interface() + demo.launch( + server_name="0.0.0.0", + server_port=7861, + share=False, + show_error=True + ) diff --git a/configuration_lottie_decoder.py b/configuration_lottie_decoder.py new file mode 100644 index 0000000..497c74d --- /dev/null +++ b/configuration_lottie_decoder.py @@ -0,0 +1,54 @@ +""" +OmniLottie Decoder Configuration +""" +from transformers import PretrainedConfig +from typing import Optional + + +class LottieDecoderConfig(PretrainedConfig): + """ + Configuration class for LottieDecoder model, inheriting from PretrainedConfig + + Stores configuration parameters for the LottieDecoder model, + supporting Hugging Face's standard save and load mechanisms. + + Args: + pix_len (int): Maximum length for image/video tokens, default 4560 + text_len (int): Maximum length for text tokens, default 1500 + base_model_path (str): Path or name of base Qwen2.5-VL model + vocab_size (int): Vocabulary size, extended to 192400 to support Lottie tokens + bos_token_id (int): Beginning-of-sequence token ID for Lottie + eos_token_id (int): End-of-sequence token ID for Lottie + pad_token_id (int): Padding token ID + torch_dtype (str): Model weight data type, default "bfloat16" + attn_implementation (str): Attention implementation method, default "eager" + """ + + model_type = "lottie_decoder" + + def __init__( + self, + pix_len: int = 4560, + text_len: int = 1500, + base_model_path: str = "Qwen/Qwen2.5-VL-3B-Instruct", + vocab_size: int = 192400, + bos_token_id: int = 192398, + eos_token_id: int = 192399, + pad_token_id: int = 151643, + torch_dtype: str = "bfloat16", + attn_implementation: str = "eager", + **kwargs + ): + super().__init__( + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + **kwargs + ) + + self.pix_len = pix_len + self.text_len = text_len + self.base_model_path = base_model_path + self.vocab_size = vocab_size + self.torch_dtype = torch_dtype + self.attn_implementation = attn_implementation diff --git a/decoder_hf.py b/decoder_hf.py new file mode 100644 index 0000000..cd1fd62 --- /dev/null +++ b/decoder_hf.py @@ -0,0 +1,179 @@ +import torch +import torch.nn as nn +from transformers import Qwen2_5_VLForConditionalGeneration, AutoConfig, PreTrainedModel +from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLCausalLMOutputWithPast +from typing import Any, Dict, List, Optional, Tuple, Union +import os + +import transformers.models.qwen2_5_vl.modeling_qwen2_5_vl as qwen_modeling + +from configuration_lottie_decoder import LottieDecoderConfig + + +class LottieDecoder(PreTrainedModel): + """ + Autoregressive generative model for OmniLottie + + Lottie animation generation model based on Qwen2.5-VL, + supports generating Lottie JSON code from videos. + """ + + config_class = LottieDecoderConfig + base_model_prefix = "lottie_decoder" + supports_gradient_checkpointing = True + + def __init__(self, config: LottieDecoderConfig): + """ + Initialize LottieDecoder model + + Args: + config (LottieDecoderConfig): Model configuration object + """ + super().__init__(config) + + self.config = config + self.pix_len = config.pix_len + self.text_len = config.text_len + self.vocab_size = config.vocab_size + self.bos_token_id = config.bos_token_id + self.eos_token_id = config.eos_token_id + self.pad_token_id = config.pad_token_id + + print(f"Initializing LottieDecoder with base model: {config.base_model_path}") + + # Create base model configuration + qwen_config = AutoConfig.from_pretrained( + config.base_model_path, + vocab_size=self.vocab_size, + bos_token_id=self.bos_token_id, + eos_token_id=self.eos_token_id, + pad_token_id=self.pad_token_id, + trust_remote_code=True + ) + + # Load base Qwen2.5-VL model + self.transformer = Qwen2_5_VLForConditionalGeneration.from_pretrained( + config.base_model_path, + config=qwen_config, + torch_dtype=getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype, + attn_implementation=config.attn_implementation, + ignore_mismatched_sizes=True + ) + + # Extend vocabulary to support Lottie tokens + self.transformer.resize_token_embeddings(self.vocab_size) + + # Set to training mode initially (same as original decoder) + self.train() + + print(f"LottieDecoder initialized successfully. Vocab size: {self.vocab_size}") + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): + """ + Load LottieDecoder from pretrained model path + + Supports two loading methods: + 1. Load from Hugging Face standard format (recommended) + 2. Load from old format pytorch_model.bin (backward compatible) + """ + # Check if it's old format (contains pytorch_model.bin) + if os.path.isdir(pretrained_model_name_or_path): + old_format_path = os.path.join(pretrained_model_name_or_path, 'pytorch_model.bin') + if os.path.exists(old_format_path) and not os.path.exists(os.path.join(pretrained_model_name_or_path, 'config.json')): + print(f"Detected old format model, loading from {old_format_path}...") + return cls._from_old_format(pretrained_model_name_or_path, **kwargs) + + # Use standard Hugging Face loading process + return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) + + @classmethod + def _from_old_format(cls, checkpoint_path, **kwargs): + """ + Load model from old format (pytorch_model.bin) + + Args: + checkpoint_path: Directory path containing pytorch_model.bin + """ + # Extract configuration parameters + pix_len = kwargs.pop('pix_len', 4560) + text_len = kwargs.pop('text_len', 1500) + base_model_path = kwargs.pop('base_model_path', 'Qwen/Qwen2.5-VL-3B-Instruct') + + # Create configuration + config = LottieDecoderConfig( + pix_len=pix_len, + text_len=text_len, + base_model_path=base_model_path + ) + + # Initialize model + model = cls(config) + + # Load weights + model_file = os.path.join(checkpoint_path, 'pytorch_model.bin') + if os.path.exists(model_file): + state_dict = torch.load(model_file, map_location='cpu') + model.load_state_dict(state_dict, strict=False) + print(f"Successfully loaded weights from {model_file}") + else: + print(f"Warning: Model file not found {model_file}") + + return model + + def forward( + self, + input_ids=None, + attention_mask=None, + pixel_values=None, + image_grid_thw=None, + pixel_values_videos=None, + video_grid_thw=None, + labels=None, + past_key_values=None, + use_cache=False, + **kwargs + ): + """ + Forward pass - currently for inference only, needs implementation for training + """ + return self.transformer( + input_ids=input_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + pixel_values_videos=pixel_values_videos, + video_grid_thw=video_grid_thw, + labels=labels, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs + ) + + def generate(self, *args, **kwargs): + """ + Generate Lottie tokens + + Directly calls the underlying transformer's generate method + """ + return self.transformer.generate(*args, **kwargs) + + def get_input_embeddings(self): + """Get input embeddings""" + return self.transformer.get_input_embeddings() + + def set_input_embeddings(self, value): + """Set input embeddings""" + self.transformer.set_input_embeddings(value) + + def get_output_embeddings(self): + """Get output embeddings""" + return self.transformer.get_output_embeddings() + + def set_output_embeddings(self, new_embeddings): + """Set output embeddings""" + self.transformer.set_output_embeddings(new_embeddings) + + def resize_token_embeddings(self, new_num_tokens: Optional[int] = None): + """Resize token embeddings""" + return self.transformer.resize_token_embeddings(new_num_tokens) diff --git a/inference_hf.py b/inference_hf.py new file mode 100644 index 0000000..08d85b8 --- /dev/null +++ b/inference_hf.py @@ -0,0 +1,461 @@ +""" +OmniLottie Inference Script - Hugging Face Compatible Version + +Uses decoder_hf.py with from_pretrained() to load models. +Supports automatic model downloading from Hugging Face Hub. + +Usage: + # Text-to-Lottie (from HF Hub) + python inference_hf.py --model_path OmniLottie/OmniLottie --text "A bouncing ball" + + # Video-to-Lottie (local model) + python inference_hf.py --model_path ./model --video video.mp4 + + # Image-to-Lottie + python inference_hf.py --model_path OmniLottie/OmniLottie --image image.png --text "rotating animation" +""" + +import os +import torch +import argparse +import json +import re +from pathlib import Path +from PIL import Image +import numpy as np +from decord import VideoReader, cpu + +# Import HF-compatible model +from decoder_hf import LottieDecoder +from transformers import AutoProcessor +from qwen_vl_utils import process_vision_info + +# Import Lottie conversion tools +from lottie.objects.lottie_tokenize import LottieTensor +from lottie.objects.lottie_param import ( + from_sequence, ShapeLayer, NullLayer, PreCompLayer, TextLayer, + SolidColorLayer, Font, Chars, + shape_layer_to_json, null_layer_to_json, precomp_layer_to_json, + text_layer_to_json, solid_layer_to_json, font_to_json, char_to_json +) + +# Constants +SYSTEM_PROMPT = "You are a Lottie animation expert." +VIDEO_PROMPT = "Turn this video into Lottie code." +LOTTIE_BOS = 192398 +LOTTIE_EOS = 192399 +PAD_TOKEN = 151643 + + +def simplify_to_animation_description(text): + """Simplify text prompt to animation description""" + if not text or not isinstance(text, str): + return text + + prefixes = [ + r'^The video features?\s+', r'^The scene shows?\s+', + r'^An animation of\s+', r'^There is\s+', r'^It shows?\s+' + ] + + for pattern in prefixes: + text = re.sub(pattern, '', text, flags=re.IGNORECASE) + + if text: + text = text[0].upper() + text[1:] + + return text.strip() + + +def load_frames_from_video(video_path, num_frames=8, max_size=336): + """Load frames from video file (matches app_hf.py)""" + ext = os.path.splitext(video_path)[1].lower() + frames = [] + + if ext in ('.gif', '.webp'): + try: + img = Image.open(video_path) + total_frames = getattr(img, 'n_frames', 1) + + if total_frames < 1: + raise ValueError(f"No frames in {ext.upper()}: {video_path}") + + indices = np.linspace(0, total_frames - 1, min(num_frames, total_frames)).astype(int) + + for idx in indices: + img.seek(idx) + frame = img.convert('RGB') + if max(frame.size) > max_size: + frame.thumbnail((max_size, max_size), Image.LANCZOS) + frames.append(frame) + + img.close() + + except Exception as e: + raise ValueError(f"Failed to load {ext.upper()}: {str(e)}") + + else: + try: + vr = VideoReader(video_path, ctx=cpu(0)) + total_frames = len(vr) + + if total_frames < 1: + raise ValueError(f"Video has no frames: {video_path}") + + indices = np.linspace(0, total_frames - 1, num_frames).astype(int) + frames_np = vr.get_batch(indices).asnumpy() + + for f in frames_np: + img = Image.fromarray(f) + if max(img.size) > max_size: + img.thumbnail((max_size, max_size), Image.LANCZOS) + frames.append(img) + + except Exception as e: + raise ValueError(f"Failed to load video: {str(e)}") + + while len(frames) < num_frames: + frames.append(frames[-1].copy()) + + return frames + + +def build_messages(task_type, text_prompt=None, image=None, video_frames=None): + """Build messages for inference (matches app_hf.py)""" + messages = [{"role": "system", "content": SYSTEM_PROMPT}] + + if task_type == "text": + text = simplify_to_animation_description(text_prompt) + messages.append({ + "role": "user", + "content": [{"type": "text", "text": f"Generate Lottie code: {text}"}] + }) + + elif task_type == "image": + text = simplify_to_animation_description(text_prompt) if text_prompt else "A simple animation" + messages.append({ + "role": "user", + "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": f"Animate this image: {text}"} + ] + }) + + elif task_type == "video": + messages.append({ + "role": "user", + "content": [ + {"type": "video", "video": video_frames, "fps": 8.0}, + {"type": "text", "text": VIDEO_PROMPT} + ] + }) + + return messages + + +def prepare_inference_input(processor, messages, device): + """Prepare input for inference (matches app_hf.py exactly)""" + text_input = processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + + image_inputs, video_inputs = process_vision_info(messages) + + inputs = processor( + text=[text_input], + images=image_inputs if image_inputs else None, + videos=video_inputs if video_inputs else None, + padding=False, + return_tensors="pt" + ) + + input_ids = inputs['input_ids'] + attention_mask = inputs['attention_mask'] + target_len = 1500 + + if input_ids.shape[1] < target_len: + pad_len = target_len - input_ids.shape[1] + input_ids = torch.cat([ + torch.full((1, pad_len), PAD_TOKEN, dtype=torch.long), + input_ids + ], dim=1) + attention_mask = torch.cat([ + torch.zeros((1, pad_len), dtype=torch.long), + attention_mask + ], dim=1) + + result = { + 'input_ids': input_ids.to(device), + 'attention_mask': attention_mask.to(device), + 'pixel_values': inputs.get('pixel_values').to(device) if inputs.get('pixel_values') is not None else None, + 'image_grid_thw': inputs.get('image_grid_thw').to(device) if inputs.get('image_grid_thw') is not None else None, + 'pixel_values_videos': inputs.get('pixel_values_videos').to(device) if inputs.get('pixel_values_videos') is not None else None, + 'video_grid_thw': inputs.get('video_grid_thw').to(device) if inputs.get('video_grid_thw') is not None else None, + } + + return result + + +def generate_lottie(model, inputs, max_tokens, device, use_sampling=False, temperature=0.9, top_p=0.25, top_k=5): + """Generate Lottie tokens (matches app_hf.py exactly)""" + model.transformer.rope_deltas = None + position_ids, _ = model.transformer.get_rope_index( + input_ids=inputs['input_ids'], + attention_mask=inputs['attention_mask'], + image_grid_thw=inputs.get('image_grid_thw'), + video_grid_thw=inputs.get('video_grid_thw'), + ) + position_ids = position_ids * inputs['attention_mask'][None, ] + + kwargs = { + 'input_ids': inputs['input_ids'], + 'attention_mask': inputs['attention_mask'], + 'pixel_values': inputs.get('pixel_values'), + 'image_grid_thw': inputs.get('image_grid_thw'), + 'pixel_values_videos': inputs.get('pixel_values_videos'), + 'video_grid_thw': inputs.get('video_grid_thw'), + 'position_ids': position_ids, + 'max_new_tokens': max_tokens, + 'eos_token_id': LOTTIE_EOS, + 'pad_token_id': PAD_TOKEN, + 'use_cache': True, + } + + if use_sampling: + kwargs.update({'do_sample': True, 'temperature': temperature, 'top_p': top_p, 'top_k': top_k}) + else: + kwargs.update({'do_sample': False, 'num_beams': 1}) + + with torch.no_grad(): + outputs = model.transformer.generate(**kwargs) + + input_len = inputs['input_ids'].shape[1] + generated_ids = outputs[0][input_len:].tolist() + + del outputs, kwargs, position_ids + + if generated_ids and generated_ids[0] == LOTTIE_BOS: + generated_ids = generated_ids[1:] + if LOTTIE_EOS in generated_ids: + generated_ids = generated_ids[:generated_ids.index(LOTTIE_EOS)] + + return generated_ids + + +def tokens_to_lottie_json(generated_ids): + """Convert generated tokens to Lottie JSON format""" + reconstructed_tensor = LottieTensor.from_list(generated_ids) + reconstructed_sequence = reconstructed_tensor.to_sequence() + reconstructed = from_sequence(reconstructed_sequence) + + json_animation = { + "v": reconstructed.get("v", "5.5.2"), + "fr": reconstructed.get("fr", 8), + "ip": reconstructed.get("ip", 0), + "op": reconstructed.get("op", 16), + "w": reconstructed.get("w", 512), + "h": reconstructed.get("h", 512), + "nm": reconstructed.get("nm", "Animation"), + "ddd": reconstructed.get("ddd", 0), + "assets": [], + "layers": [], + } + + # Process fonts + if "fonts" in reconstructed and reconstructed["fonts"]: + fonts_data = reconstructed["fonts"] + if isinstance(fonts_data, dict) and "list" in fonts_data: + fonts_json = {"list": []} + for font in fonts_data["list"]: + if isinstance(font, Font): + fonts_json["list"].append(font_to_json(font)) + else: + fonts_json["list"].append(font) + json_animation["fonts"] = fonts_json + + # Process chars + if "chars" in reconstructed and reconstructed["chars"]: + chars_json = [] + for char in reconstructed["chars"]: + if isinstance(char, Chars): + chars_json.append(char_to_json(char)) + else: + chars_json.append(char) + json_animation["chars"] = chars_json + + # Process assets + for asset in reconstructed.get("assets", []): + asset_json = dict(asset) + if "layers" in asset: + asset_json["layers"] = [] + for layer in asset["layers"]: + if isinstance(layer, ShapeLayer): + asset_json["layers"].append(shape_layer_to_json(layer)) + elif isinstance(layer, NullLayer): + asset_json["layers"].append(null_layer_to_json(layer)) + elif isinstance(layer, PreCompLayer): + asset_json["layers"].append(precomp_layer_to_json(layer)) + elif isinstance(layer, TextLayer): + asset_json["layers"].append(text_layer_to_json(layer)) + elif isinstance(layer, SolidColorLayer): + asset_json["layers"].append(solid_layer_to_json(layer)) + else: + asset_json["layers"].append(layer) + json_animation["assets"].append(asset_json) + + # Process layers + for layer in reconstructed.get("layers", []): + if isinstance(layer, ShapeLayer): + json_animation["layers"].append(shape_layer_to_json(layer)) + elif isinstance(layer, NullLayer): + json_animation["layers"].append(null_layer_to_json(layer)) + elif isinstance(layer, PreCompLayer): + json_animation["layers"].append(precomp_layer_to_json(layer)) + elif isinstance(layer, TextLayer): + json_animation["layers"].append(text_layer_to_json(layer)) + elif isinstance(layer, SolidColorLayer): + json_animation["layers"].append(solid_layer_to_json(layer)) + else: + json_animation["layers"].append(layer) + + return json_animation + + +def main(): + parser = argparse.ArgumentParser(description='OmniLottie Inference (HF Compatible)') + + # Model arguments + parser.add_argument('--model_path', type=str, required=True, + help='Model path (local path or HF Hub ID, e.g. OmniLottie/OmniLottie)') + parser.add_argument('--processor_path', type=str, default='/mnt/jfs-test/Qwen2.5-VL-3B-Instruct', + help='Processor path (local path or HF Hub ID)') + + # Input arguments (choose one) + parser.add_argument('--text', type=str, help='Text prompt') + parser.add_argument('--image', type=str, help='Image path') + parser.add_argument('--video', type=str, help='Video path') + + # Output arguments + parser.add_argument('--output', type=str, default='output.json', + help='Output Lottie JSON file path') + + # Generation arguments + parser.add_argument('--max_tokens', type=int, default=4096, + help='Maximum number of tokens to generate') + parser.add_argument('--do_sample', action='store_true', + help='Enable sampling (otherwise use greedy decoding)') + parser.add_argument('--temperature', type=float, default=0.9, + help='Sampling temperature') + parser.add_argument('--top_p', type=float, default=0.25, + help='Top-p sampling') + parser.add_argument('--top_k', type=int, default=5, + help='Top-k sampling') + + # Device arguments + parser.add_argument('--device', type=str, default='cuda', + help='Device (cuda/cpu)') + + args = parser.parse_args() + + # Validate input + if not (args.text or args.image or args.video): + parser.error("Must provide --text, --image, or --video") + + # Set device + device = torch.device(args.device if torch.cuda.is_available() else "cpu") + + # Load model + print("="*60) + print("Loading OmniLottie model...") + print("="*60) + + print(f"\n1. Loading model from: {args.model_path}") + model = LottieDecoder.from_pretrained( + args.model_path, + torch_dtype=torch.bfloat16, + trust_remote_code=True + ) + model = model.to(device).eval() + print(f" āœ“ Model loaded (vocab_size: {model.vocab_size})") + + print(f"\n2. Loading processor from: {args.processor_path}") + processor = AutoProcessor.from_pretrained( + args.processor_path, + padding_side="left", + trust_remote_code=True + ) + print(f" āœ“ Processor loaded") + + # Prepare inputs + print("\n" + "="*60) + print("Preparing inputs...") + print("="*60) + + if args.text: + print(f"\nMode: Text-to-Lottie") + print(f"Prompt: {args.text}") + messages = build_messages("text", text_prompt=args.text) + + elif args.image: + print(f"\nMode: Image-to-Lottie") + print(f"Image: {args.image}") + image = Image.open(args.image) + if image.mode != 'RGB': + image = image.convert('RGB') + image = image.resize((448, 448), Image.LANCZOS) + messages = build_messages("image", text_prompt=args.text, image=image) + + elif args.video: + print(f"\nMode: Video-to-Lottie") + print(f"Video: {args.video}") + frames = load_frames_from_video(args.video, num_frames=8) + messages = build_messages("video", video_frames=frames) + + # Prepare inference input + inputs = prepare_inference_input(processor, messages, device) + + # Generate + print("\n" + "="*60) + print("Generating Lottie animation...") + print("="*60) + + lottie_tokens = generate_lottie( + model=model, + inputs=inputs, + max_tokens=args.max_tokens, + device=device, + use_sampling=args.do_sample, + temperature=args.temperature, + top_p=args.top_p, + top_k=args.top_k + ) + + print(f"\nāœ“ Generated {len(lottie_tokens)} Lottie tokens") + + # Convert to JSON + print("\nConverting tokens to Lottie JSON...") + lottie_json = tokens_to_lottie_json(lottie_tokens) + + # Save + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + + with open(output_path, 'w', encoding='utf-8') as f: + json.dump(lottie_json, f, indent=2) + + print("\n" + "="*60) + print("āœ“ Generation complete!") + print("="*60) + print(f"\nOutput saved to: {output_path}") + print(f"Animation info:") + print(f" - Size: {lottie_json['w']} x {lottie_json['h']}") + print(f" - Frame rate: {lottie_json['fr']} fps") + print(f" - Duration: {lottie_json['op'] - lottie_json['ip']} frames") + print(f" - Layers: {len(lottie_json.get('layers', []))}") + + print(f"\nšŸ’” You can now use this Lottie file with:") + print(f" - lottie-web: https://airbnb.io/lottie/") + print(f" - LottieFiles: https://lottiefiles.com/") + + +if __name__ == "__main__": + main()