Views
No views yet
uv add transformers torch accelerate1from transformers import AutoModel, AutoTokenizer
2from transformers import SiglipProcessor
3
4# Load the model
5model = AutoModel.from_pretrained("OSTswiss/LiteViT5", trust_remote_code=True)
6
7# Load tokenizer and processor
8tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-base")
9processor = SiglipProcessor.from_pretrained("google/siglip2-base-patch16-512")1from PIL import Image
2import torch
3
4from transformers import AutoModel, AutoTokenizer
5from transformers import SiglipProcessor
6
7# Load the model
8model = AutoModel.from_pretrained("OSTswiss/LiteViT5", trust_remote_code=True, device_map="auto")
9
10# Load tokenizer and processor
11tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-base")
12processor = SiglipProcessor.from_pretrained("google/siglip2-base-patch16-512")
13
14# Preprocess image (split into 4 parts + full image = 5 views)
15def prepare_image(image_path: str, processor):
16 """
17 Prepare image with 5 views (4 quarters + full).
18
19 Args:
20 image_path: Path to the image file
21 processor: SigLIP processor
22
23 Returns:
24 Tensor of shape [5, 3, 512, 512]
25 """
26 image = Image.open(image_path).convert("RGB")
27
28 # Split into 4 quarters
29 width, height = image.size
30 quarters = [
31 image.crop((0, 0, width//2, height//2)), # top-left
32 image.crop((width//2, 0, width, height//2)), # top-right
33 image.crop((0, height//2, width//2, height)), # bottom-left
34 image.crop((width//2, height//2, width, height)), # bottom-right
35 ]
36
37 # Process all views
38 processed = [
39 processor(images=q, return_tensors="pt")["pixel_values"]
40 for q in quarters
41 ]
42 # Add full image
43 processed.append(
44 processor(images=image, return_tensors="pt")["pixel_values"]
45 )
46
47 pixel_values = torch.cat(processed, dim=0)
48 return pixel_values
49
50def generate_text(model, pixel_values, tokenizer, max_length=512):
51 """
52 Generate text from image.
53
54 Args:
55 model: LiteVit5 model
56 pixel_values: Preprocessed image tensor
57 tokenizer: Tokenizer for decoding
58 max_length: Maximum generation length
59
60 Returns:
61 Generated text string
62 """
63 with torch.no_grad():
64 output_ids = model.generate(pixel_values, max_length=max_length)
65
66 text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
67 return text
68
69device = next(model.parameters()).device
70
71# Process images
72pixel_values = prepare_image("./image_13.png", processor)
73pixel_values = pixel_values.to(device)
74print("\nGenerating HTML from image_13.png...")
75output_ids = model.generate(pixel_values, max_length=2024)
76text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
77print(f"Generated: {text}")