Update app.py

app.py

@@ -1,122 +1,491 @@
-# app.py
+# universal_lora_trainer_accelerate.py
+"""
+Universal LoRA Trainer (Accelerate + PEFT) with Gradio UI.
+
+- Real LoRA training (UNet / ChronoEdit transformer / prompt enhancer)
+- Dataset: local folder or HF repo id containing dataset.csv with columns: file_name,text
+- HF_TOKEN is read from environment for uploads
+- QwenEdit / prompt-enhancer LoRA optional
+"""
+
 import os
-import gradio as gr
+import math
+import time
+import tempfile
+from pathlib import Path
+from typing import Optional, List, Tuple
+
 import torch
-from huggingface_hub import HfApi, Repository, upload_folder
+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, HfApi
+
 from diffusers import DiffusionPipeline
-from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-# =========================
-# 🧩 Helper Functions
-# =========================
-def load_diffusion_model(model_name, dtype=torch.float16):
-    print(f"Loading base model: {model_name}")
-    pipe = DiffusionPipeline.from_pretrained(model_name, torch_dtype=dtype)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    pipe.to(device)
-    return pipe
 
-def apply_lora(pipe, lora_path=None, lora_rank=8, alpha=16):
-    if lora_path and os.path.exists(lora_path):
-        print(f"Loading LoRA weights from {lora_path}")
-        pipe.load_lora_weights(lora_path)
+# optional ChronoEdit
+try:
+    from chronoedit_diffusers.pipeline_chronoedit import ChronoEditPipeline
+    CHRONOEDIT_AVAILABLE = True
+except Exception:
+    CHRONOEDIT_AVAILABLE = False
+
+# PEFT + Accelerate
+from peft import LoraConfig, get_peft_model
+from accelerate import Accelerator
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
+VIDEO_EXTS = {".mp4", ".mov", ".avi", ".mkv"}
+
+
+# -------------------------
+# Utilities
+# -------------------------
+def is_hub_repo_like(s: str) -> bool:
+    # simple heuristic: contains a slash and no local path separators
+    return "/" in s and not os.path.exists(s)
+
+
+def hf_download_file(repo_id: str, filename: str, local_cache_dir: Optional[str] = None, token: Optional[str] = None) -> str:
+    """Download a single file from HF repo to a temporary local path and return path."""
+    token = token or os.environ.get("HF_TOKEN")
+    out = hf_hub_download(repo_id=repo_id, filename=filename, use_auth_token=token)
+    return out
+
+
+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: local or HF repo
+# -------------------------
+class MediaTextDataset(Dataset):
+    """
+    CSV must have columns: file_name, text
+    If file_name is a local path (exists), loads from local; otherwise if dataset_repo is provided it downloads from hub.
+    """
+
+    def __init__(self, dataset_dir_or_repo: 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_dir_or_repo
+        self.is_hub = is_hub_repo_like(dataset_dir_or_repo)
+        self.df = None
+        self.root = None
+        self.tmpdir = None
+        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")
+
+        if self.is_hub:
+            # download CSV into temp dir
+            self.tmpdir = Path(tempfile.mkdtemp(prefix="dataset_hf_"))
+            csv_local = hf_download_file(self.source, csv_name, token=self.hub_token)
+            # hf_hub_download returns path inside local cache; copy into tmpdir for consistent file reads
+            csv_df = pd.read_csv(csv_local)
+            self.df = csv_df
+            # ensure we will download each referenced file on demand by storing repo id
+            self.root = None
+        else:
+            self.root = Path(dataset_dir_or_repo)
+            csv_path = self.root / csv_name
+            if not csv_path.exists():
+                raise FileNotFoundError(f"{csv_path} not found")
+            self.df = pd.read_csv(csv_path)
+
+        # 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:
+        # returns a local path for the file (cached)
+        # if local path exists return as-is
+        if self.root is not None:
+            p = self.root / file_name
+            if p.exists():
+                return str(p)
+        # else download from hub repo
+        local_path = hf_hub_download(repo_id=self.source, filename=file_name, use_auth_token=self.hub_token)
+        return local_path
+
+    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)
+        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 loading 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)
     else:
-        print("Initializing new LoRA config.")
-        config = LoraConfig(r=lora_rank, lora_alpha=alpha)
-        pipe.unet = get_peft_model(pipe.unet, config)
+        print(f"Loading standard Diffusers pipeline: {base_model_id}")
+        pipe = DiffusionPipeline.from_pretrained(base_model_id, torch_dtype=torch_dtype)
     return pipe
 
