Upload folder using huggingface_hub

README.md

@@ -1,12 +1,11 @@
 ---
-title: Talking Head Lora Train
-emoji: 💻
-colorFrom: red
-colorTo: purple
+title: Talking Head - LoRA Train
+emoji: 🎨
+colorFrom: yellow
+colorTo: red
 sdk: gradio
-sdk_version: 6.9.0
+sdk_version: 5.9.1
 app_file: app.py
 pinned: false
+hardware: a100-large
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,323 @@
+"""Space 4: Train LoRA (Flux.1-dev + PEFT)
+
+Downloads frames from Hub -> LoRA training on Flux.1 -> saves adapter to Hub.
+GPU: A100 (Flux.1-dev full pipeline + LoRA)
+"""
+import gc
+import json
+import logging
+import os
+import shutil
+import traceback
+from pathlib import Path
+
+import gradio as gr
+import torch
+from PIL import Image
+from torchvision import transforms
+
+from hub_utils import download_step, upload_step
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
+logger = logging.getLogger(__name__)
+
+# ── Config ──
+IS_HF_SPACE = os.environ.get("SPACE_ID") is not None
+_data_path = Path("/data")
+if IS_HF_SPACE and _data_path.exists() and os.access(_data_path, os.W_OK):
+    BASE_DIR = _data_path
+else:
+    BASE_DIR = Path("data")
+
+FRAMES_DIR = BASE_DIR / "frames"
+LORA_MODEL_DIR = BASE_DIR / "lora_model"
+TEMP_DIR = BASE_DIR / "temp"
+HF_CACHE_DIR = BASE_DIR / "hf_cache"
+
+for d in [FRAMES_DIR, LORA_MODEL_DIR, TEMP_DIR, HF_CACHE_DIR]:
+    d.mkdir(parents=True, exist_ok=True)
+
+os.environ["HF_HOME"] = str(HF_CACHE_DIR)
+os.environ["TRANSFORMERS_CACHE"] = str(HF_CACHE_DIR)
+
+FLUX_MODEL_ID = "black-forest-labs/FLUX.1-dev"
+LORA_TRIGGER_WORD = "alvaro_person"
+LORA_RANK = 16
+LORA_ALPHA = 16
+LORA_LR = 1e-4
+LORA_STEPS = 1500
+LORA_BATCH_SIZE = 1
+LORA_RESOLUTION = 1024
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+APP_VERSION = "1.0.0"
+
+
+def _clear_cache():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+
+# ── Dataset preparation ──
+
+def _prepare_dataset(image_dir, trigger_word):
+    dataset_dir = TEMP_DIR / "lora_dataset"
+    if dataset_dir.exists():
+        shutil.rmtree(dataset_dir)
+    dataset_dir.mkdir(parents=True)
+
+    images = sorted(image_dir.glob("*.jpg")) + sorted(image_dir.glob("*.png"))
+    captions_file = image_dir / "captions.json"
+    captions = {}
+    if captions_file.exists():
+        with open(captions_file) as f:
+            captions = json.load(f)
+
+    for i, img_path in enumerate(images):
+        dst_img = dataset_dir / f"img_{i:04d}{img_path.suffix}"
+        shutil.copy2(img_path, dst_img)
+        caption = captions.get(img_path.name, "a photo of a person")
+        full_caption = f"{trigger_word}, {caption}"
+        dst_img.with_suffix(".txt").write_text(full_caption)
+
+    logger.info(f"Prepared {len(images)} images with trigger word '{trigger_word}'")
+    return dataset_dir
+
+
+# ── LoRA training ──
+
+def _train_lora(
+    dataset_dir, output_dir, rank, alpha, learning_rate,
+    max_steps, batch_size, resolution, progress_callback=None,
+):
+    from diffusers import FluxPipeline
+    from peft import LoraConfig, get_peft_model
+    from torch.utils.data import Dataset, DataLoader
+
+    class CaptionedImageDataset(Dataset):
+        def __init__(self, root_dir, res):
+            self.root = Path(root_dir)
+            self.images = sorted(self.root.glob("*.jpg")) + sorted(self.root.glob("*.png"))
+            self.transform = transforms.Compose([
+                transforms.Resize((res, res)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5], [0.5]),
+            ])
+        def __len__(self):
+            return len(self.images)
+        def __getitem__(self, idx):
+            img_path = self.images[idx]
+            image = Image.open(img_path).convert("RGB")
+            image = self.transform(image)
+            txt_path = img_path.with_suffix(".txt")
+            caption = txt_path.read_text().strip() if txt_path.exists() else ""
+            return {"image": image, "caption": caption}
+
+    logger.info(f"Loading Flux.1 from {FLUX_MODEL_ID}...")
+    pipe = FluxPipeline.from_pretrained(
+        FLUX_MODEL_ID, torch_dtype=torch.bfloat16,
+        token=os.environ.get("HF_TOKEN"),
+    )
+
+    lora_config = LoraConfig(
+        r=rank, lora_alpha=alpha,
+        target_modules=["to_q", "to_k", "to_v", "to_out.0"],
+        lora_dropout=0.0,
+    )
+
+    pipe.transformer = get_peft_model(pipe.transformer, lora_config)
+    pipe.transformer.to(DEVICE, dtype=torch.bfloat16)
+    pipe.transformer.train()
+
+    trainable_params = sum(p.numel() for p in pipe.transformer.parameters() if p.requires_grad)
+    logger.info(f"Trainable LoRA parameters: {trainable_params:,}")
+
+    dataset = CaptionedImageDataset(dataset_dir, resolution)
+    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0)
+
+    optimizer = torch.optim.AdamW(
+        [p for p in pipe.transformer.parameters() if p.requires_grad],
+        lr=learning_rate, weight_decay=0.01,
+    )
+
+    pipe.text_encoder.to(DEVICE, dtype=torch.bfloat16)
+    if hasattr(pipe, "text_encoder_2") and pipe.text_encoder_2 is not None:
+        pipe.text_encoder_2.to(DEVICE, dtype=torch.bfloat16)
+    pipe.vae.to(DEVICE, dtype=torch.bfloat16)
+
+    global_step = 0
+    for epoch in range(max_steps // max(1, len(dataset)) + 1):
+        for batch in loader:
+            if global_step >= max_steps:
+                break
+            images_batch = batch["image"].to(DEVICE, dtype=torch.bfloat16)
+            captions_batch = batch["caption"]
+
+            with torch.no_grad():
+                latents = pipe.vae.encode(images_batch).latent_dist.sample()
+                latents = (latents - pipe.vae.config.shift_factor) * pipe.vae.config.scaling_factor
+                prompt_embeds, pooled_prompt_embeds, text_ids = pipe.encode_prompt(
+                    prompt=captions_batch, prompt_2=captions_batch,
+                )
+                prompt_embeds = prompt_embeds.to(dtype=torch.bfloat16)
+                pooled_prompt_embeds = pooled_prompt_embeds.to(dtype=torch.bfloat16)
+
+            noise = torch.randn_like(latents)
+            timesteps = torch.randint(0, 1000, (latents.shape[0],), device=latents.device)
+            sigmas = (timesteps.float() / 1000.0).to(dtype=torch.bfloat16).view(-1, 1, 1, 1)
+            noisy_latents = (1 - sigmas) * latents + sigmas * noise
+
+            bs, ch, h, w = noisy_latents.shape
+            noisy_packed = pipe._pack_latents(noisy_latents, bs, ch, h, w)
+            latent_image_ids = pipe._prepare_latent_image_ids(bs, h // 2, w // 2, DEVICE, torch.bfloat16)
+            guidance = torch.full((bs,), 3.5, device=DEVICE, dtype=torch.bfloat16)
+
+            noise_pred = pipe.transformer(
+                hidden_states=noisy_packed, timestep=timesteps, guidance=guidance,
+                encoder_hidden_states=prompt_embeds, pooled_projections=pooled_prompt_embeds,
+                txt_ids=text_ids, img_ids=latent_image_ids, return_dict=False,
+            )[0]
+
+            target = noise - latents
+            target_packed = pipe._pack_latents(target, bs, ch, h, w)
+            loss = torch.nn.functional.mse_loss(noise_pred, target_packed)
+
+            loss.backward()
+            optimizer.step()
+            optimizer.zero_grad()
+
+            global_step += 1
+            if global_step % 50 == 0:
+                logger.info(f"Step {global_step}/{max_steps}, Loss: {loss.item():.4f}")
+                if progress_callback:
+                    prog = 0.1 + (global_step / max_steps) * 0.85
+                    progress_callback(prog, f"Step {global_step}/{max_steps}, Loss: {loss.item():.4f}")
+
+        if global_step >= max_steps:
+            break
+
+    pipe.transformer.save_pretrained(str(output_dir))
+    logger.info(f"LoRA saved to {output_dir}")
+    del pipe
+    _clear_cache()
+
+
+# ── Gradio handlers ──
+
+def download_frames_from_hub(project_name):
+    if not project_name or not project_name.strip():
+        return "Error: Debes introducir un nombre de proyecto"
+    name = project_name.strip()
+    try:
+        if FRAMES_DIR.exists():
+            shutil.rmtree(FRAMES_DIR)
+        FRAMES_DIR.mkdir(parents=True)
+
+        download_step(name, "step1_frames", str(BASE_DIR))
+        src = BASE_DIR / name / "step1_frames"
+        if src.exists():
+            for f in src.iterdir():
+                shutil.move(str(f), str(FRAMES_DIR / f.name))
+            shutil.rmtree(BASE_DIR / name, ignore_errors=True)
+
+        frames = sorted(FRAMES_DIR.glob("*.jpg"))
+        return f"OK - Descargados {len(frames)} frames"
+    except Exception as e:
+        return f"Error: {e}"
+
+
+def train_lora_handler(project_name, trigger_word, rank, lr, steps, progress=gr.Progress()):
+    if not project_name or not project_name.strip():
+        return "Error: Debes introducir un nombre de proyecto"
+
+    images = list(FRAMES_DIR.glob("*.jpg")) + list(FRAMES_DIR.glob("*.png"))
+    if len(images) < 10:
+        return f"Error: Se necesitan al menos 10 imagenes, encontradas {len(images)}. Descarga frames primero."
+
+    logger.info(f"=== LoRA Training Started === trigger={trigger_word}, rank={rank}, steps={steps}")
+    try:
+        _clear_cache()
+        progress(0.05, desc="Preparando dataset...")
+        dataset_dir = _prepare_dataset(FRAMES_DIR, trigger_word)
+
+        progress(0.1, desc="Iniciando entrenamiento LoRA...")
+        _train_lora(
+            dataset_dir=dataset_dir, output_dir=LORA_MODEL_DIR,
+            rank=int(rank), alpha=int(rank), learning_rate=lr,
+            max_steps=int(steps), batch_size=LORA_BATCH_SIZE,
+            resolution=LORA_RESOLUTION,
+            progress_callback=lambda p, m: progress(p, desc=m),
+        )
+
+        config = {
+            "base_model": FLUX_MODEL_ID, "trigger_word": trigger_word,
+            "rank": int(rank), "alpha": int(rank), "steps": int(steps),
+            "resolution": LORA_RESOLUTION,
+        }
+        with open(LORA_MODEL_DIR / "lora_config.json", "w") as f:
+            json.dump(config, f, indent=2)
+
+        shutil.rmtree(dataset_dir, ignore_errors=True)
+        _clear_cache()
+
+        logger.info("=== LoRA Training Complete ===")
+        return f"OK - LoRA guardado en: {LORA_MODEL_DIR}"
+    except Exception as e:
+        logger.error(f"=== LoRA Training Failed ===\n{traceback.format_exc()}")
+        return f"Error: {e}"
+
+
+def save_to_hub(project_name):
+    if not project_name or not project_name.strip():
+        return "Error: Debes introducir un nombre de proyecto"
+    name = project_name.strip()
+    models = list(LORA_MODEL_DIR.glob("*.safetensors")) + list(LORA_MODEL_DIR.glob("adapter_model.*"))
+    if not models:
+        return "Error: No hay modelo LoRA para guardar. Entrena primero."
+    try:
+        return upload_step(name, "step4_lora", str(LORA_MODEL_DIR))
+    except Exception as e:
+        return f"Error: {e}"
+
+
+# ── UI ──
+
+with gr.Blocks(title="Talking Head - LoRA Train", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(f"# Talking Head - Entrenar LoRA `v{APP_VERSION}`\nFlux.1-dev LoRA training con tus imagenes")
+
+    project_name = gr.Textbox(
+        label="Nombre del proyecto",
+        placeholder="mi_proyecto",
+        info="Obligatorio. Se usa como carpeta en el Hub.",
+    )
+
+    gr.Markdown("### 1. Descargar frames del Hub")
+    download_btn = gr.Button("Descargar frames del Hub", variant="secondary")
+    download_status = gr.Textbox(label="Estado descarga", interactive=False)
+
+    gr.Markdown("### 2. Entrenar LoRA")
+    with gr.Row():
+        trigger_word = gr.Textbox(value=LORA_TRIGGER_WORD, label="Trigger Word")
+        lora_rank = gr.Slider(4, 64, value=LORA_RANK, step=4, label="LoRA Rank")
+    with gr.Row():
+        lora_lr = gr.Number(value=LORA_LR, label="Learning Rate")
+        lora_steps = gr.Slider(500, 5000, value=LORA_STEPS, step=100, label="Training Steps")
+    train_btn = gr.Button("Entrenar LoRA", variant="primary")
+    train_status = gr.Textbox(label="Estado entrenamiento", interactive=False)
+
+    gr.Markdown("### 3. Guardar modelo en Hub")
+    save_btn = gr.Button("Guardar en Hub", variant="secondary")
+    save_status = gr.Textbox(label="Estado guardado", interactive=False)
+
+    download_btn.click(download_frames_from_hub, inputs=[project_name], outputs=[download_status])
+    train_btn.click(
+        train_lora_handler,
+        inputs=[project_name, trigger_word, lora_rank, lora_lr, lora_steps],
+        outputs=[train_status],
+    )
+    save_btn.click(save_to_hub, inputs=[project_name], outputs=[save_status])
+
+if __name__ == "__main__":
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860)

hub_utils.py

@@ -0,0 +1,64 @@
+"""Hub utilities for uploading/downloading step data to HF Dataset repo."""
+import os
+import logging
+from pathlib import Path
+from huggingface_hub import HfApi, hf_hub_download, list_repo_tree
+
+logger = logging.getLogger(__name__)
+
+HF_DATASET_REPO_ID = "baenacoco/talking-head-avatar"
+
+
+def _get_api():
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        raise ValueError("HF_TOKEN no encontrado en variables de entorno")
+    api = HfApi(token=token)
+    api.create_repo(repo_id=HF_DATASET_REPO_ID, repo_type="dataset", exist_ok=True)
+    return api
+
+
+def upload_step(name: str, step_folder: str, local_dir: str):
+    """Upload a local directory to {name}/{step_folder}/ in the dataset repo."""
+    api = _get_api()
+    api.upload_folder(
+        folder_path=local_dir,
+        path_in_repo=f"{name}/{step_folder}",
+        repo_id=HF_DATASET_REPO_ID,
+        repo_type="dataset",
+    )
+    logger.info(f"Uploaded {local_dir} -> {name}/{step_folder}")
+    return f"Subido a Hub: {name}/{step_folder}"
+
+
+def download_step(name: str, step_folder: str, local_dir: str):
+    """Download {name}/{step_folder}/ from the dataset repo to a local directory."""
+    from huggingface_hub import snapshot_download
+    token = os.environ.get("HF_TOKEN")
+    snapshot_download(
+        repo_id=HF_DATASET_REPO_ID,
+        repo_type="dataset",
+        local_dir=local_dir,
+        allow_patterns=[f"{name}/{step_folder}/**"],
+        token=token,
+    )
+    logger.info(f"Downloaded {name}/{step_folder} -> {local_dir}")
+    return f"Descargado de Hub: {name}/{step_folder}"
+
+
+def list_projects() -> list[str]:
+    """List project names (top-level folders) in the dataset repo."""
+    token = os.environ.get("HF_TOKEN")
+    try:
+        api = HfApi(token=token)
+        entries = list(api.list_repo_tree(
+            repo_id=HF_DATASET_REPO_ID, repo_type="dataset", path_in_repo="",
+        ))
+        return sorted(set(
+            e.rfilename.split("/")[0] if hasattr(e, "rfilename") else e.path.split("/")[0]
+            for e in entries
+            if ("/" in getattr(e, "rfilename", "")) or hasattr(e, "path")
+        ))
+    except Exception as e:
+        logger.warning(f"Could not list projects: {e}")
+        return []

packages.txt

@@ -0,0 +1,2 @@
+libgl1-mesa-glx
+libglib2.0-0

requirements.txt

@@ -0,0 +1,13 @@
+setuptools>=69.0.0
+gradio>=5.9.1
+torch>=2.1.0
+torchvision>=0.16.0
+transformers>=4.36.0,<5.0.0
+diffusers>=0.25.0
+accelerate>=0.25.0
+safetensors>=0.4.0
+peft>=0.7.0
+huggingface_hub>=0.20.0
+Pillow>=10.0.0
+sentencepiece>=0.1.99
+protobuf>=3.20.0