# --- Lens en offload complet pour GPU 16 Go ---

!git clone https://github.com/microsoft/Lens.git
!pip install diffusers transformers accelerate sentencepiece

import sys
sys.path.append("/content/Lens")

import torch
from lens import LensPipeline

pipe = LensPipeline.from_pretrained(
    "microsoft/Lens",
    torch_dtype=torch.bfloat16,
    device_map="cpu",               # <-- clé pour 16 Go
    offload_folder="offload",         # <-- swap CPU
    low_cpu_mem_usage=True
)

pipe.enable_model_cpu_offload()       # <-- plus fort que sequential
pipe.enable_vae_slicing()             # <-- réduit la VRAM
pipe.enable_vae_tiling()              # <-- réduit encore la VRAM

image = pipe(
    prompt="A cat holding a sign that says 'hello world'",
    base_resolution=1024,             # <-- 1440 trop lourd pour 16 Go
    aspect_ratio="1:1",
    num_inference_steps=20,
    guidance_scale=5.0,
    generator=torch.Generator(device="cuda").manual_seed(0),
).images[0]

image.save("lens.png")
print("Image générée : lens.png")