Upload extract_sdxl_embeddings.py

extract_sdxl_embeddings.py

@@ -0,0 +1,100 @@
+# extract_sdxl_embeddings.py
+import argparse
+from pathlib import Path
+from typing import List
+
+import torch
+from safetensors.torch import save_file
+from diffusers import StableDiffusionXLPipeline
+
+
+def read_prompts(txt_path: str) -> List[str]:
+    with open(txt_path, "r", encoding="utf-8") as f:
+        return [line.rstrip("\n") for line in f]
+
+
+def load_sdxl(checkpoint_path: str, precision: str):
+    precision = precision.lower()
+    if precision == "bf16":
+        dtype = torch.bfloat16
+    else:
+        dtype = torch.float16  # T4 suele ir mejor así para SDXL
+
+    path = Path(checkpoint_path)
+    if path.is_dir():
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            checkpoint_path,
+            torch_dtype=dtype,
+            use_safetensors=True,
+        )
+    else:
+        # Útil para .safetensors / .ckpt de un solo archivo.
+        pipe = StableDiffusionXLPipeline.from_single_file(
+            checkpoint_path,
+            torch_dtype=dtype,
+        )
+
+    pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+    pipe.set_progress_bar_config(disable=True)
+    pipe.eval()
+    return pipe
+
+
+@torch.no_grad()
+def encode_batch(pipe: StableDiffusionXLPipeline, batch_prompts: List[str]):
+    device = pipe._execution_device if hasattr(pipe, "_execution_device") else next(pipe.text_encoder.parameters()).device
+
+    # Diffusers soporta prompt_embeds y pooled_prompt_embeds en SDXL. :contentReference[oaicite:3]{index=3}
+    prompt_embeds, pooled_prompt_embeds = pipe.encode_prompt(
+        prompt=batch_prompts,
+        prompt_2=batch_prompts,
+        device=device,
+        num_images_per_prompt=1,
+        do_classifier_free_guidance=False,
+    )[:2]
+
+    return prompt_embeds.detach().cpu(), pooled_prompt_embeds.detach().cpu()
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--sdxl_checkpoint", type=str, required=True,
+                        help="Ruta al .safetensors / .ckpt o directorio Diffusers de SDXL.")
+    parser.add_argument("--prompts_txt", type=str, required=True)
+    parser.add_argument("--out_dir", type=str, default="output_embeddings")
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--precision", type=str, default="fp16", choices=["fp16", "bf16"])
+    parser.add_argument("--pad_width", type=int, default=5)
+    args = parser.parse_args()
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    prompts = read_prompts(args.prompts_txt)
+    pipe = load_sdxl(args.sdxl_checkpoint, args.precision)
+
+    n = len(prompts)
+    print(f"Procesando {n} prompts...")
+
+    for i in range(0, n, args.batch_size):
+        batch_prompts = prompts[i:i + args.batch_size]
+        text_embeds, pooled_text_embeds = encode_batch(pipe, batch_prompts)
+
+        for b in range(text_embeds.shape[0]):
+            file_idx = i + b
+            file_name = f"{file_idx:0{args.pad_width}d}.safetensors"
+            save_path = out_dir / file_name
+
+            sample = {
+                "text_embeds": text_embeds[b:b+1].contiguous().to(torch.float16),
+                "pooled_text_embeds": pooled_text_embeds[b:b+1].contiguous().to(torch.float16),
+            }
+            save_file(sample, str(save_path))
+
+        print(f"Guardado hasta {min(i + args.batch_size - 1, n - 1):0{args.pad_width}d}")
+
+    print(f"Listo. Salida en: {out_dir}")
+
+
+if __name__ == "__main__":
+    main()