app.py

# universal_lora_trainer_accelerate_singlefile_dynamic.py
"""
Universal Dynamic LoRA Trainer (Accelerate + PEFT + Gradio)
- Auto-detects base model type (Flux, SD, ChronoEdit, QwenEdit, etc.)
- Auto-selects correct adapter target (unet, transformer, text_encoder)
- Supports CSV and Parquet datasets
- Uploads adapter to HF Hub using HF_TOKEN (env only)
"""

import os, torch, gradio as gr, pandas as pd, numpy as np
from pathlib import Path
from tqdm.auto import tqdm
from huggingface_hub import create_repo, upload_folder, hf_hub_download
from diffusers import DiffusionPipeline
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as T, torchvision
from peft import LoraConfig, get_peft_model
from accelerate import Accelerator
import torch.nn as nn

# Optional: ChronoEdit + QwenEdit
try:
    from chronoedit_diffusers.pipeline_chronoedit import ChronoEditPipeline
    CHRONOEDIT_AVAILABLE = True
except Exception:
    CHRONOEDIT_AVAILABLE = False

try:
    from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPipeline
    QWENEDIT_AVAILABLE = True
except Exception:
    QWENEDIT_AVAILABLE = False

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp"}
VIDEO_EXTS = {".mp4", ".mov", ".avi", ".mkv"}

def is_hub_repo_like(s): return "/" in s and not Path(s).exists()

def download_from_hf(repo_id, filename, token=None):
    token = token or os.environ.get("HF_TOKEN")
    return hf_hub_download(repo_id=repo_id, filename=filename, repo_type="dataset", token=token)

# ---------------- Dataset ----------------
class MediaTextDataset(Dataset):
    def __init__(self, source, csv_name="dataset.csv", max_frames=5):
        self.is_hub = is_hub_repo_like(source)
        self.source = source
        token = os.environ.get("HF_TOKEN")
        if self.is_hub:
            file_path = download_from_hf(source, csv_name, token)
        else:
            file_path = Path(source) / csv_name
        if not Path(file_path).exists():
            alt = Path(str(file_path).replace(".csv", ".parquet"))
            if alt.exists(): file_path = alt
        self.df = pd.read_parquet(file_path) if str(file_path).endswith(".parquet") else pd.read_csv(file_path)
        self.root = Path(source) if not self.is_hub else None
        self.img_tf = T.Compose([T.ToPILImage(), T.Resize((512,512)), T.ToTensor(), T.Normalize([0.5]*3, [0.5]*3)])
        self.video_tf = T.Compose([T.ToPILImage(), T.Resize((128,256)), T.ToTensor(), T.Normalize([0.5]*3, [0.5]*3)])
        self.max_frames = max_frames
    def __len__(self): return len(self.df)
    def _maybe_dl(self, fname): return str(Path(self.root)/fname) if self.root else download_from_hf(self.source, fname)
    def __getitem__(self, i):
        rec = self.df.iloc[i]
        p = Path(self._maybe_dl(rec["file_name"]))
        if p.suffix.lower() in IMAGE_EXTS:
            img = torchvision.io.read_image(str(p))
            if isinstance(img, torch.Tensor): img = img.permute(1,2,0).numpy()
            return {"type": "image", "image": self.img_tf(img), "caption": rec["text"]}
        elif p.suffix.lower() in VIDEO_EXTS:
            vid,_,_ = torchvision.io.read_video(str(p))
            total, idxs = len(vid), []
            if total == 0: return {"type":"video","frames":torch.zeros((self.max_frames,3,128,256))}
            if total < self.max_frames: idxs = list(range(total))+[total-1]*(self.max_frames-total)
            else: idxs = np.linspace(0,total-1,self.max_frames).round().astype(int)
            frames = torch.stack([self.video_tf(vid[j].numpy()) for j in idxs])
            return {"type": "video", "frames": frames, "caption": rec["text"]}
        else: raise RuntimeError(f"Unsupported file {p}")

# ---------------- Dynamic pipeline loader ----------------
def load_pipeline_auto(base_model, dtype=torch.float16):
    low = base_model.lower()
    if "chrono" in low and CHRONOEDIT_AVAILABLE:
        print(f"Using ChronoEdit pipeline for {base_model}")
        return ChronoEditPipeline.from_pretrained(base_model, torch_dtype=dtype)
    elif "qwen" in low and QWENEDIT_AVAILABLE:
        print(f"Using QwenEdit pipeline for {base_model}")
        return QwenImageEditPipeline.from_pretrained(base_model, torch_dtype=dtype)
    else:
        print(f"Using Diffusion pipeline for {base_model}")
        return DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype)

