app_quant.py

# universal_lora_trainer_quant_dynamic.py
"""
Universal Dynamic LoRA Trainer (Accelerate + PEFT) with optional QLoRA 4-bit support.

- Supports CSV and Parquet dataset files (columns: file_name, text)
- Accepts dataset from a local folder or Hugging Face dataset repo id (username/repo)
- Real LoRA training (PEFT) for:
    * text->image (UNet)
    * text->video (ChronoEdit transformer)
    * prompt-enhancer (text_encoder / QwenEdit)
- Optional:
    * 4-bit quantization (bitsandbytes / QLoRA)
    * xFormers / FlashAttention
    * AdaLoRA (if available)
- Uses HF_TOKEN from environment for upload
- Use `accelerate launch` for multi-GPU / optimized run
"""

import os
import math
import tempfile
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, list_repo_files

from diffusers import DiffusionPipeline

# optional pip installs - guard imports
try:
    from chronoedit_diffusers.pipeline_chronoedit import ChronoEditPipeline
    CHRONOEDIT_AVAILABLE = True
except Exception:
    CHRONOEDIT_AVAILABLE = False

# Qwen image edit optional
try:
    from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPipeline  # optional name
    QWENEDIT_AVAILABLE = True
except Exception:
    QWENEDIT_AVAILABLE = False

# BitsAndBytes (quantization)
try:
    from transformers import BitsAndBytesConfig
    BNB_AVAILABLE = True
except Exception:
    BitsAndBytesConfig = None
    BNB_AVAILABLE = False

# xFormers
try:
    import xformers  # noqa
    XFORMERS_AVAILABLE = True
except Exception:
    XFORMERS_AVAILABLE = False

# PEFT / AdaLoRA
try:
    from peft import LoraConfig, get_peft_model
    try:
        from peft import AdaLoraConfig  # optional
        ADALORA_AVAILABLE = True
    except Exception:
        AdaLoraConfig = None
        ADALORA_AVAILABLE = False
except Exception as e:
    raise RuntimeError("Install peft: pip install peft") from e

# Accelerate
try:
    from accelerate import Accelerator
except Exception as e:
    raise RuntimeError("Install accelerate: pip install accelerate") from e

# ------------------------
# Config
# ------------------------
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:
    return "/" in s and not Path(s).exists()

def download_from_hf(repo_id: str, filename: str, token: Optional[str] = None, repo_type: str = "dataset") -> str:
    token = token or os.environ.get("HF_TOKEN")
    return hf_hub_download(repo_id=repo_id, filename=filename, use_auth_token=token, repo_type=repo_type)

def try_list_repo_files(repo_id: str, repo_type: str = "dataset", token: Optional[str] = None):
    token = token or os.environ.get("HF_TOKEN")
    try:
        return list_repo_files(repo_id, token=token, repo_type=repo_type)
    except Exception:
        return []

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 class (CSV/Parquet)
# ------------------------
class MediaTextDataset(Dataset):
    """
    Loads records from CSV or Parquet with columns:
      - file_name (relative path in folder or filename inside HF dataset repo)
      - text
    """
    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")

        # load dataframe (CSV or parquet)
        if self.is_hub:
            # try CSV then parquet; specify repo_type="dataset"
            searched = try_list_repo_files(self.source, repo_type="dataset", token=self.hub_token)
            # prefer exact csv_name
            try:
                csv_local = download_from_hf(self.source, csv_name, token=self.hub_token, repo_type="dataset")
            except Exception:
                # try .parquet variant
                alt = csv_name.replace(".csv", ".parquet") if csv_name.endswith(".csv") else csv_name + ".parquet"
                csv_local = download_from_hf(self.source, alt, token=self.hub_token, repo_type="dataset")
            if str(csv_local).endswith(".parquet"):
                df = pd.read_parquet(csv_local)
            else:
                df = pd.read_csv(csv_local)
            self.df = df
            self.root = None
        else:
            root = Path(dataset_source)
            csv_path = root / csv_name
            parquet_path = root / csv_name.replace(".csv", ".parquet") if csv_name.endswith(".csv") else root / (csv_name + ".parquet")
            if csv_path.exists():
                self.df = pd.read_csv(csv_path)
            elif parquet_path.exists():
                self.df = pd.read_parquet(parquet_path)
            else:
                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 dataset repo
        return download_from_hf(self.source, file_name, token=self.hub_token, repo_type="dataset")

    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 loader with optional quantization