-def enhance_prompt(prompt, model_name="Qwen/Qwen2.5-1.5B-Instruct"):
-    print(f"Enhancing prompt with {model_name}")
-    tok = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16).to("cuda" if torch.cuda.is_available() else "cpu")
-    inputs = tok(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=100)
-    enhanced = tok.decode(outputs[0], skip_special_tokens=True)
-    return enhanced
-
-def train_lora(base_model, dataset_path, output_dir, steps=100, lr=1e-4, progress=gr.Progress(track_tqdm=True)):
-    progress(0, desc="Loading model...")
-    pipe = load_diffusion_model(base_model)
-    pipe = apply_lora(pipe)
-    progress(0.2, desc="Preparing dataset...")
-
-    # Dummy dataset loader (replace with your dataset logic)
-    import pandas as pd
-    df = pd.read_csv(dataset_path)
-    prompts = df['text'].tolist()
-
-    progress(0.3, desc="Training...")
-    for i, text in enumerate(prompts):
-        progress(0.3 + 0.6*(i/len(prompts)), desc=f"Training on sample {i+1}/{len(prompts)}")
-        # Simulate training step
-        torch.cuda.empty_cache()
-        torch.manual_seed(i)
-        _ = pipe(prompt=text, num_inference_steps=1)
-
-    os.makedirs(output_dir, exist_ok=True)
-    pipe.save_pretrained(output_dir)
-    progress(1, desc="Training complete ✅")
-    return output_dir
-
-def upload_to_hub(model_path, repo_id, token):
-    print(f"Uploading {model_path} to {repo_id}...")
-    api = HfApi()
-    upload_folder(repo_id=repo_id, folder_path=model_path, token=token)
-    return f"✅ Model uploaded to: https://huggingface.co/{repo_id}"
-
-def test_model(model_path, prompt):
-    pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16)
-    pipe.to("cuda" if torch.cuda.is_available() else "cpu")
-    image = pipe(prompt=prompt, num_inference_steps=8).images[0]
-    return image
-
-# =========================
-# 🎨 Gradio UI
-# =========================
-def gradio_app():
-    with gr.Blocks(title="Universal Diffusion Trainer") as demo:
-        gr.Markdown("## 🌌 Universal Diffusion Fine-tuner & Tester\nTrain or Test any Diffusion Model (T2I, T2V, LoRA, Prompt Enhancer)")
-
-        with gr.Tab("🔧 Training"):
-            base_model = gr.Textbox(label="Base Model (e.g., nvidia/ChronoEdit-14B-Diffusers)", value="runwayml/stable-diffusion-v1-5")
-            dataset = gr.Textbox(label="CSV Dataset Path (with columns file_name,text)", value="data.csv")
-            steps = gr.Slider(10, 1000, 100, step=10, label="Training Steps")
-            output_dir = gr.Textbox(label="Output Folder", value="./trained_model")
-            hf_repo = gr.Textbox(label="Upload to HF Repo (e.g., rahul7star/my-lora-model)")
-            hf_token = gr.Textbox(label="Hugging Face Token", type="password")
-            run_train = gr.Button("🚀 Start Training")
-            log = gr.Textbox(label="Logs")
-            progress = gr.HTML()
-            
-            def train_and_upload(base_model, dataset, steps, output_dir, hf_repo, hf_token, progress=gr.Progress(track_tqdm=True)):
-                output_path = train_lora(base_model, dataset, output_dir, steps, progress=progress)
-                if hf_repo and hf_token:
-                    url = upload_to_hub(output_path, hf_repo, hf_token)
-                    return f"Training done ✅\nUploaded to: {url}"
+
+def attach_lora(pipe, target: str, r: int = 8, alpha: int = 16, dropout: float = 0.0):
+    """
+    Attach LoRA to pipe.unet (image), pipe.transformer (video), or pipe.text_encoder (prompt)
+    Returns: modified pipe and the attribute name used
+    """
+    if target == "unet":
+        if not hasattr(pipe, "unet"):
+            raise RuntimeError("Chosen pipeline has no UNet")
+        target_module = pipe.unet
+        attr = "unet"
+    elif target == "transformer":
+        if not hasattr(pipe, "transformer"):
+            raise RuntimeError("Chosen pipeline has no transformer")
+        target_module = pipe.transformer
+        attr = "transformer"
+    elif target == "text_encoder":
+        if not hasattr(pipe, "text_encoder"):
+            raise RuntimeError("Chosen pipeline has no text_encoder")
+        target_module = pipe.text_encoder
+        attr = "text_encoder"
+    else:
+        raise RuntimeError("Unknown target for LoRA")
+
+    target_modules = find_target_modules(target_module)
+    print("LoRA target_modules:", 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)
+
+    # set back into pipeline
+    setattr(pipe, attr, peft_model)
+    return pipe, attr
+
+
+# -------------------------
+# Training loop (Accelerate)
+# -------------------------
+def train_lora_accelerate(base_model_id: str,
+                          dataset_dir_or_repo: 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,
+                          hub_token: Optional[str] = None,
+                          save_every_steps: int = 200) -> Tuple[str, List[str]]:
+    """
+    Run training using Accelerate. Returns (output_dir, logs)
+    """
+
+    accelerator = Accelerator()
+    device = accelerator.device
+
+    pipe = load_pipeline_auto(base_model_id, torch_dtype=torch.float16 if device.type == "cuda" else torch.float32)
+
+    dataset = MediaTextDataset(dataset_dir_or_repo, csv_name=csv_name, max_frames=max_frames)
+    dataloader = DataLoader(dataset, batch_size=1, shuffle=True, collate_fn=lambda x: x)
+
+    # attach LoRA to the chosen target
+    pipe, attr = attach_lora(pipe, adapter_target, r=lora_r, alpha=lora_alpha)
+    # Move model parts to device via accelerator.prepare
+    # For simplicity, we'll collect parameters to optimize
+    if adapter_target == "unet":
+        peft_module = pipe.unet
+    elif adapter_target == "transformer":
+        peft_module = pipe.transformer
+    else:
+        peft_module = pipe.text_encoder
+
+    # Collect trainable params
+    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)
+
+    # Also move pipeline core bits to device if required (VAE, scheduler) - only for inference functions
+    # We'll call pipeline components when needed, moving them to device manually
+    logs = []
+    global_step = 0
+    loss_fn = nn.MSELoss()
+
+    # prepare scheduler timesteps if available
+    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:
+            example = batch[0]
+            if example["type"] == "image":
+                # image flow
+                img = example["image"].unsqueeze(0).to(device)
+                caption = [example["caption"]]
+                if not hasattr(pipe, "encode_prompt"):
+                    raise RuntimeError("Pipeline has no encode_prompt")
+                # CALL encode_prompt on CPU side: move text encoder temporarily
+                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 missing VAE")
+                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)
+                # UNet forward (peft_module is already prepared and on device by accelerator)
+                # For accelerate we must call through the pipeline's UNet wrapper; if we replaced pipe.unet earlier,
+                # ensure the call signature matches: many UNets return a ModelOutput; for simplicity attempt common API
+                noise_pred = peft_module(noisy_latents, t.expand(noisy_latents.shape[0]), encoder_hidden_states=prompt_embeds)[0] if isinstance(peft_module(noisy_latents, t.expand(noisy_latents.shape[0]), encoder_hidden_states=prompt_embeds), tuple) else peft_module(noisy_latents, t.expand(noisy_latents.shape[0]), encoder_hidden_states=prompt_embeds).sample
+                loss = loss_fn(noise_pred, noise)
+            else:
+                # video flow (ChronoEdit simplified)
+                if not CHRONOEDIT_AVAILABLE:
+                    raise RuntimeError("ChronoEdit pipeline not available in this environment")
+                frames = example["frames"].unsqueeze(0).to(device)  # [1, T, C, H, W]
+                # preprocess frames into pipeline expected format
+                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:
-                    return f"Training done ✅\nModel saved at {output_path}"
+                    latent_model_input = torch.cat([noisy_latents, condition], dim=1)
+                # transformer forward
+                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 = out[0]
+                loss = loss_fn(noise_pred, noise)
+
+            accelerator.backward(loss)
+            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}"})
 
