Fix pydantic 2 errors

python/cog/base_input.py

@@ -1,19 +1,22 @@
 from pathlib import Path
 
+import pydantic
 from pydantic import BaseModel
 
-from .types import (
-    URLPath,
-)
+from .types import PYDANTIC_V2, URLPath
 
 
 # Base class for inputs, constructed dynamically in get_input_type().
 # (This can't be a docstring or it gets passed through to the schema.)
 class BaseInput(BaseModel):
-    class Config:
-        # When using `choices`, the type is converted into an enum to validate
-        # But, after validation, we want to pass the actual value to predict(), not the enum object
-        use_enum_values = True
+    if PYDANTIC_V2:
+        model_config = pydantic.ConfigDict(use_enum_values=True)  # type: ignore
+    else:
+
+        class Config:
+            # When using `choices`, the type is converted into an enum to validate
+            # But, after validation, we want to pass the actual value to predict(), not the enum object
+            use_enum_values = True
 
     def cleanup(self) -> None:
         """

python/cog/server/http.py

@@ -29,16 +29,7 @@
 from ..json import upload_files
 from ..logging import setup_logging
 from ..mode import Mode
-from ..predictor import (
-    get_input_type,
-    get_output_type,
-    get_predictor_ref,
-    get_training_input_type,
-    get_training_output_type,
-    load_config,
-    load_slim_predictor_from_ref,
-)
-from ..types import PYDANTIC_V2, CogConfig
+from ..types import PYDANTIC_V2
 
 if PYDANTIC_V2:
     from .helpers import (