# ------------------------
def load_pipeline_auto(base_model_id: str, use_4bit: bool = False, bnb_config: Optional[object] = None, torch_dtype=torch.float16):
    low = base_model_id.lower()
    is_chrono = "chrono" in low or "wan" in low or "video" in low
    is_qwen = "qwen" in low or "qwenimage" in low
    # choose pipeline
    if is_chrono and CHRONOEDIT_AVAILABLE:
        print("Loading ChronoEdit pipeline")
        # ChronoEdit may not accept quant config; try with safer call
        if use_4bit and bnb_config is not None:
            pipe = ChronoEditPipeline.from_pretrained(base_model_id, torch_dtype=torch.float16)  # quantized loading of chronoedit not widely supported
        else:
            pipe = ChronoEditPipeline.from_pretrained(base_model_id, torch_dtype=torch_dtype)
    elif is_qwen and QWENEDIT_AVAILABLE:
        print("Loading QWEN image-edit pipeline")
        pipe = QwenImageEditPipeline.from_pretrained(base_model_id, torch_dtype=torch_dtype)
    else:
        # fallback to DiffusionPipeline - supports quantization_config for diffusers+transformers
        print("Loading standard DiffusionPipeline:", base_model_id, "use_4bit=", use_4bit)
        if use_4bit and BNB_AVAILABLE and bnb_config is not None:
            pipe = DiffusionPipeline.from_pretrained(base_model_id, quantization_config=bnb_config, torch_dtype=torch.float16)
        else:
            pipe = DiffusionPipeline.from_pretrained(base_model_id, torch_dtype=torch_dtype)
    return pipe

# ------------------------
# Auto infer adapter target
# ------------------------
def infer_target_for_task(task_type: str, model_name: str) -> str:
    low = model_name.lower()
    if task_type == "prompt-lora" or "qwen" in low or "qwenedit" in low:
        return "text_encoder"
    if task_type == "text-video" or "chrono" in low or "wan" in low:
        return "transformer"
    # default
    return "unet"

# ------------------------
# LoRA attach (supports AdaLoRA if available)
# ------------------------
def attach_lora(pipe, adapter_target: str, r: int = 8, alpha: int = 16, dropout: float = 0.0, use_adalora: bool = False):
    if adapter_target == "unet":
        if not hasattr(pipe, "unet"):
            raise RuntimeError("Pipeline has no UNet to attach LoRA")
        target_module = pipe.unet
        attr = "unet"
    elif adapter_target == "transformer":
        if not hasattr(pipe, "transformer"):
            raise RuntimeError("Pipeline has no transformer to attach LoRA")
        target_module = pipe.transformer
        attr = "transformer"
    elif adapter_target == "text_encoder":
        if not hasattr(pipe, "text_encoder"):
            # some models name it differently; try encoder attribute fallback
            if hasattr(pipe, "text_encoder"):
                target_module = pipe.text_encoder
                attr = "text_encoder"
            else:
                raise RuntimeError("Pipeline has no text_encoder for prompt-loRA")
        else:
            target_module = pipe.text_encoder
            attr = "text_encoder"
    else:
        raise RuntimeError("Unknown adapter_target")

    target_modules = find_target_modules(target_module)
    print("Detected target_modules for LoRA:", target_modules)

    if use_adalora and ADALORA_AVAILABLE:
        lora_config = AdaLoraConfig(
            r=r,
            lora_alpha=alpha,
            target_modules=target_modules,
            init_r=4,
            lora_dropout=dropout,
        )
    else:
        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

