Merge pull request #7 from livepeer/ryan/feat/model-detection

ryanontheinside · web-flow · commit a1f2c692f5c7 · 2025-06-30T18:37:43.000-04:00
Ryan/feat/model detection
diff --git a/README_CONFIG.md b/README_CONFIG.md
@@ -59,7 +59,6 @@ output_image = wrapper(input_image)
 ```yaml
 # Model and pipeline
 model_id: "stabilityai/sd-turbo"
-pipeline_type: "sdturbo"
 
 # Generation mode
 mode: "txt2img"  # or "img2img"
@@ -118,7 +117,6 @@ controlnets:
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `model_id` | str | *required* | Model ID or local path |
-| `pipeline_type` | str | `"sd1.5"` | Pipeline type: `"sd1.5"`, `"sdturbo"`, `"sdxlturbo"` |
 | `mode` | str | `"img2img"` | Generation mode: `"img2img"` or `"txt2img"` |
 | `width` | int | `512` | Output image width |
 | `height` | int | `512` | Output image height |
diff --git a/README_CONTROLNET.md b/README_CONTROLNET.md
@@ -141,7 +141,6 @@ config = {
                 'high_threshold': 200
             },
             'enabled': True,
-            'pipeline_type': 'sd1.5',
             'control_guidance_start': 0.0,
             'control_guidance_end': 1.0,
         }
@@ -173,7 +172,6 @@ controlnet_configs = [
             'high_threshold': 200
         },
         'enabled': True,
-        'pipeline_type': 'sd1.5',
         'control_guidance_start': 0.0,
         'control_guidance_end': 1.0,
     }
@@ -226,7 +224,6 @@ output = wrapper(input_image)
 ```yaml
 # Basic SD-Turbo + Canny ControlNet
 model_id: "stabilityai/sd-turbo"
-pipeline_type: "sdturbo"
 mode: "img2img"
 prompt: "a beautiful landscape, highly detailed"
 