def infer_target_for_task(task_type, model_name):
    if task_type == "prompt-lora" or "qwen" in model_name.lower():
        return "text_encoder"
    elif task_type == "text-video" or "chrono" in model_name.lower() or "wan" in model_name.lower():
        return "transformer"
    else:
        return "unet"

def find_target_modules(model):
    names = [n for n,_ in model.named_modules()]
    targets = [n.split(".")[-1] for n in names if any(k in n for k in ["to_q","to_k","to_v","q_proj","v_proj"])]
    return targets or ["to_q","to_k","to_v","to_out"]

# ---------------- Training ----------------
def train_lora(base_model, dataset_src, csv_name, task_type, output_dir, epochs=1, lr=1e-4, r=8, alpha=16):
    accelerator = Accelerator()
    pipe = load_pipeline_auto(base_model)
    target = infer_target_for_task(task_type, base_model)
    if not hasattr(pipe, target): raise RuntimeError(f"Pipeline has no {target}")
    lcfg = LoraConfig(r=r, lora_alpha=alpha, target_modules=find_target_modules(getattr(pipe, target)), lora_dropout=0.0)
    lora_module = get_peft_model(getattr(pipe, target), lcfg)
    dataset = MediaTextDataset(dataset_src, csv_name)
    loader = DataLoader(dataset, batch_size=1, shuffle=True)
    lora_module, opt, loader = accelerator.prepare(lora_module, torch.optim.AdamW(lora_module.parameters(), lr=lr), loader)
    mse = nn.MSELoss(); logs=[]
    for ep in range(epochs):
        for i,b in enumerate(tqdm(loader, desc=f"Epoch {ep+1}")):
            ex = b[0]; loss=torch.tensor(0.0, device=DEVICE)
            if ex["type"]=="image" and hasattr(pipe,"vae"):
                img=ex["image"].unsqueeze(0).to(DEVICE)
                lat=pipe.vae.encode(img).latent_dist.sample()*pipe.vae.config.scaling_factor
                noise=torch.randn_like(lat); loss=mse(lat,noise)
            accelerator.backward(loss); opt.step(); opt.zero_grad()
            logs.append(f"step {i} loss {loss.item():.4f}")
    Path(output_dir).mkdir(exist_ok=True)
    lora_module.save_pretrained(output_dir)
    return output_dir, logs[-20:]

# ---------------- Upload ----------------
def upload_adapter(local, repo_id):
    token=os.environ.get("HF_TOKEN")
    if not token: raise RuntimeError("HF_TOKEN missing")
    create_repo(repo_id, exist_ok=True)
    upload_folder(local, repo_id=repo_id, repo_type="model", token=token)
    return f"https://huggingface.co/{repo_id}"

# ---------------- Gradio UI ----------------
def run_ui():
    with gr.Blocks() as demo:
        gr.Markdown("# 🌐 Universal Dynamic LoRA Trainer (Flux / ChronoEdit / QwenEdit)")
        with gr.Row():
            base_model=gr.Textbox(label="Base model", value="black-forest-labs/FLUX.1-dev")
            dataset=gr.Textbox(label="Dataset folder or HF repo", value="./dataset")
            csvname=gr.Textbox(label="CSV/Parquet file", value="dataset.csv")
            task=gr.Dropdown(["text-image","text-video","prompt-lora"], label="Task type", value="text-image")
            out=gr.Textbox(label="Output dir", value="./adapter_out")
            repo=gr.Textbox(label="Upload HF repo (optional)", value="")
        with gr.Row():
            r=gr.Slider(1,64,value=8,label="LoRA rank"); a=gr.Slider(1,64,value=16,label="LoRA alpha")
            ep=gr.Number(value=1,label="Epochs"); lr=gr.Number(value=1e-4,label="Learning rate")
            btn=gr.Button("🚀 Start Training")
        logs=gr.Textbox(label="Logs", lines=12)
        img=gr.Image(label="Sample Output (optional)")
        def launch(bm,ds,csv,t,out_dir,r_,a_,ep_,lr_,repo_):
            try:
                out,log=train_lora(bm,ds,csv,t,out_dir,int(ep_),float(lr_),int(r_),int(a_))
                link=upload_adapter(out,repo_) if repo_ else None
                return "\n".join(log), None, link
            except Exception as e:
                return f"❌ {e}", None, None
        btn.click(launch,[base_model,dataset,csvname,task,out,r,a,ep,lr,repo],[logs,img,gr.Textbox()])
    return demo

if __name__=="__main__":
    run_ui().launch(server_name="0.0.0.0",server_port=7860)