# ------------------------
# Training loop (Accelerate-aware)
# ------------------------
def train_lora_accelerate(base_model_id: str,
                          dataset_source: str,
                          csv_name: str,
                          task_type: str,
                          adapter_target_override: Optional[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,
                          use_4bit: bool = False,
                          enable_xformers: bool = False,
                          use_adalora: bool = False,
                          gradient_accumulation_steps: int = 1,
                          mixed_precision: Optional[str] = None,
                          save_every_steps: int = 200,
                          max_frames: int = 5):

    # Setup Accelerator
    accelerator = Accelerator(mixed_precision=mixed_precision or ("fp16" if torch.cuda.is_available() else "no")),
    # Note: Accelerator is returned as a tuple if trailing comma; fix:
    accelerator = accelerator if isinstance(accelerator, Accelerator) else accelerator[0]
    device = accelerator.device

    # prepare bitsandbytes config if requested
    bnb_conf = None
    if use_4bit and BNB_AVAILABLE:
        bnb_conf = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
        )

    # Load pipeline (supports quant for standard diffusers)
    pipe = load_pipeline_auto(base_model_id, use_4bit=use_4bit, bnb_config=bnb_conf, torch_dtype=torch.float16 if device.type == "cuda" else torch.float32)

    # optionally enable memory efficient attention
    if enable_xformers:
        try:
            if hasattr(pipe, "enable_xformers_memory_efficient_attention"):
                pipe.enable_xformers_memory_efficient_attention()
            elif hasattr(pipe, "enable_attention_slicing"):
                pipe.enable_attention_slicing()
            print("xFormers / memory efficient attention enabled.")
        except Exception as e:
            print("Could not enable xformers:", e)

    # infer adapter target automatically if not overridden
    adapter_target = adapter_target_override if adapter_target_override else infer_target_for_task(task_type, base_model_id)
    print("Adapter target set to:", adapter_target)

    # attach LoRA
    pipe, attr = attach_lora(pipe, adapter_target, r=lora_r, alpha=lora_alpha, dropout=0.0, use_adalora=use_adalora)
    # pick the peft module for optimization
    peft_module = getattr(pipe, attr)

    # dataset + dataloader (we use batch_size=1 collate)
    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)

    # optimizer
    trainable_params = [p for n,p in peft_module.named_parameters() if p.requires_grad]
    optimizer = torch.optim.AdamW(trainable_params, lr=lr)

    # prepare objects with accelerator
    peft_module, optimizer, dataloader = accelerator.prepare(peft_module, optimizer, dataloader)

    # training loop
    logs = []
    global_step = 0
    loss_fn = nn.MSELoss()

    # scheduler setup if available
    if hasattr(pipe, "scheduler"):
        try:
            pipe.scheduler.set_timesteps(50, device=device)
            timesteps = pipe.scheduler.timesteps
        except Exception:
            timesteps = None
    else:
        timesteps = None

    # Training
    for epoch in range(int(epochs)):
        pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{epochs}")
        for batch in pbar:
            example = batch[0]
            # image flow
            if example["type"] == "image":
                img = example["image"].unsqueeze(0).to(device)
                caption = [example["caption"]]

                if not hasattr(pipe, "encode_prompt"):
                    raise RuntimeError("Pipeline lacks encode_prompt - cannot encode 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,
                )

                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)
                if timesteps is None:
                    t = torch.tensor(1, device=device)
                else:
                    t = pipe.scheduler.timesteps[torch.randint(0, len(pipe.scheduler.timesteps), (1,)).item()].to(device)
                noisy_latents = pipe.scheduler.add_noise(latents, noise, t)

                # forward through peft_module (unet)
                out = peft_module(noisy_latents, t.expand(noisy_latents.shape[0]), encoder_hidden_states=prompt_embeds)
                if hasattr(out, "sample"):
                    noise_pred = out.sample
                elif isinstance(out, tuple):
                    noise_pred = out[0]
                else:
                    noise_pred = out

                loss = loss_fn(noise_pred, noise)

            else:
                # video flow (ChronoEdit simplified)
                if not CHRONOEDIT_AVAILABLE:
                    raise RuntimeError("ChronoEdit training requested but not installed in environment")
                frames = example["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)

                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] if isinstance(out, tuple) else out
                loss = loss_fn(noise_pred, noise)

            # backward and optimizer step (accelerator)
            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}"})

            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 (best-effort)
# ------------------------
def test_generation_load_and_run(base_model_id: str, adapter_dir: Optional[str], adapter_target: str, prompt: str, use_4bit: bool = False):
    # load base pipeline (no heavy quant config)
    bnb_conf = None
    if use_4bit and BNB_AVAILABLE:
        bnb_conf = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")
    pipe = load_pipeline_auto(base_model_id, use_4bit=use_4bit, bnb_config=bnb_conf, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)

    # attempt to load adapter into target module (best-effort)
    try:
        if adapter_target == "unet" and hasattr(pipe, "unet"):
            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 warning", e)

    pipe.to(DEVICE)
    out = pipe(prompt=prompt, num_inference_steps=8)
    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("No images/frames returned")

# ------------------------
# 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}"

# ------------------------
# UI: Boost info helper
# ------------------------
def boost_info_text(use_4bit: bool, enable_xformers: bool, mixed_precision: Optional[str], device_type: str):
    lines = []
    lines.append(f"Device: {device_type.upper()}")
    if use_4bit and BNB_AVAILABLE:
        lines.append("4-bit QLoRA enabled: ~4x memory saving (bitsandbytes NF4 + double quant).")
    else:
        lines.append("QLoRA disabled or bitsandbytes not installed.")
    if enable_xformers and XFORMERS_AVAILABLE:
        lines.append("xFormers/FlashAttention: memory-efficient attention enabled (faster & lower mem).")
    else:
        lines.append("xFormers not enabled or not installed.")
    if mixed_precision:
        lines.append(f"Mixed precision: {mixed_precision}")
    else:
        lines.append("Mixed precision: default (no automatic FP16/BF16).")
    lines.append("Expected: GPU + 4-bit + xFormers = fastest. CPU + 4-bit possible but slow.")
    return "\n".join(lines)