@@ -245,7 +242,6 @@ controlnets:
 ```yaml
 # Multiple ControlNets with different strengths
 model_id: "stabilityai/sd-turbo"
-pipeline_type: "sdturbo"
 prompt: "a person in a detailed environment"
 
 controlnets:
@@ -363,7 +359,6 @@ controlnet_config = {
     'preprocessor': str,               # Preprocessor name
     'preprocessor_params': dict,       # Preprocessor-specific parameters
     'enabled': bool,                   # Enable/disable this ControlNet
-    'pipeline_type': str,              # Pipeline type (sd1.5, sdturbo, soon sdxlturbo)
     'control_guidance_start': float,   # Guidance start timestep (0.0-1.0)
     'control_guidance_end': float,     # Guidance end timestep (0.0-1.0)
 }
diff --git a/configs/sd15_canny_depth.yaml.example b/configs/sd15_canny_depth.yaml.example
@@ -16,7 +16,6 @@ frame_buffer_size: 1
 delta: 0.7
 
 # Advanced parameters
-pipeline_type: "sdturbo"
 use_lcm_lora: false
 use_tiny_vae: true
 acceleration: "tensorrt"
diff --git a/configs/sd15_depth_trt_example.yaml.example b/configs/sd15_depth_trt_example.yaml.example
@@ -16,7 +16,6 @@ delta: 0.7
 frame_buffer_size: 1
 
 # Advanced parameters
-pipeline_type: "sd1.5"
 use_lcm_lora: true
 use_tiny_vae: true
 acceleration: "xformers"
diff --git a/configs/sd15_tile.yaml.example b/configs/sd15_tile.yaml.example
@@ -12,7 +12,6 @@ guidance_scale: 1.1
 num_inference_steps: 50
 
 # Advanced parameters
-pipeline_type: "sd1.5"
 use_lcm_lora: true
 use_tiny_vae: true
 acceleration: "xformers"
diff --git a/configs/sdturbo_color.yaml.example b/configs/sdturbo_color.yaml.example
@@ -17,7 +17,6 @@ use_denoising_batch: true
 delta: 0.7
 
 # Advanced parameters
-pipeline_type: "sdturbo"
 use_lcm_lora: true
 use_tiny_vae: true
 acceleration: "tensorrt"
diff --git a/configs/sdturbo_mediapipe_pose_depth_trt.yaml.example b/configs/sdturbo_mediapipe_pose_depth_trt.yaml.example
@@ -17,7 +17,6 @@ use_denoising_batch: true
 delta: 0.7
 
 # Advanced parameters
-pipeline_type: "sdturbo"
 use_lcm_lora: true
 use_tiny_vae: true
 acceleration: "tensorrt"
diff --git a/configs/sdturbo_yolonas_pose_depth_trt.yaml.example b/configs/sdturbo_yolonas_pose_depth_trt.yaml.example
@@ -17,7 +17,6 @@ use_denoising_batch: true
 delta: 0.7
 
 # Advanced parameters
-pipeline_type: "sdturbo"
 use_lcm_lora: true
 use_tiny_vae: true
 acceleration: "tensorrt"
diff --git a/configs/sdxlturbo_canny.yaml.example b/configs/sdxlturbo_canny.yaml.example
@@ -1,5 +1,4 @@
 model_id: "stabilityai/sdxl-turbo"
-pipeline_type: "sdxlturbo"
 t_index_list: [32]  # Controls denoising strength - lower values = less denoising, higher = more denoising
 width: 1024
 height: 1024
@@ -18,7 +17,6 @@ use_taesd: true  # Use Tiny AutoEncoder XL for faster decoding
 safety_checker: false
 
 # Advanced parameters
-pipeline_type: "sdxlturbo"
 use_lcm_lora: false  # SD-XL Turbo already has optimizations built-in
 use_tiny_vae: true
 acceleration: "none"  # Can be "tensorrt" for additional speed
diff --git a/demo/realtime-img2img/frontend/src/lib/components/PreprocessorDocs.svelte b/demo/realtime-img2img/frontend/src/lib/components/PreprocessorDocs.svelte
@@ -40,7 +40,6 @@
       use_denoising_batch: true,
       delta: 0.7,
       frame_buffer_size: 1,
-      pipeline_type: "sd1.5",
       use_lcm_lora: true,
       use_tiny_vae: true,
       acceleration: "xformers",
diff --git a/demo/realtime-img2img/main.py b/demo/realtime-img2img/main.py
@@ -686,7 +686,6 @@ async def get_preprocessors_info():
                 "use_denoising_batch": True,
                 "delta": 0.7,
                 "frame_buffer_size": 1,
-                "pipeline_type": "sd1.5",
                 "use_lcm_lora": True,
                 "use_tiny_vae": True,
                 "acceleration": "xformers",
diff --git a/examples/controlnet/controlnet_video_test.py b/examples/controlnet/controlnet_video_test.py
@@ -21,7 +21,7 @@
 from collections import deque
 
 
-def process_video(config_path, input_video, output_dir, resolution=None, engine_only=False):
+def process_video(config_path, input_video, output_dir, engine_only=False):
     """Process video through ControlNet pipeline"""
     print(f"process_video: Loading config from {config_path}")
     
@@ -32,15 +32,11 @@ def process_video(config_path, input_video, output_dir, resolution=None, engine_
     # Load configuration
     config = load_config(config_path)
     
-    # Set default resolution based on pipeline type if not specified
-    if resolution is None:
-        pipeline_type = config.get('pipeline_type', 'sd1.5')
-        if pipeline_type == 'sdxlturbo':
-            resolution = 1024
-        else:
-            resolution = 512
+    # Get width and height from config (with defaults)
+    width = config.get('width', 512)
+    height = config.get('height', 512)
     
-    print(f"process_video: Using resolution: {resolution}x{resolution}")
+    print(f"process_video: Using dimensions: {width}x{height}")
     
     # Create output directory
     output_dir = Path(output_dir)
@@ -55,8 +51,8 @@ def process_video(config_path, input_video, output_dir, resolution=None, engine_
     shutil.copy2(input_video, input_copy_path)
     print(f"process_video: Copied input video to {input_copy_path}")
     
-    # Create wrapper using the built-in function
-    wrapper = create_wrapper_from_config(config, width=resolution, height=resolution)
+    # Create wrapper using the built-in function (width/height from config)
+    wrapper = create_wrapper_from_config(config)
     
     if engine_only:
         print("Engine-only mode: TensorRT engines have been built (if needed). Exiting.")
@@ -76,7 +72,7 @@ def process_video(config_path, input_video, output_dir, resolution=None, engine_
     # Setup output video writer
     output_video_path = output_dir / "output_video.mp4"
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(str(output_video_path), fourcc, fps, (resolution * 2, resolution))
+    out = cv2.VideoWriter(str(output_video_path), fourcc, fps, (width + width, height))
     
     # Performance tracking
     frame_times = []
@@ -93,7 +89,7 @@ def process_video(config_path, input_video, output_dir, resolution=None, engine_
         frame_start_time = time.time()
         
         # Resize frame
-        frame_resized = cv2.resize(frame, (resolution, resolution))
+        frame_resized = cv2.resize(frame, (width, height))
         
         # Convert frame to PIL
         frame_rgb = cv2.cvtColor(frame_resized, cv2.COLOR_BGR2RGB)
@@ -112,7 +108,7 @@ def process_video(config_path, input_video, output_dir, resolution=None, engine_
         
         # Add labels
         cv2.putText(combined, "Input", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
-        cv2.putText(combined, "Generated", (resolution + 10, 30), 
+        cv2.putText(combined, "Generated", (width + 10, 30), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
         
         # Write frame
@@ -148,7 +144,8 @@ def process_video(config_path, input_video, output_dir, resolution=None, engine_
     metrics = {
         "input_video": str(input_video),
         "config_file": str(config_path),
-        "resolution": f"{resolution}x{resolution}",
+        "width": width,
+        "height": height,
         "total_frames": frame_idx,
         "total_time_seconds": total_time,
         "avg_fps": avg_fps,
@@ -157,7 +154,6 @@ def process_video(config_path, input_video, output_dir, resolution=None, engine_
         "avg_frame_time_seconds": avg_frame_time,
         "min_frame_time_seconds": min_frame_time,
         "max_frame_time_seconds": max_frame_time,
-        "pipeline_type": config.get('pipeline_type', 'sd1.5'),
         "model_id": config['model_id'],
         "acceleration": config.get('acceleration', 'none'),
         "frame_buffer_size": config.get('frame_buffer_size', 1),
@@ -179,7 +175,6 @@ def process_video(config_path, input_video, output_dir, resolution=None, engine_
     
     return metrics
 
-
 def main():
     parser = argparse.ArgumentParser(description="ControlNet Video Test Demo")
     
@@ -193,8 +188,6 @@ def main():
                        help="Path to input video file")
     parser.add_argument("--output-dir", type=str, default=None,
                        help="Output directory for results (default: creates timestamped directory)")
-    parser.add_argument("--resolution", type=int, default=None,
-                       help="Video resolution (auto-detects from pipeline type if not specified)")
     parser.add_argument("--engine-only", action="store_true", help="Only build TensorRT engines and exit (no video processing)")
     
     args = parser.parse_args()
@@ -222,7 +215,7 @@ def main():
     print(f"main: Output directory: {args.output_dir}")
     
     try:
-        metrics = process_video(args.config, args.input_video, args.output_dir, args.resolution, engine_only=args.engine_only)
+        metrics = process_video(args.config, args.input_video, args.output_dir, engine_only=args.engine_only)
         if args.engine_only:
             print("main: Engine-only mode completed successfully!")
             return 0
diff --git a/src/streamdiffusion/acceleration/tensorrt/controlnet_models.py b/src/streamdiffusion/acceleration/tensorrt/controlnet_models.py
@@ -9,7 +9,6 @@ class ControlNetTRT(BaseModel):
     """TensorRT model definition for ControlNet compilation"""
     
     def __init__(self, 
-                 controlnet_type: str = "canny",
                  fp16: bool = True,
                  device: str = "cuda",
                  max_batch: int = 2,
@@ -25,9 +24,8 @@ def __init__(self,
             embedding_dim=embedding_dim,
             **kwargs
         )
-        self.controlnet_type = controlnet_type
         self.unet_dim = unet_dim
-        self.name = f"ControlNet-{controlnet_type}"
+        self.name = "ControlNet"
         
     def get_input_names(self) -> List[str]:
         """Get input names for ControlNet TensorRT engine"""
@@ -116,7 +114,7 @@ class ControlNetSDXLTRT(ControlNetTRT):
     def __init__(self, **kwargs):
         kwargs.setdefault('embedding_dim', 2048)
         super().__init__(**kwargs)
-        self.name = f"ControlNet-SDXL-{self.controlnet_type}"
+        self.name = "ControlNet-SDXL"
     
     def get_input_names(self) -> List[str]:
         """SDXL ControlNet has additional conditioning inputs"""
@@ -171,10 +169,9 @@ def get_sample_input(self, batch_size, image_height, image_width):
 
 
 def create_controlnet_model(model_type: str = "sd15", 
-                           controlnet_type: str = "canny",
                            **kwargs) -> ControlNetTRT:
     """Factory function to create appropriate ControlNet TensorRT model"""
     if model_type.lower() in ["sdxl", "sdxl-turbo"]:
-        return ControlNetSDXLTRT(controlnet_type=controlnet_type, **kwargs)
+        return ControlNetSDXLTRT(**kwargs)
     else:
-        return ControlNetTRT(controlnet_type=controlnet_type, **kwargs) 
+        return ControlNetTRT(**kwargs) 
diff --git a/src/streamdiffusion/acceleration/tensorrt/engine_pool.py b/src/streamdiffusion/acceleration/tensorrt/engine_pool.py
@@ -74,7 +74,6 @@ def _discover_existing_engines(self) -> None:
     def get_or_load_engine(self, 
                           model_id: str,
                           pytorch_model: Any,
-                          controlnet_type: str = "canny",
                           model_type: str = "sd15",
                           batch_size: int = 1) -> HybridControlNet:
         """Get or load ControlNet engine with TensorRT/PyTorch fallback"""
@@ -97,7 +96,7 @@ def get_or_load_engine(self,
                 print(f"ControlNetEnginePool.get_or_load_engine: Architecture detection failed: {e}, using provided type: {model_type}")
             
             success = self._compile_controlnet(
-                pytorch_model, controlnet_type, model_type, str(engine_path), batch_size
+                pytorch_model, model_type, str(engine_path), batch_size
             )
             
             compilation_time = time.time() - compilation_start
@@ -122,13 +121,12 @@ def get_or_load_engine(self,
     
     def _compile_controlnet(self, 
                            pytorch_model: Any,
-                           controlnet_type: str, 
                            model_type: str,
                            engine_path: str,
                            batch_size: int) -> bool:
         """Compile ControlNet to TensorRT"""
         try:
-            print(f"ControlNetEnginePool._compile_controlnet: Starting ControlNet compilation: {controlnet_type} ({model_type}) for {self.image_width}x{self.image_height}")
+            print(f"ControlNetEnginePool._compile_controlnet: Starting ControlNet compilation: {model_type} for {self.image_width}x{self.image_height}")
             
             # Use actual image dimensions instead of hardcoded 512x512
             height = self.image_height
@@ -143,7 +141,6 @@ def _compile_controlnet(self,
             
             controlnet_model = create_controlnet_model(
                 model_type=model_type,
-                controlnet_type=controlnet_type,
                 max_batch=batch_size,
                 min_batch_size=1,
                 embedding_dim=embedding_dim
diff --git a/src/streamdiffusion/config.py b/src/streamdiffusion/config.py
@@ -108,7 +108,6 @@ def _extract_prepare_params(config: Dict[str, Any]) -> Dict[str, Any]:
 def _prepare_controlnet_configs(config: Dict[str, Any]) -> List[Dict[str, Any]]:
     """Prepare ControlNet configurations for wrapper"""
     controlnet_configs = []
-    pipeline_type = config.get('pipeline_type', 'sd1.5')
 
     for cn_config in config['controlnets']:
         controlnet_config = {
@@ -117,7 +116,6 @@ def _prepare_controlnet_configs(config: Dict[str, Any]) -> List[Dict[str, Any]]:
             'conditioning_scale': cn_config.get('conditioning_scale', 1.0),
             'enabled': cn_config.get('enabled', True),
             'preprocessor_params': cn_config.get('preprocessor_params'),
-            'pipeline_type': pipeline_type,
             'control_guidance_start': cn_config.get('control_guidance_start', 0.0),
             'control_guidance_end': cn_config.get('control_guidance_end', 1.0),
         }
diff --git a/src/streamdiffusion/controlnet/base_controlnet_pipeline.py b/src/streamdiffusion/controlnet/base_controlnet_pipeline.py
diff --git a/src/streamdiffusion/wrapper.py b/src/streamdiffusion/wrapper.py