-            run_train.click(train_and_upload, 
-                            inputs=[base_model, dataset, steps, output_dir, hf_repo, hf_token],
-                            outputs=[log])
+            if max_train_steps and global_step >= max_train_steps:
+                break
 
-        with gr.Tab("🧪 Test Model"):
-            test_model_path = gr.Textbox(label="Model Path or Repo", value="./trained_model")
-            test_prompt = gr.Textbox(label="Prompt", value="A futuristic city with flying cars at sunset")
-            test_btn = gr.Button("🖼️ Generate")
-            test_output = gr.Image(label="Generated Output")
-            test_btn.click(test_model, inputs=[test_model_path, test_prompt], outputs=test_output)
+            if global_step % save_every_steps == 0:
+                # save PEFT adapter
+                out_sub = Path(output_dir) / f"lora_step_{global_step}"
+                out_sub.mkdir(parents=True, exist_ok=True)
+                try:
+                    # try to call save_pretrained on peft wrapper
+                    peft_module.save_pretrained(str(out_sub))
+                except Exception as e:
+                    # fallback to saving state_dict
+                    torch.save({k: v.cpu() for k, v in peft_module.state_dict().items()}, str(out_sub / "adapter_state_dict.pt"))
+                print(f"Saved intermediate adapter at {out_sub}")
 
