dataset fix and tag fix

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

@@ -598,28 +598,55 @@ def generate_model_card_readme(repo_id: str, config: dict, model_name: str, cura
         tags = []
         lower_arch = (arch or "").lower()
         lower_model_name = (model_config.get("name_or_path", "") or "").lower()
+        base_model_lower = (base_model or "").lower()
 
         instruction_arches = {'flux_kontext', 'hidream_e1', 'qwen_image_edit'}
         is_instruction = lower_arch in instruction_arches or 'kontext' in lower_model_name
+        is_edit_model = (
+            'edit' in lower_arch
+            or 'image_edit' in lower_arch
+            or 'edit' in lower_model_name
+            or 'image-edit' in lower_model_name
+            or 'edit' in base_model_lower
+            or 'image-edit' in base_model_lower
+        )
 
         datasets_config = config.get('config', {}).get('process', [{}])[0].get('datasets', [])
-        is_wan_video = lower_arch.startswith('wan2') or 'wan2' in lower_model_name
-        is_video = (
-            'video' in lower_arch
-            or process_config.get('type') == 'video'
-            or any(isinstance(dataset, dict) and dataset.get('do_i2v') for dataset in datasets_config)
-            or is_wan_video
+        sample_config = process_config.get('sample', {}) if isinstance(process_config, dict) else {}
+        sample_num_frames = sample_config.get('num_frames') if isinstance(sample_config, dict) else None
+
+        video_indicators = [
+            'video' in lower_arch,
+            'video' in lower_model_name,
+            'video' in base_model_lower,
+            't2v' in lower_arch,
+            't2v' in lower_model_name,
+            't2v' in base_model_lower,
+            'i2v' in lower_arch,
+            'i2v' in lower_model_name,
+            'i2v' in base_model_lower,
+            process_config.get('type') == 'video' if isinstance(process_config, dict) else False,
+            any(isinstance(dataset, dict) and dataset.get('do_i2v') for dataset in datasets_config),
+            isinstance(sample_num_frames, (int, float)) and sample_num_frames and sample_num_frames > 1,
+            lower_arch.startswith('wan'),
+            'wan' in lower_model_name,
+            'wan' in base_model_lower,
+        ]
+        is_video = any(video_indicators)
+        is_i2v = any(
+            indicator
+            for indicator in [
+                'i2v' in lower_arch,
+                'i2v' in lower_model_name,
+                'i2v' in base_model_lower,
+                any(isinstance(dataset, dict) and dataset.get('do_i2v') for dataset in datasets_config),
+            ]
         )
 
-        if lower_arch.startswith('qwen_image_edit') or 'image-edit' in lower_model_name:
-            tags.append("image-to-image")
-        elif is_instruction:
+        if is_instruction or is_edit_model:
             tags.append("image-to-image")
         elif is_video:
-            if is_wan_video and ('i2v' in model_name.lower() or 'i2v' in lower_arch or 'i2v' in lower_model_name):
-                tags.append("image-to-video")
-            else:
-                tags.append("text-to-video")
+            tags.append("image-to-video" if is_i2v else "text-to-video")
         else:
             tags.append("text-to-image")
 

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

@@ -13,6 +13,7 @@ import { Button } from '@headlessui/react';
 import { ChevronDown, ChevronUp, Upload } from 'lucide-react';
 import { apiClient } from '@/utils/api';
 import { addUserDataset, updateUserDatasetPath } from '@/utils/storage/datasetStorage';
+import DatasetImageCard from '@/components/DatasetImageCard';
 
 type DatasetMode = 'upload' | 'existing';
 
@@ -53,6 +54,12 @@ const buildDatasetName = (base: string, suffix: string) => {
   return `${slug}${suffix}`;
 };
 
+const extractFolderName = (value?: string | null) => {
+  if (!value) return '';
+  const parts = value.split(/[\\/]+/).filter(Boolean);
+  return parts.length ? parts[parts.length - 1] : '';
+};
+
 type SimplifiedJobProps = {
   jobConfig: JobConfig;
   setJobConfig: (value: any, key: string) => void;
@@ -135,6 +142,8 @@ export default function SimplifiedJob({
   const [trainUploading, setTrainUploading] = useState(false);
   const [trainUploadInfo, setTrainUploadInfo] = useState<string | null>(null);
   const [trainUploadError, setTrainUploadError] = useState<string | null>(null);
+  const [trainImages, setTrainImages] = useState<{ img_path: string }[]>([]);
+  const [trainImagesStatus, setTrainImagesStatus] = useState<'idle' | 'loading' | 'success' | 'error'>('idle');
 
   const [controlDatasetMode, setControlDatasetMode] = useState<DatasetMode>('upload');
   const [controlModeTouched, setControlModeTouched] = useState(hasCustomControlPath);
@@ -152,6 +161,8 @@ export default function SimplifiedJob({
   const [controlUploading, setControlUploading] = useState(false);
   const [controlUploadInfo, setControlUploadInfo] = useState<string | null>(null);
   const [controlUploadError, setControlUploadError] = useState<string | null>(null);
+  const [controlImages, setControlImages] = useState<{ img_path: string }[]>([]);
+  const [controlImagesStatus, setControlImagesStatus] = useState<'idle' | 'loading' | 'success' | 'error'>('idle');
 
   const modelArch = useMemo(() => {
     return modelArchs.find(arch => arch.name === process.model.arch);
@@ -486,6 +497,80 @@ export default function SimplifiedJob({
     return datasetOptions;
   }, [datasetOptions]);
 
+  const trainDatasetIdentifier = useMemo(() => {
+    return (
+      trainDatasetResolvedName || extractFolderName(trainDatasetPath) || trainDatasetName || ''
+    ).trim();
+  }, [trainDatasetName, trainDatasetPath, trainDatasetResolvedName]);
+
+  const controlDatasetIdentifier = useMemo(() => {
+    if (!requiresControlDataset) {
+      return '';
+    }
+    return (
+      controlDatasetResolvedName || extractFolderName(controlDatasetPath) || controlDatasetName || ''
+    ).trim();
+  }, [controlDatasetName, controlDatasetPath, controlDatasetResolvedName, requiresControlDataset]);
+
+  const canLoadTrainImages = useMemo(
+    () => Boolean(trainDatasetIdentifier) && (trainDatasetMode === 'existing' || Boolean(trainDatasetResolvedName)),
+    [trainDatasetIdentifier, trainDatasetMode, trainDatasetResolvedName],
+  );
+
+  const canLoadControlImages = useMemo(
+    () =>
+      requiresControlDataset &&
+      Boolean(controlDatasetIdentifier) &&
+      (controlDatasetMode === 'existing' || Boolean(controlDatasetResolvedName)),
+    [controlDatasetIdentifier, controlDatasetMode, controlDatasetResolvedName, requiresControlDataset],
+  );
+
+  const refreshTrainDatasetImages = useCallback(() => {
+    if (!canLoadTrainImages) {
+      setTrainImages([]);
+      setTrainImagesStatus('idle');
+      return;
+    }
+    setTrainImagesStatus('loading');
+    apiClient
+      .post('/api/datasets/listImages', { datasetName: trainDatasetIdentifier })
+      .then(res => {
+        setTrainImages(res.data?.images || []);
+        setTrainImagesStatus('success');
+      })
+      .catch(error => {
+        console.error('Failed to load training dataset images:', error);
+        setTrainImagesStatus('error');
+      });
+  }, [canLoadTrainImages, trainDatasetIdentifier]);
+
+  const refreshControlDatasetImages = useCallback(() => {
+    if (!canLoadControlImages) {
+      setControlImages([]);
+      setControlImagesStatus('idle');
+      return;
+    }
+    setControlImagesStatus('loading');
+    apiClient
+      .post('/api/datasets/listImages', { datasetName: controlDatasetIdentifier })
+      .then(res => {
+        setControlImages(res.data?.images || []);
+        setControlImagesStatus('success');
+      })
+      .catch(error => {
+        console.error('Failed to load control dataset images:', error);
+        setControlImagesStatus('error');
+      });
+  }, [canLoadControlImages, controlDatasetIdentifier]);
+
+  useEffect(() => {
+    refreshTrainDatasetImages();
+  }, [refreshTrainDatasetImages]);
+
+  useEffect(() => {
+    refreshControlDatasetImages();
+  }, [refreshControlDatasetImages]);
+
   return (
     <form onSubmit={handleSubmit} className="space-y-8">
       <Card title="Training Setup">
@@ -685,6 +770,44 @@ export default function SimplifiedJob({
           </div>
         )}
 
+        {canLoadTrainImages && (
+          <div className="mt-6">
+            <div className="flex items-center justify-between mb-3">
+              <h3 className="text-sm font-semibold text-gray-200">Training Dataset Files</h3>
+              <button
+                type="button"
+                className="text-xs text-blue-400 hover:text-blue-300"
+                onClick={refreshTrainDatasetImages}
+              >
+                Refresh
+              </button>
+            </div>
+            {trainImagesStatus === 'loading' && (
+              <p className="text-xs text-gray-400">Loading dataset preview…</p>
+            )}
+            {trainImagesStatus === 'error' && (
+              <p className="text-xs text-red-400">
+                Unable to load files for this dataset. Ensure the folder exists and contains images or videos.
+              </p>
+            )}
+            {trainImagesStatus === 'success' && trainImages.length === 0 && (
+              <p className="text-xs text-gray-400">No files detected yet. Upload images to see them here.</p>
+            )}
+            {trainImagesStatus === 'success' && trainImages.length > 0 && (
+              <div className="grid grid-cols-1 sm:grid-cols-2 md:grid-cols-3 gap-4">
+                {trainImages.map(image => (
+                  <DatasetImageCard
+                    key={image.img_path}
+                    imageUrl={image.img_path}
+                    alt="Training dataset file"
+                    onDelete={refreshTrainDatasetImages}
+                  />
+                ))}
+              </div>
+            )}
+          </div>
+        )}
+
         <div className="grid grid-cols-1 md:grid-cols-2 gap-6 mt-6">
           <NumberInput
             label="Dataset Weight"
@@ -771,6 +894,44 @@ export default function SimplifiedJob({
                 />
               </div>
             )}
+
+            {canLoadControlImages && (
+              <div className="mt-6">
+                <div className="flex items-center justify-between mb-3">
+                  <h3 className="text-sm font-semibold text-gray-200">Control Dataset Files</h3>
+                  <button
+                    type="button"
+                    className="text-xs text-blue-400 hover:text-blue-300"
+                    onClick={refreshControlDatasetImages}
+                  >
+                    Refresh
+                  </button>
+                </div>
+                {controlImagesStatus === 'loading' && (
+                  <p className="text-xs text-gray-400">Loading control dataset preview…</p>
+                )}
+                {controlImagesStatus === 'error' && (
+                  <p className="text-xs text-red-400">
+                    Unable to load files for this control dataset. Ensure the folder exists and has matching filenames.
+                  </p>
+                )}
+                {controlImagesStatus === 'success' && controlImages.length === 0 && (
+                  <p className="text-xs text-gray-400">No files detected yet. Upload matching control images.</p>
+                )}
+                {controlImagesStatus === 'success' && controlImages.length > 0 && (
+                  <div className="grid grid-cols-1 sm:grid-cols-2 md:grid-cols-3 gap-4">
+                    {controlImages.map(image => (
+                      <DatasetImageCard
+                        key={image.img_path}
+                        imageUrl={image.img_path}
+                        alt="Control dataset file"
+                        onDelete={refreshControlDatasetImages}
+                      />
+                    ))}
+                  </div>
+                )}
+              </div>
+            )}
           </div>
         )}
       </Card>