Update app.py

app.py

@@ -1,507 +1,168 @@
-# universal_lora_trainer_accelerate_singlefile.py
+# universal_lora_trainer_accelerate_singlefile_dynamic.py
 """
-Universal LoRA Trainer (Accelerate + PEFT) single-file app with Gradio UI.
-
-Features:
- - Supports CSV and Parquet dataset files (columns: file_name, text)
- - Accepts dataset from a local folder or Hugging Face repo id (username/repo)
- - Real LoRA training (PEFT) for: text->image (UNet), text->video (ChronoEdit transformer),
-   and prompt-enhancer LoRA (QwenEdit/text_encoder)
- - Uses accelerate for device orchestration (recommended: use `accelerate launch ...` for multi-GPU)
- - Shows logs and sample generation in Gradio
- - Uploads adapter to HF Hub using HF_TOKEN from environment (not UI)
-
-Requirements:
- pip install torch torchvision diffusers transformers accelerate peft huggingface_hub gradio pandas tqdm
-
-Optional (ChronoEdit speedups): pip install chronoedit-diffusers flash-attn
+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
-import tempfile
+import os, torch, gradio as gr, pandas as pd, numpy as np
 from pathlib import Path
-from typing import Optional, Tuple, List
-
-import torch
-import torch.nn as nn
-from torch.utils.data import Dataset, DataLoader
-import torchvision
-import torchvision.transforms as T
-import pandas as pd
-import numpy as np
-import gradio as gr
 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
+# Optional: ChronoEdit + QwenEdit
 try:
     from chronoedit_diffusers.pipeline_chronoedit import ChronoEditPipeline
     CHRONOEDIT_AVAILABLE = True
 except Exception:
     CHRONOEDIT_AVAILABLE = False
 
-# PEFT + Accelerate
-try:
-    from peft import LoraConfig, get_peft_model
-except Exception as e:
-    raise RuntimeError("Install peft (pip install peft)") from e
-
 try:
-    from accelerate import Accelerator
-except Exception as e:
-    raise RuntimeError("Install accelerate (pip install accelerate)") from e
+    from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPipeline
+    QWENEDIT_AVAILABLE = True
+except Exception:
+    QWENEDIT_AVAILABLE = False
 
-# ---------- config ----------
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
+IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp"}
 VIDEO_EXTS = {".mp4", ".mov", ".avi", ".mkv"}
-# ---------------------------
 
-def is_hub_repo_like(s: str) -> bool:
-    return "/" in s and not Path(s).exists()
+def is_hub_repo_like(s): return "/" in s and not Path(s).exists()
 
-def download_from_hf(repo_id: str, filename: str, token: Optional[str] = None) -> str:
+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")
-   
+    return hf_hub_download(repo_id=repo_id, filename=filename, repo_type="dataset", token=token)
 
-
-def find_target_modules(model, candidates=("q_proj", "k_proj", "v_proj", "o_proj", "to_q", "to_k", "to_v", "proj_out", "to_out")):
-    names = [n for n, _ in model.named_modules()]
-    selected = set()
-    for cand in candidates:
-        for n in names:
-            if cand in n:
-                selected.add(n.split(".")[-1])
-    if not selected:
-        return ["to_q", "to_k", "to_v", "to_out"]
-    return list(selected)
-
-# -------------------------
-# Dataset: CSV or Parquet
-# -------------------------
+# ---------------- Dataset ----------------
 class MediaTextDataset(Dataset):
-    """
-    Loads records from CSV or parquet with columns:
-      - file_name
-      - text
-    file_name can be a local path relative to dataset_dir, or a filename when using HF repo.
-    """
-    def __init__(self, dataset_source: str, csv_name: str = "dataset.csv", max_frames: int = 5,
-                 image_size=(512,512), video_frame_size=(128,256), hub_token: Optional[str] = None):
-        self.source = dataset_source
-        self.is_hub = is_hub_repo_like(dataset_source)
-        self.max_frames = max_frames
-        self.image_size = image_size
-        self.video_frame_size = video_frame_size
-        self.hub_token = hub_token or os.environ.get("HF_TOKEN")
-        self.tmpdir = None
-
-        # load df (csv or parquet)
+    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:
-            # download file from hub to local cache (hf_hub_download returns cached path)
-            csv_local = download_from_hf(self.source, csv_name, token=self.hub_token)
-            # load via pandas (auto-detect extension)
-            if csv_local.endswith(".parquet"):
-                df = pd.read_parquet(csv_local)
-            else:
-                df = pd.read_csv(csv_local)
-            self.df = df
-            self.root = None
+            file_path = download_from_hf(source, csv_name, token)
         else:
-            root = Path(dataset_source)
-            csv_path_csv = root / csv_name
-            csv_path_parquet = root / csv_name.replace(".csv", ".parquet") if csv_name.endswith(".csv") else root / (csv_name + ".parquet")
-            if csv_path_csv.exists():
-                self.df = pd.read_csv(csv_path_csv)
-            elif csv_path_parquet.exists():
-                self.df = pd.read_parquet(csv_path_parquet)
-            else:
-                # try given csv_name as parquet/csv
-                p = root / csv_name
-                if p.exists():
-                    if p.suffix.lower() == ".parquet":
-                        self.df = pd.read_parquet(p)
-                    else:
-                        self.df = pd.read_csv(p)
-                else:
-                    raise FileNotFoundError(f"Can't find {csv_name} in {dataset_source}")
-            self.root = root
-
-        # transforms
-        self.image_transform = T.Compose([T.ToPILImage(), T.Resize(image_size), T.ToTensor(), T.Normalize([0.5]*3, [0.5]*3)])
-        self.video_transform = T.Compose([T.ToPILImage(), T.Resize(video_frame_size), T.ToTensor(), T.Normalize([0.5]*3, [0.5]*3)])
-
-    def __len__(self):
-        return len(self.df)
-
-    def _maybe_download_from_hub(self, file_name: str) -> str:
-        if self.root is not None:
-            p = self.root / file_name
-            if p.exists():
-                return str(p)
-        # else download from hub
-        return download_from_hf(self.source, file_name, token=self.hub_token)
-
-    def _read_video_frames(self, path: str, num_frames: int):
-        video_frames, _, _ = torchvision.io.read_video(str(path), pts_unit='sec')
-        total = len(video_frames)
-        if total == 0:
-            C, H, W = 3, self.video_frame_size[0], self.video_frame_size[1]
-            return torch.zeros((num_frames, C, H, W), dtype=torch.float32)
-        if total < num_frames:
-            idxs = list(range(total)) + [total-1]*(num_frames-total)
-        else:
-            idxs = np.linspace(0, total-1, num_frames).round().astype(int).tolist()
-        frames = []
-        for i in idxs:
-            arr = video_frames[i].numpy() if hasattr(video_frames[i], "numpy") else np.array(video_frames[i])
-            frames.append(self.video_transform(arr))
-        frames = torch.stack(frames, dim=0)  # [T, C, H, W]
-        return frames
-
-    def __getitem__(self, idx):
-        rec = self.df.iloc[idx]
-        file_name = rec["file_name"]
-        caption = rec["text"]
-        if self.is_hub:
-            local_path = self._maybe_download_from_hub(file_name)
-        else:
-            local_path = str(Path(self.root) / file_name)
-        p = Path(local_path)
-        suffix = p.suffix.lower()
-        if suffix in IMAGE_EXTS:
-            img = torchvision.io.read_image(local_path)  # [C,H,W]
-            if isinstance(img, torch.Tensor):
-                img = img.permute(1,2,0).numpy()
-            return {"type": "image", "image": self.image_transform(img), "caption": caption, "file_name": file_name}
-        elif suffix in VIDEO_EXTS:
-            frames = self._read_video_frames(local_path, self.max_frames)  # [T,C,H,W]
-            return {"type": "video", "frames": frames, "caption": caption, "file_name": file_name}
-        else:
-            raise RuntimeError(f"Unsupported media type: {local_path}")
-
-# -------------------------
-# Pipeline / LoRA helpers
-# -------------------------
-def load_pipeline_auto(base_model_id: str, torch_dtype=torch.float16):
-    is_chrono = "chrono" in base_model_id.lower() or "chronoedit" in base_model_id.lower()
-    if CHRONOEDIT_AVAILABLE and is_chrono:
-        print(f"Loading ChronoEdit pipeline: {base_model_id}")
-        pipe = ChronoEditPipeline.from_pretrained(base_model_id, torch_dtype=torch_dtype)
+            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"Loading standard Diffusers pipeline: {base_model_id}")
-        pipe = DiffusionPipeline.from_pretrained(base_model_id, torch_dtype=torch_dtype)
-    return pipe
-
-def attach_lora(pipe, target: str, r: int = 8, alpha: int = 16, dropout: float = 0.0):
-    if target == "unet":
-        if not hasattr(pipe, "unet"):
-            raise RuntimeError("Pipeline has no UNet for this model")
-        target_module = pipe.unet
-        attr = "unet"
-    elif target == "transformer":
-        if not hasattr(pipe, "transformer"):
-            raise RuntimeError("Pipeline has no transformer for this model")
-        target_module = pipe.transformer
-        attr = "transformer"
-    elif target == "text_encoder":
-        if not hasattr(pipe, "text_encoder"):
-            raise RuntimeError("Pipeline has no text_encoder for this model")
-        target_module = pipe.text_encoder
-        attr = "text_encoder"
+        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:
-        raise RuntimeError("Unknown adapter target")
+        return "unet"
 
-    target_modules = find_target_modules(target_module)
-    print("LoRA target sub-module names detected:", target_modules)
-    lora_config = LoraConfig(r=r, lora_alpha=alpha, target_modules=target_modules, lora_dropout=dropout, bias="none", task_type="SEQ_2_SEQ_LM")
-    peft_model = get_peft_model(target_module, lora_config)
-    setattr(pipe, attr, peft_model)
-    return pipe, attr
+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 loop (Accelerate)
-# -------------------------
-def train_lora_accelerate(base_model_id: str,
-                          dataset_source: str,
-                          csv_name: str,
-                          adapter_target: str,
-                          output_dir: str,
-                          epochs: int = 1,
-                          batch_size: int = 1,
-                          lr: float = 1e-4,
-                          max_train_steps: Optional[int] = None,
-                          lora_r: int = 8,
-                          lora_alpha: int = 16,
-                          max_frames: int = 5,
-                          save_every_steps: int = 200) -> Tuple[str, List[str]]:
+# ---------------- 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()
-    device = accelerator.device
-
-    # load pipeline
-    pipe = load_pipeline_auto(base_model_id, torch_dtype=torch.float16 if device.type == "cuda" else torch.float32)
-
-    dataset = MediaTextDataset(dataset_source, csv_name=csv_name, max_frames=max_frames)
-    dataloader = DataLoader(dataset, batch_size=1, shuffle=True, collate_fn=lambda x: x)
-
-    # attach LoRA
-    pipe, _ = attach_lora(pipe, adapter_target, r=lora_r, alpha=lora_alpha)
-    # select peft module for optimizer
-    if adapter_target == "unet":
-        peft_module = pipe.unet
-    elif adapter_target == "transformer":
-        peft_module = pipe.transformer
-    else:
-        peft_module = pipe.text_encoder
-
-    trainable_params = [p for _, p in peft_module.named_parameters() if p.requires_grad]
-    optimizer = torch.optim.AdamW(trainable_params, lr=lr)
-
-    # Prepare with accelerator
-    peft_module, optimizer, dataloader = accelerator.prepare(peft_module, optimizer, dataloader)
-
-    # If the pipeline has other parts required during training (vae, scheduler...), we'll call them on CPU/GPU directly.
-    logs = []
-    global_step = 0
-    loss_fn = nn.MSELoss()
-
-    if hasattr(pipe, "scheduler"):
-        pipe.scheduler.set_timesteps(50, device=device)
-        timesteps = pipe.scheduler.timesteps
-    else:
-        timesteps = None
-
-    for epoch in range(epochs):
-        pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{epochs}")
-        for batch in pbar:
-            ex = batch[0]
-            if ex["type"] == "image":
-                # image training flow (SD-like)
-                img = ex["image"].unsqueeze(0).to(device)
-                caption = [ex["caption"]]
-
-                if not hasattr(pipe, "encode_prompt"):
-                    raise RuntimeError("Pipeline lacks encode_prompt (can't encode text prompts)")
-
-                prompt_embeds, negative_prompt_embeds = pipe.encode_prompt(prompt=caption, negative_prompt=None, do_classifier_free_guidance=True, num_videos_per_prompt=1, prompt_embeds=None, negative_prompt_embeds=None, max_sequence_length=512, device=device)
-
-                # VAE encode
-                if not hasattr(pipe, "vae"):
-                    raise RuntimeError("Pipeline lacks VAE required for latent conversion")
-                with torch.no_grad():
-                    latents = pipe.vae.encode(img.to(device)).latent_dist.sample() * pipe.vae.config.scaling_factor
-
-                noise = torch.randn_like(latents).to(device)
-                t = pipe.scheduler.timesteps[torch.randint(0, len(pipe.scheduler.timesteps), (1,)).item()].to(device)
-                noisy_latents = pipe.scheduler.add_noise(latents, noise, t)
-
-                # call peft_module (unet) - adapt to common return types
-                # peft_module was prepared by accelerator and is on device
-                unet_out = peft_module(noisy_latents, t.expand(noisy_latents.shape[0]), encoder_hidden_states=prompt_embeds)
-                # unet can return ModelOutput with .sample or tuple
-                if hasattr(unet_out, "sample"):
-                    noise_pred = unet_out.sample
-                elif isinstance(unet_out, tuple):
-                    noise_pred = unet_out[0]
-                else:
-                    # Try to find tensor in object
-                    noise_pred = unet_out
-
-                loss = loss_fn(noise_pred, noise)
-
-            else:
-                # video training (ChronoEdit simplified)
-                if not CHRONOEDIT_AVAILABLE:
-                    raise RuntimeError("ChronoEdit training requested but chronoedit_diffusers not installed")
-                frames = ex["frames"].unsqueeze(0).to(device)  # [1, T, C, H, W]
-                frames_np = frames.squeeze(0).permute(0,2,3,1).cpu().numpy().tolist()
-                video_tensor = pipe.video_processor.preprocess(frames_np, height=frames.shape[-2], width=frames.shape[-1]).to(device)
-                latents_out = pipe.prepare_latents(video_tensor, batch_size=1, num_channels_latents=pipe.vae.config.z_dim, height=video_tensor.shape[-2], width=video_tensor.shape[-1], num_frames=frames.shape[1], dtype=video_tensor.dtype, device=device, generator=None, latents=None, last_image=None)
-                if pipe.config.expand_timesteps:
-                    latents, condition, first_frame_mask = latents_out
-                else:
-                    latents, condition = latents_out
-                    first_frame_mask = None
-
-                noise = torch.randn_like(latents).to(device)
-                t = pipe.scheduler.timesteps[torch.randint(0, len(pipe.scheduler.timesteps), (1,)).item()].to(device)
-                noisy_latents = pipe.scheduler.add_noise(latents, noise, t)
-
-                if pipe.config.expand_timesteps:
-                    latent_model_input = (1 - first_frame_mask) * condition + first_frame_mask * noisy_latents
-                else:
-                    latent_model_input = torch.cat([noisy_latents, condition], dim=1)
-
-                # transformer forward (peft_module)
-                trans_out = peft_module(hidden_states=latent_model_input, timestep=t.unsqueeze(0).expand(latent_model_input.shape[0]), encoder_hidden_states=None, encoder_hidden_states_image=None, return_dict=False)
-                noise_pred = trans_out[0] if isinstance(trans_out, tuple) else trans_out
-                loss = loss_fn(noise_pred, noise)
-
-            # backward + step via accelerator
-            accelerator.backward(loss) if 'accelerator' in globals() else loss.backward()
-            optimizer.step()
-            optimizer.zero_grad()
-            global_step += 1
-
-            logs.append(f"step {global_step} loss {loss.item():.6f}")
-            pbar.set_postfix({"loss": f"{loss.item():.6f}"})
-
-            if max_train_steps and global_step >= max_train_steps:
-                break
-
-            if global_step % save_every_steps == 0:
-                out_sub = Path(output_dir) / f"lora_step_{global_step}"
-                out_sub.mkdir(parents=True, exist_ok=True)
-                try:
-                    peft_module.save_pretrained(str(out_sub))
-                except Exception:
-                    torch.save({k: v.cpu() for k, v in peft_module.state_dict().items()}, str(out_sub / "adapter_state_dict.pt"))
-                print(f"Saved adapter at {out_sub}")
-
-        if max_train_steps and global_step >= max_train_steps:
-            break
-
-    # final save
-    Path(output_dir).mkdir(parents=True, exist_ok=True)
-    try:
-        peft_module.save_pretrained(output_dir)
-    except Exception:
-        torch.save({k: v.cpu() for k, v in peft_module.state_dict().items()}, str(Path(output_dir) / "adapter_state_dict.pt"))
-
-    return output_dir, logs
-
-# -------------------------
-# Test generation
-# -------------------------
-def test_generation_load_and_run(base_model_id: str, adapter_dir: Optional[str], adapter_target: str, prompt: str, num_inference_steps: int = 8):
-    pipe = load_pipeline_auto(base_model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
-    # if adapter_dir provided, try to load adapter weights into target module (best-effort)
-    if adapter_dir:
-        try:
-            if adapter_target == "unet" and hasattr(pipe, "unet"):
-                # wrap unet with a matching peft config and load
-                lcfg = LoraConfig(r=8, lora_alpha=16, target_modules=find_target_modules(pipe.unet))
-                pipe.unet = get_peft_model(pipe.unet, lcfg)
-                try:
-                    pipe.unet.load_state_dict(torch.load(Path(adapter_dir) / "pytorch_model.bin"), strict=False)
-                except Exception:
-                    try:
-                        pipe.unet.load_adapter(adapter_dir)
-                    except Exception:
-                        pass
-            elif adapter_target == "transformer" and hasattr(pipe, "transformer"):
-                lcfg = LoraConfig(r=8, lora_alpha=16, target_modules=find_target_modules(pipe.transformer))
-                pipe.transformer = get_peft_model(pipe.transformer, lcfg)
-            elif adapter_target == "text_encoder" and hasattr(pipe, "text_encoder"):
-                lcfg = LoraConfig(r=8, lora_alpha=16, target_modules=find_target_modules(pipe.text_encoder))
-                pipe.text_encoder = get_peft_model(pipe.text_encoder, lcfg)
-        except Exception as e:
-            print("Adapter load attempt warning:", e)
-
-    pipe.to(DEVICE)
-    out = pipe(prompt=prompt, num_inference_steps=num_inference_steps)
-    if hasattr(out, "images"):
-        return out.images[0]
-    elif hasattr(out, "frames"):
-        frames = out.frames[0]
-        from PIL import Image
-        return Image.fromarray((frames[-1] * 255).clip(0,255).astype("uint8"))
-    else:
-        raise RuntimeError("No images or frames returned")
-
-# -------------------------
-# Upload adapter
-# -------------------------
-def upload_adapter(local_dir: str, repo_id: str) -> str:
-    token = os.environ.get("HF_TOKEN")
-    if token is None:
-        raise RuntimeError("HF_TOKEN not set in environment for upload")
+    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(folder_path=local_dir, repo_id=repo_id, repo_type="model", token=token)
+    upload_folder(local, repo_id=repo_id, repo_type="model", token=token)
     return f"https://huggingface.co/{repo_id}"
 
-# -------------------------
-# Gradio UI
-# -------------------------
-def run_all_ui(base_model_id: str,
-               dataset_source: str,
-               csv_name: str,
-               task_type: str,
-               adapter_target: str,
-               lora_r: int,
-               lora_alpha: int,
-               epochs: int,
-               batch_size: int,
-               lr: float,
-               max_train_steps: int,
-               output_dir: str,
-               upload_repo: str,
-               save_every_steps: int):
-    # minor mapping: QwenEdit/ prompt-lora -> text_encoder
-    if task_type == "prompt-lora":
-        adapter_target = "text_encoder"
-
-    try:
-        out_dir, logs = train_lora_accelerate(base_model_id, dataset_source, csv_name, adapter_target, output_dir,
-                                              epochs=epochs, batch_size=batch_size, lr=lr, max_train_steps=(max_train_steps if max_train_steps>0 else None),
-                                              lora_r=lora_r, lora_alpha=lora_alpha, max_frames=5, save_every_steps=save_every_steps)
-    except Exception as e:
-        return f"Training failed: {e}", None, None
-
-    link = None
-    if upload_repo:
-        try:
-            link = upload_adapter(out_dir, upload_repo)
-        except Exception as e:
-            link = f"Upload failed: {e}"
-
-    # quick test: use first prompt from dataset
-    try:
-        ds = MediaTextDataset(dataset_source, csv_name=csv_name, max_frames=5)
-        test_prompt = ds.df.iloc[0]["text"] if len(ds.df) > 0 else "A cat on a skateboard"
-    except Exception:
-        test_prompt = "A cat on a skateboard"
-
-    test_img = None
-    try:
-        test_img = test_generation_load_and_run(base_model_id, out_dir, adapter_target, test_prompt)
-    except Exception as e:
-        print("Test gen failed:", e)
-
-    return "\n".join(logs[-200:]), test_img, link
-
-def build_ui():
+# ---------------- Gradio UI ----------------
+def run_ui():
     with gr.Blocks() as demo:
-        gr.Markdown("# Universal LoRA Trainer (single-file) — Accelerate + PEFT")
+        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():
-            with gr.Column(scale=2):
-                base_model = gr.Textbox(label="Base model id (Diffusers)", value="runwayml/stable-diffusion-v1-5")
-                dataset_source = gr.Textbox(label="Dataset folder or HF repo (e.g. user/repo)", value="./dataset")
-                csv_name = gr.Textbox(label="CSV/Parquet filename", value="dataset.csv")
-                task_type = gr.Dropdown(label="Task type", choices=["text-image", "text-video", "prompt-lora"], value="text-image")
-                adapter_target = gr.Dropdown(label="Adapter target (unet/transformer/text_encoder)", choices=["unet", "transformer", "text_encoder"], value="unet")
-                lora_r = gr.Slider(1, 32, value=8, step=1, label="LoRA rank (r)")
-                lora_alpha = gr.Slider(1, 64, value=16, step=1, label="LoRA alpha")
-                epochs = gr.Number(label="Epochs", value=1)
-                batch_size = gr.Number(label="Batch size (per device)", value=1)
-                lr = gr.Number(label="Learning rate", value=1e-4)
-                max_train_steps = gr.Number(label="Max train steps (0 = unlimited)", value=0)
-                save_every_steps = gr.Number(label="Save every steps", value=200)
-                output_dir = gr.Textbox(label="Local output dir for adapter", value="./adapter_out")
-                upload_repo = gr.Textbox(label="Upload adapter to HF repo (optional, username/repo)", value="")
-                start_btn = gr.Button("Start training")
-            with gr.Column(scale=1):
-                logs = gr.Textbox(label="Training logs (tail)", lines=18)
-                sample_image = gr.Image(label="Sample generated frame after training")
-        def on_start(base_model_id, dataset_source, csv_name, task_type, adapter_target, lora_r, lora_alpha, epochs, batch_size, lr, max_train_steps, output_dir, upload_repo, save_every_steps):
-            return run_all_ui(base_model_id, dataset_source, csv_name, task_type, adapter_target, int(lora_r), int(lora_alpha), int(epochs), int(batch_size), float(lr), int(max_train_steps), output_dir, upload_repo, int(save_every_steps))
-        start_btn.click(on_start, inputs=[base_model, dataset_source, csv_name, task_type, adapter_target, lora_r, lora_alpha, epochs, batch_size, lr, max_train_steps, output_dir, upload_repo, save_every_steps], outputs=[logs, sample_image, gr.Textbox()])
+            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__":
-    demo = build_ui()
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+if __name__=="__main__":
+    run_ui().launch(server_name="0.0.0.0",server_port=7860)