default text-to-image

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")