Upload 103 files

ui/src/app/api/datasets/create/route.tsx

@@ -3,22 +3,80 @@ import fs from 'fs';
 import path from 'path';
 import { getDatasetsRoot } from '@/server/settings';
 
+const sanitizeSegment = (value: string): string => {
+  if (!value || typeof value !== 'string') {
+    return '';
+  }
+  return value
+    .toLowerCase()
+    .replace(/[^a-z0-9]+/g, '_')
+    .replace(/^_+|_+$/g, '');
+};
+
+const ensureDirectory = (dirPath: string) => {
+  if (!fs.existsSync(dirPath)) {
+    fs.mkdirSync(dirPath, { recursive: true });
+  }
+};
+
+const resolveDatasetName = (rootPath: string, desiredName: string, namespace?: string | null) => {
+  const baseName = sanitizeSegment(desiredName) || 'dataset';
+  const namespaceSuffix = sanitizeSegment(namespace || '');
+
+  const datasetExists = (candidate: string) => fs.existsSync(path.join(rootPath, candidate));
+
+  if (!datasetExists(baseName)) {
+    return { name: baseName, path: path.join(rootPath, baseName) };
+  }
+
+  if (namespaceSuffix) {
+    let candidate = sanitizeSegment(`${baseName}_${namespaceSuffix}`) || `${baseName}_${namespaceSuffix}`;
+    let attempts = 0;
+    while (datasetExists(candidate)) {
+      attempts += 1;
+      if (attempts > 50) {
+        throw new Error('Unable to allocate unique dataset name');
+      }
+      candidate = sanitizeSegment(`${candidate}_${namespaceSuffix}`) || `${candidate}_${namespaceSuffix}`;
+    }
+    return { name: candidate, path: path.join(rootPath, candidate) };
+  }
+
+  let counter = 2;
+  while (counter < 1000) {
+    const candidate = sanitizeSegment(`${baseName}_${counter}`) || `${baseName}_${counter}`;
+    if (!datasetExists(candidate)) {
+      return { name: candidate, path: path.join(rootPath, candidate) };
+    }
+    counter += 1;
+  }
+
+  throw new Error('Unable to allocate unique dataset name');
+};
+
 export async function POST(request: Request) {
   try {
     const body = await request.json();
-    let { name } = body;
-    // clean name by making lower case,  removing special characters, and replacing spaces with underscores
-    name = name.toLowerCase().replace(/[^a-z0-9]+/g, '_');
-
-    let datasetsPath = await getDatasetsRoot();
-    let datasetPath = path.join(datasetsPath, name);
+    const { name, namespace } = body ?? {};
 
-    // if folder doesnt exist, create it
-    if (!fs.existsSync(datasetPath)) {
-      fs.mkdirSync(datasetPath, { recursive: true });
+    if (!name || typeof name !== 'string') {
+      throw new Error('Dataset name is required');
     }
 
-    return NextResponse.json({ success: true, name: name, path: datasetPath });
+    const datasetsPath = await getDatasetsRoot();
+    const { name: resolvedName, path: datasetPath } = resolveDatasetName(
+      datasetsPath,
+      name,
+      typeof namespace === 'string' ? namespace : null,
+    );
+
+    ensureDirectory(datasetPath);
+
+    return NextResponse.json({
+      success: true,
+      name: resolvedName,
+      path: datasetPath,
+    });
   } catch (error: any) {
     console.error('Dataset create error:', error);
     return NextResponse.json({ error: error?.message || 'Failed to create dataset' }, { status: 500 });

ui/src/app/api/hf-jobs/route.ts

@@ -401,14 +401,13 @@ def upload_results(output_path: str, model_name: str, namespace: str, token: str
     import tempfile
     import shutil
     import glob
-    import re
-    import yaml
     from datetime import datetime
     from huggingface_hub import create_repo, upload_file, HfApi
-    
+    from collections import deque
+
     try:
         repo_id = f"{namespace}/{model_name}"
-        
+
         # Create repository
         create_repo(repo_id=repo_id, token=token, exist_ok=True)
         
@@ -453,30 +452,86 @@ def upload_results(output_path: str, model_name: str, namespace: str, token: str
                     uploaded_files.append(filename)
                     config_files_uploaded.append(filename)
             
-            # 2. Handle sample images
-            samples_uploaded = []
+            def prepare_sample_metadata(samples_directory: str, sample_conf: dict):
+                if not samples_directory or not os.path.isdir(samples_directory):
+                    return [], []
+
+                allowed_ext = {'.jpg', '.jpeg', '.png', '.webp'}
+                image_records = []
+                for root, _, files in os.walk(samples_directory):
+                    for filename in files:
+                        ext = os.path.splitext(filename)[1].lower()
+                        if ext not in allowed_ext:
+                            continue
+                        abs_path = os.path.join(root, filename)
+                        try:
+                            mtime = os.path.getmtime(abs_path)
+                        except Exception:
+                            mtime = 0
+                        image_records.append((abs_path, mtime))
+
+                if not image_records:
+                    return [], []
+
+                image_records.sort(key=lambda item: (-item[1], item[0]))
+                image_queue = deque(image_records)
+
+                samples_list = sample_conf.get("samples", []) if sample_conf else []
+                if not samples_list:
+                    legacy = sample_conf.get("prompts", []) if sample_conf else []
+                    samples_list = [{"prompt": prompt} for prompt in legacy if prompt]
+
+                curated_samples = []
+                for sample in samples_list:
+                    prompt = None
+                    if isinstance(sample, dict):
+                        prompt = sample.get("prompt")
+                    elif isinstance(sample, str):
+                        prompt = sample
+
+                    if not prompt:
+                        continue
+
+                    if not image_queue:
+                        break
+
+                    image_path, _ = image_queue.popleft()
+                    repo_rel_path = f"images/{os.path.basename(image_path)}"
+                    curated_samples.append({
+                        "prompt": prompt,
+                        "local_path": image_path,
+                        "repo_path": repo_rel_path,
+                    })
+
+                all_files = [record[0] for record in image_records]
+                return curated_samples, all_files
+
             samples_dir = os.path.join(output_path, "samples")
-            if os.path.isdir(samples_dir):
+            sample_config = config.get("config", {}).get("process", [{}])[0].get("sample", {})
+            curated_samples, sample_files = prepare_sample_metadata(samples_dir, sample_config)
+
+            samples_uploaded = []
+            if sample_files:
                 print("Uploading sample images...")
-                # Create samples directory in repo
-                for filename in os.listdir(samples_dir):
-                    if filename.lower().endswith(('.jpg', '.jpeg', '.png', '.webp')):
-                        file_path = os.path.join(samples_dir, filename)
-                        repo_path = f"samples/{filename}"
-                        api.upload_file(
-                            path_or_fileobj=file_path,
-                            path_in_repo=repo_path,
-                            repo_id=repo_id,
-                            token=token
-                        )
-                        samples_uploaded.append(repo_path)
-            
+                for file_path in sample_files:
+                    if not os.path.isfile(file_path):
+                        continue
+                    filename = os.path.basename(file_path)
+                    repo_path = f"images/{filename}"
+                    api.upload_file(
+                        path_or_fileobj=file_path,
+                        path_in_repo=repo_path,
+                        repo_id=repo_id,
+                        token=token
+                    )
+                    samples_uploaded.append(repo_path)
+
             # 3. Generate and upload README.md
             readme_content = generate_model_card_readme(
                 repo_id=repo_id,
                 config=config,
                 model_name=model_name,
-                samples_dir=samples_dir if os.path.isdir(samples_dir) else None,
+                curated_samples=curated_samples,
                 uploaded_files=uploaded_files
             )
             
@@ -500,12 +555,11 @@ def upload_results(output_path: str, model_name: str, namespace: str, token: str
         print(f"Failed to upload model: {e}")
         raise e
 
-def generate_model_card_readme(repo_id: str, config: dict, model_name: str, samples_dir: str = None, uploaded_files: list = None) -> str:
+def generate_model_card_readme(repo_id: str, config: dict, model_name: str, curated_samples: list = None, uploaded_files: list = None) -> str:
     """Generate README.md content for the model card based on AI Toolkit's implementation"""
-    import re
     import yaml
     import os
-    
+
     try:
         # Extract configuration details
         process_config = config.get("config", {}).get("process", [{}])[0]
@@ -545,40 +599,27 @@ def generate_model_card_readme(repo_id: str, config: dict, model_name: str, samp
         # Add LoRA-specific tags
         tags.extend(["lora", "diffusers", "template:sd-lora", "ai-toolkit"])
         
-        # Generate widgets from sample images and prompts
+        # Generate widgets and gallery section from sample images
+        curated_samples = curated_samples or []
+
         widgets = []
-        if samples_dir and os.path.isdir(samples_dir):
-            sample_prompts = sample_config.get("samples", [])
-            if not sample_prompts:
-                # Fallback to old format
-                sample_prompts = [{"prompt": p} for p in sample_config.get("prompts", [])]
-            
-            # Get sample image files
-            sample_files = []
-            if os.path.isdir(samples_dir):
-                for filename in os.listdir(samples_dir):
-                    if filename.lower().endswith(('.jpg', '.jpeg', '.png', '.webp')):
-                        # Parse filename pattern: timestamp__steps_index.jpg
-                        match = re.search(r"__(\d+)_(\d+)\.jpg$", filename)
-                        if match:
-                            steps, index = int(match.group(1)), int(match.group(2))
-                            # Only use samples from final training step
-                            final_steps = train_config.get("steps", 1000)
-                            if steps == final_steps:
-                                sample_files.append((index, f"samples/{filename}"))
-            
-            # Sort by index and create widgets
-            sample_files.sort(key=lambda x: x[0])
-            
-            for i, prompt_obj in enumerate(sample_prompts):
-                prompt = prompt_obj.get("prompt", "") if isinstance(prompt_obj, dict) else str(prompt_obj)
-                if i < len(sample_files):
-                    _, image_path = sample_files[i]
-                    widgets.append({
-                        "text": prompt,
-                        "output": {"url": image_path}
-                    })
-        
+        prompt_bullets = []
+        for sample in curated_samples:
+            prompt_text = str(sample.get("prompt", "")).strip()
+            repo_path = sample.get("repo_path")
+            if not prompt_text or not repo_path:
+                continue
+            widgets.append({
+                "text": prompt_text,
+                "output": {"url": repo_path}
+            })
+            prompt_md = prompt_text.replace("`", "\`")
+            prompt_bullets.append(f"- `{prompt_md}`")
+
+        gallery_section = "<Gallery />\n\n"
+        if prompt_bullets:
+            gallery_section += "### Prompts\n\n" + "\n".join(prompt_bullets) + "\n\n"
+
         # Determine torch dtype based on model
         dtype = "torch.bfloat16" if "flux" in arch.lower() else "torch.float16"
         
@@ -598,6 +639,16 @@ def generate_model_card_readme(repo_id: str, config: dict, model_name: str, samp
         
         if widgets:
             frontmatter["widget"] = widgets
+
+        inference_params = {}
+        sample_width = sample_config.get("width") if isinstance(sample_config, dict) else None
+        sample_height = sample_config.get("height") if isinstance(sample_config, dict) else None
+        if sample_width:
+            inference_params["width"] = sample_width
+        if sample_height:
+            inference_params["height"] = sample_height
+        if inference_params:
+            frontmatter["inference"] = {"parameters": inference_params}
         
         if trigger_word:
             frontmatter["instance_prompt"] = trigger_word
@@ -623,7 +674,7 @@ def generate_model_card_readme(repo_id: str, config: dict, model_name: str, samp
 
 Model trained with [AI Toolkit by Ostris](https://github.com/ostris/ai-toolkit)
 
-<Gallery />
+{gallery_section}
 
 ## Trigger words
 

ui/src/app/datasets/page.tsx

@@ -22,7 +22,7 @@ export default function Datasets() {
   const { datasets, status, refreshDatasets } = useDatasetList();
   const [newDatasetName, setNewDatasetName] = useState('');
   const [isNewDatasetModalOpen, setIsNewDatasetModalOpen] = useState(false);
-  const { status: authStatus } = useAuth();
+  const { status: authStatus, namespace } = useAuth();
   const isAuthenticated = authStatus === 'authenticated';
 
   // Transform datasets array into rows with objects
@@ -85,7 +85,9 @@ export default function Datasets() {
       return;
     }
     try {
-      const data = await apiClient.post('/api/datasets/create', { name: newDatasetName }).then(res => res.data);
+      const data = await apiClient
+        .post('/api/datasets/create', { name: newDatasetName, namespace })
+        .then(res => res.data);
       console.log('New dataset created:', data);
       if (usingBrowserDb && data?.name) {
         addUserDataset(data.name, data?.path || '');
@@ -117,7 +119,9 @@ export default function Datasets() {
           return;
         }
         try {
-          const data = await apiClient.post('/api/datasets/create', { name }).then(res => res.data);
+          const data = await apiClient
+            .post('/api/datasets/create', { name, namespace })
+            .then(res => res.data);
           console.log('New dataset created:', data);
           if (usingBrowserDb && data?.name) {
             addUserDataset(data.name, data?.path || '');

ui/src/app/jobs/new/jobConfig.ts

@@ -74,7 +74,7 @@ export const defaultJobConfig: JobConfig = {
             use_ema: false,
             ema_decay: 0.99,
           },
-          skip_first_sample: false,
+          skip_first_sample: true,
           disable_sampling: false,
           dtype: 'bf16',
           diff_output_preservation: false,
@@ -94,7 +94,7 @@ export const defaultJobConfig: JobConfig = {
         },
         sample: {
           sampler: 'flowmatch',
-          sample_every: 250,
+          sample_every: 1500,
           width: 1024,
           height: 1024,
           samples: [

ui/src/app/jobs/new/page.tsx

@@ -31,7 +31,7 @@ export default function TrainingForm() {
   const router = useRouter();
   const searchParams = useSearchParams();
   const runId = searchParams.get('id');
-  const { status: authStatus } = useAuth();
+  const { status: authStatus, namespace: authNamespace } = useAuth();
   const isAuthenticated = authStatus === 'authenticated';
   const [gpuIDs, setGpuIDs] = useState<string | null>(null);
   const { settings, isSettingsLoaded } = useSettings();
@@ -67,7 +67,7 @@ export default function TrainingForm() {
             } else {
               try {
                 const response = await apiClient
-                  .post('/api/datasets/create', { name })
+                  .post('/api/datasets/create', { name, namespace: authNamespace })
                   .then(res => res.data);
                 if (response?.path) {
                   datasetPath = response.path;