# ------------------------
# Gradio UI wiring
# ------------------------
def run_all_ui(base_model_id: str,
               dataset_source: str,
               csv_name: str,
               task_type: str,
               adapter_target_override: str,
               lora_r: int,
               lora_alpha: int,
               epochs: int,
               batch_size: int,
               lr: float,
               max_train_steps: int,
               output_dir: str,
               upload_repo: str,
               use_4bit: bool,
               enable_xformers: bool,
               use_adalora: bool,
               grad_accum: int,
               mixed_precision: str,
               save_every_steps: int):
    # map task_type -> adapter_target if override empty
    adapter_target = adapter_target_override or infer_target_for_task(task_type, base_model_id)
    try:
        out_dir, logs = train_lora_accelerate(
            base_model_id,
            dataset_source,
            csv_name,
            task_type,
            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,
            use_4bit=use_4bit,
            enable_xformers=enable_xformers,
            use_adalora=use_adalora,
            gradient_accumulation_steps=grad_accum,
            mixed_precision=(mixed_precision if mixed_precision != "none" else None),
            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 generation using first dataset prompt
    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, use_4bit=use_4bit)
    except Exception as e:
        print("Test gen failed:", e)

    return "\n".join(logs[-200:]), test_img, link

def build_ui():
    with gr.Blocks() as demo:
        gr.Markdown("# Universal LoRA Trainer — Quantization & Speedups (single-file)")
        with gr.Row():
            with gr.Column(scale=2):
                base_model = gr.Textbox(label="Base model id (Diffusers / ChronoEdit / Qwen)", value="runwayml/stable-diffusion-v1-5")
                dataset_source = gr.Textbox(label="Dataset folder or HF dataset repo (username/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_override = gr.Textbox(label="Adapter target override (leave blank for auto)", value="")
                lora_r = gr.Slider(1, 64, value=8, step=1, label="LoRA rank (r)")
                lora_alpha = gr.Slider(1, 128, 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="")
            with gr.Column(scale=1):
                gr.Markdown("## Speed / Quantization")
                use_4bit = gr.Checkbox(label="Enable 4-bit QLoRA (bitsandbytes)", value=False)
                enable_xformers = gr.Checkbox(label="Enable xFormers / memory efficient attention", value=False)
                use_adalora = gr.Checkbox(label="Use AdaLoRA (if available in peft)", value=False)
                grad_accum = gr.Number(label="Gradient accumulation steps", value=1)
                mixed_precision = gr.Radio(choices=["none", "fp16", "bf16"], value=("fp16" if torch.cuda.is_available() else "none"), label="Mixed precision")
                gr.Markdown("### Boost Info")
                boost_info = gr.Textbox(label="Expected boost / notes", value="", lines=6)
                start_btn = gr.Button("Start Training")
        with gr.Row():
            logs = gr.Textbox(label="Training logs (tail)", lines=18)
            sample_image = gr.Image(label="Sample generated frame after training")
            upload_link = gr.Textbox(label="Upload link / status")
        def on_start(base_model, dataset_source, csv_name, task_type, adapter_target_override, lora_r, lora_alpha, epochs, batch_size, lr, max_train_steps, output_dir, upload_repo, use_4bit_val, enable_xformers_val, use_adalora_val, grad_accum_val, mixed_precision_val, save_every_steps):
            boost_text = boost_info_text(use_4bit_val, enable_xformers_val, mixed_precision_val, "gpu" if torch.cuda.is_available() else "cpu")
            # start training (blocking)
            logs_out, sample, link = run_all_ui(base_model, dataset_source, csv_name, task_type, adapter_target_override, int(lora_r), int(lora_alpha), int(epochs), int(batch_size), float(lr), int(max_train_steps), output_dir, upload_repo, use_4bit_val, enable_xformers_val, use_adalora_val, int(grad_accum_val), mixed_precision_val, int(save_every_steps))
            return boost_text + "\n\n" + logs_out, sample, link
        start_btn.click(on_start, inputs=[base_model, dataset_source, csv_name, task_type, adapter_target_override, lora_r, lora_alpha, epochs, batch_size, lr, max_train_steps, output_dir, upload_repo, use_4bit, enable_xformers, use_adalora, grad_accum, mixed_precision, save_every_steps], outputs=[boost_info, sample_image, upload_link])
    return demo

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