-        with gr.Tab("✨ Prompt Enhancement"):
-            prompt_input = gr.Textbox(label="Input Prompt")
-            enhance_btn = gr.Button("Enhance")
-            enhanced_out = gr.Textbox(label="Enhanced Prompt")
-            enhance_btn.click(enhance_prompt, inputs=[prompt_input], outputs=[enhanced_out])
+        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):
+    # Load base pipeline
+    pipe = load_pipeline_auto(base_model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+    # If adapter_dir is provided, load the adapter into the same target
+    if adapter_dir:
+        if adapter_target == "unet":
+            # peft: load_pretrained onto the module
+            if hasattr(pipe, "unet"):
+                pipe.unet = get_peft_model(pipe.unet, LoraConfig(r=8, lora_alpha=16, target_modules=find_target_modules(pipe.unet)))
+                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:
+                        print("Adapter loading fallbacks")
+        elif adapter_target == "transformer":
+            if hasattr(pipe, "transformer"):
+                pipe.transformer = get_peft_model(pipe.transformer, LoraConfig(r=8, lora_alpha=16, target_modules=find_target_modules(pipe.transformer)))
+                # loader fallback
+        elif adapter_target == "text_encoder":
+            if hasattr(pipe, "text_encoder"):
+                pipe.text_encoder = get_peft_model(pipe.text_encoder, LoraConfig(r=8, lora_alpha=16, target_modules=find_target_modules(pipe.text_encoder)))
+
+    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"))
+    raise RuntimeError("Pipeline returned no images/frames")
+
+
+# -------------------------
+# Upload adapter to HF Hub
+# -------------------------
+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")
+    create_repo(repo_id, exist_ok=True)
+    upload_folder(folder_path=local_dir, repo_id=repo_id, repo_type="model", token=token)
+    return f"https://huggingface.co/{repo_id}"
+
+
+# -------------------------
+# Gradio UI wiring
+# -------------------------
+def run_all_ui(base_model_id: str,
+               dataset_source: str,
+               csv_name: str,
+               mode: 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):
+    # training
+    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
+
+    # upload (if requested)
+    link = None
+    if upload_repo:
+        try:
+            link = upload_adapter(out_dir, upload_repo)
+        except Exception as e:
+            link = f"Upload failed: {e}"
+
+    # test generation with first prompt in 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 generation error:", e)
+
+    return "\n".join(logs[-200:]), test_img, link
+
+
+def build_ui():
+    with gr.Blocks() as demo:
+        gr.Markdown("# Universal LoRA Trainer (Accelerate + PEFT)")
+        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. username/repo)", value="./dataset")
+                csv_name = gr.Textbox(label="CSV filename", value="dataset.csv")
+                mode = gr.Radio(["text-image", "text-video", "prompt-lora"], label="Mode", 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, user/repo)", value="")
+                start_btn = gr.Button("Start training")
+            with gr.Column(scale=1):
+                logs = gr.Textbox(label="Training logs (tail)", lines=20)
+                sample_image = gr.Image(label="Sample generated frame after training")
+        def on_start(base_model_id, dataset_source, csv_name, mode, 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, mode, 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, mode, 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()])
     return demo
 
+
 if __name__ == "__main__":
-    demo = gradio_app()
-    demo.launch()
+    demo = build_ui()
+    demo.launch(server_name="0.0.0.0", server_port=7860)