diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 0432a655..b0ba5a98 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -45,8 +45,14 @@ jobs:
       uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
+    - name: Upgrade pip
+      run: python -m pip install --upgrade pip==23.2
     - name: Install LIT package with testing dependencies
       run: python -m pip install -e '.[test]'
+    - name: Debug dependency tree
+      run: |
+        python -m pip install pipdeptree
+        pipdeptree | grep decorator -A 5 || true
     - name: Test LIT
       run: pytest -v
     - name: Setup Node ${{ matrix.node-version }}
diff --git a/lit_nlp/api/components.py b/lit_nlp/api/components.py
index b1cdcff3..93de3f2e 100644
--- a/lit_nlp/api/components.py
+++ b/lit_nlp/api/components.py
@@ -137,7 +137,7 @@ def is_compatible(self, model: lit_model.Model,
     """True if the model and dataset support metric computation."""
     for pred_spec in model.output_spec().values():
       parent_key: Optional[str] = getattr(pred_spec, 'parent', None)
-      parent_spec: Optional[types.LitType] = dataset.spec().get(parent_key)
+      parent_spec: Optional[types.LitType] = dataset.spec().get(parent_key)  # pyrefly: ignore[bad-argument-type]
       if self.is_field_compatible(pred_spec, parent_spec):
         return True
     return False
diff --git a/lit_nlp/api/dataset.py b/lit_nlp/api/dataset.py
index eb315891..9d3b75f7 100644
--- a/lit_nlp/api/dataset.py
+++ b/lit_nlp/api/dataset.py
@@ -236,7 +236,7 @@ def remap(self, field_map: Mapping[str, str]):
     """Return a copy of this dataset with some fields renamed."""
     new_spec = utils.remap_dict(self.spec(), field_map)
     new_examples = [utils.remap_dict(ex, field_map) for ex in self.examples]
-    return Dataset(new_spec, new_examples, base=self)
+    return Dataset(new_spec, new_examples, base=self)  # pyrefly: ignore[bad-argument-type]
 
   @staticmethod
   def lit_example_from_bytes(input_bytes: bytes) -> Optional[JsonDict]:
@@ -273,7 +273,7 @@ def index_inputs(
       indexed.append(
           IndexedInput(
               data=types.MappingProxyType(
-                  example | {INPUT_ID_FIELD: ex_id, INPUT_META_FIELD: ex_meta}
+                  example | {INPUT_ID_FIELD: ex_id, INPUT_META_FIELD: ex_meta}  # pyrefly: ignore[unsupported-operation]
               ),
               id=ex_id,
               meta=ex_meta,
@@ -376,6 +376,7 @@ def load(self, path: str):
       new_dataset = base_dataset.load(path) if base_dataset else None
 
       if new_dataset is not None:
+        # pyrefly: ignore[missing-attribute]
         description = (f'{len(new_dataset)} examples from '
                        f'{path}\n{self._base.description()}')
         return IndexedDataset(
@@ -438,7 +439,7 @@ def load_lit_format(
         **kw,
     )
   else:
-    return Dataset(spec=spec, examples=examples, *args, **kw)
+    return Dataset(spec=spec, examples=examples, *args, **kw)  # pyrefly: ignore[bad-argument-type]
 
 
 # TODO(b/202210900): Remove "NoneDataset" once the LIT front-end constructs its
diff --git a/lit_nlp/api/layout.py b/lit_nlp/api/layout.py
index b39a2eb7..db56ae9e 100644
--- a/lit_nlp/api/layout.py
+++ b/lit_nlp/api/layout.py
@@ -70,7 +70,7 @@ def __call__(self, **kw):
 # LINT.IfChange
 # TODO(lit-dev): consider making modules subclass this instead of LitModuleName.
 @attr.s(auto_attribs=True)
-class ModuleConfig(dtypes.DataTuple):
+class ModuleConfig(dtypes.DataTuple):  # pyrefly: ignore[invalid-inheritance]
   module: Union[str, LitModuleName]
   requiredForTab: bool = False
   # TODO(b/172979677): support title, duplicateAsRow, numCols,
@@ -88,7 +88,7 @@ class ModuleConfig(dtypes.DataTuple):
 
 
 @attr.s(auto_attribs=True)
-class LayoutSettings(dtypes.DataTuple):
+class LayoutSettings(dtypes.DataTuple):  # pyrefly: ignore[invalid-inheritance]
   hideToolbar: bool = False
   mainHeight: int = 45
   leftWidth: int = 50
@@ -96,7 +96,7 @@ class LayoutSettings(dtypes.DataTuple):
 
 
 @attr.s(auto_attribs=True)
-class LitCanonicalLayout(dtypes.DataTuple):
+class LitCanonicalLayout(dtypes.DataTuple):  # pyrefly: ignore[invalid-inheritance]
   """Frontend UI layout; should match client/lib/types.ts."""
   upper: LitTabGroupLayout
   lower: LitTabGroupLayout = attr.ib(factory=dict)
diff --git a/lit_nlp/api/model.py b/lit_nlp/api/model.py
index 79558fce..32492b1e 100644
--- a/lit_nlp/api/model.py
+++ b/lit_nlp/api/model.py
@@ -178,12 +178,12 @@ def load(self, path: str):
   @abc.abstractmethod
   def input_spec(self) -> types.Spec:
     """Return a spec describing model inputs."""
-    return
+    return  # pyrefly: ignore[bad-return]
 
   @abc.abstractmethod
   def output_spec(self) -> types.Spec:
     """Return a spec describing model outputs."""
-    return
+    return  # pyrefly: ignore[bad-return]
 
   def get_embedding_table(self) -> tuple[list[str], np.ndarray]:
     """Return the full vocabulary and embedding table.
@@ -369,7 +369,7 @@ def predict_minibatch(self, inputs: list[JsonDict]) -> list[JsonDict]:
     Returns:
       list of outputs, following model.output_spec()
     """
-    return
+    return  # pyrefly: ignore[bad-return]
 
 
 class ProjectorModel(BatchedModel, metaclass=abc.ABCMeta):
diff --git a/lit_nlp/api/types.py b/lit_nlp/api/types.py
index 3e7dbfd7..3851e0d3 100644
--- a/lit_nlp/api/types.py
+++ b/lit_nlp/api/types.py
@@ -155,10 +155,10 @@ def from_json(d: JsonDict):
 
     base_cls = globals().get("LitType")
     cls = globals().get(type_name)  # class by name from this module
-    if cls is None or not issubclass(cls, base_cls):
+    if cls is None or not issubclass(cls, base_cls):  # pyrefly: ignore[bad-argument-type]
       raise NameError(f"{type_name} is not a valid LitType.")
 
-    return cls(**{k: d[k] for k in d if k != "__name__"})
+    return cls(**{k: d[k] for k in d if k != "__name__"})  # pyrefly: ignore[not-callable]
 
 Spec = dict[str, LitType]
 
@@ -174,7 +174,7 @@ def _remap_leaf(leaf: LitType, keymap: dict[str, str]) -> LitType:
       k: (keymap.get(v, v) if k in FIELD_REF_ATTRIBUTES else v)
       for k, v in d.items()
   }
-  return leaf.__class__(**d)
+  return leaf.__class__(**d)  # pyrefly: ignore[bad-argument-type]
 
 
 def remap_spec(spec: Spec, keymap: dict[str, str]) -> Spec:
@@ -201,7 +201,7 @@ class StringLitType(LitType):
   Mainly used for string inputs that have special formatting, and should only
   be edited manually.
   """
-  default: str = ""
+  default: str = ""  # pyrefly: ignore[bad-override]
 
   def validate_input(self, value, spec: Spec, example: Input):
     if not isinstance(value, str):
@@ -260,13 +260,13 @@ def validate_output(self, value, output_spec: Spec, output_dict: JsonDict,
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
 class ListLitType(LitType):
   """List type."""
-  default: Sequence[Any] = None
+  default: Sequence[Any] = None  # pyrefly: ignore[bad-assignment, bad-override]
 
 
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
 class _StringCandidateList(ListLitType):
   """A list of (text, score) tuples."""
-  default: ScoredTextCandidates = None
+  default: ScoredTextCandidates = None  # pyrefly: ignore[bad-assignment]
 
   def validate_output(self, value, output_spec: Spec, output_dict: JsonDict,
                       input_spec: Spec, dataset_spec: Spec,
@@ -373,9 +373,9 @@ class TokenTopKPreds(ListLitType):
 
   The inner list should contain (word, probability) in descending order.
   """
-  default: Sequence[ScoredTextCandidates] = None
+  default: Sequence[ScoredTextCandidates] = None  # pyrefly: ignore[bad-assignment]
 
-  align: str = None  # name of a Tokens field in the model output
+  align: str = None  # name of a Tokens field in the model output  # pyrefly: ignore[bad-assignment]
   parent: Optional[str] = None
 
   def _validate_scored_candidates(self, scored_candidates):
@@ -389,7 +389,7 @@ def _validate_scored_candidates(self, scored_candidates):
       if scored_candidate[1] is not None:
         if not isinstance(scored_candidate[1], NumericTypes):
           raise ValueError(f"{scored_candidate} second element is not a num")
-        if prev_val < scored_candidate[1]:
+        if prev_val < scored_candidate[1]:  # pyrefly: ignore[unsupported-operation]
           raise ValueError(
               "TokenTopKPreds candidates are not in descending order")
         else:
@@ -414,7 +414,7 @@ class Scalar(LitType):
   """Scalar value, a single float or int."""
   min_val: float = 0
   max_val: float = 1
-  default: float = 0
+  default: float = 0  # pyrefly: ignore[bad-override]
   step: float = .01
 
   def validate_input(self, value, spec: Spec, example: Input):
@@ -464,7 +464,7 @@ class TokenScores(_FloatList):
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
 class ReferenceScores(ListLitType):
   """Score of one or more target sequences."""
-  default: Sequence[float] = None
+  default: Sequence[float] = None  # pyrefly: ignore[bad-assignment]
 
   # name of a TextSegment or ReferenceTexts field in the input
   parent: Optional[str] = None
@@ -504,7 +504,7 @@ def validate_input(self, value, spec: Spec, example: Input):
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
 class _Tensor(LitType):
   """A tensor type."""
-  default: Sequence[float] = None
+  default: Sequence[float] = None  # pyrefly: ignore[bad-assignment, bad-override]
 
   def validate_input(self, value, spec: Spec, example: Input):
     if isinstance(value, list):
@@ -585,7 +585,7 @@ class SpanLabels(ListLitType):
   Span labels can cover more than one token, may not cover all tokens in the
   sentence, and may overlap with each other.
   """
-  default: Sequence[dtypes.SpanLabel] = None
+  default: Sequence[dtypes.SpanLabel] = None  # pyrefly: ignore[bad-assignment]
   align: str  # name of Tokens field
   parent: Optional[str] = None
 
@@ -609,7 +609,7 @@ class EdgeLabels(ListLitType):
   https://github.com/nyu-mll/jiant/tree/master/probing#data-format for more
   details.
   """
-  default: Sequence[dtypes.EdgeLabel] = None
+  default: Sequence[dtypes.EdgeLabel] = None  # pyrefly: ignore[bad-assignment]
   align: str  # name of Tokens field
 
   def validate_output(self, value, output_spec: Spec, output_dict: JsonDict,
@@ -636,7 +636,7 @@ class MultiSegmentAnnotations(ListLitType):
   TODO(lit-dev): by default, spans are treated as bytes in this context.
   Make this configurable, if some spans need to refer to tokens instead.
   """
-  default: Sequence[dtypes.AnnotationCluster] = None
+  default: Sequence[dtypes.AnnotationCluster] = None  # pyrefly: ignore[bad-assignment]
   exclusive: bool = False  # if true, treat as candidate list
   background: bool = False  # if true, don't emphasize in visualization
 
@@ -775,7 +775,7 @@ class SubwordOffsets(ListLitType):
 
   offsets[i] should be the index of the first wordpiece for input token i.
   """
-  default: Sequence[int] = None
+  default: Sequence[int] = None  # pyrefly: ignore[bad-assignment]
   align_in: str  # name of field in data spec
   align_out: str  # name of field in model output spec
 
@@ -793,7 +793,7 @@ class SparseMultilabelPreds(_StringCandidateList):
 
   The tuples are of the label and the score.
   """
-  default: ScoredTextCandidates = None
+  default: ScoredTextCandidates = None  # pyrefly: ignore[bad-assignment]
   vocab: Optional[Sequence[str]] = None  # label names
   parent: Optional[str] = None
 
@@ -816,7 +816,7 @@ class SingleFieldMatcher(FieldMatcher):
 
   UI will materialize this to a dropdown-list.
   """
-  default: str = None
+  default: str = None  # pyrefly: ignore[bad-assignment, bad-override]
 
 
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
@@ -826,7 +826,7 @@ class MultiFieldMatcher(FieldMatcher):
   UI will materialize this to multiple checkboxes. Use this when the user needs
   to pick more than one field in UI.
   """
-  default: Sequence[str] = []  # default names of selected items.
+  default: Sequence[str] = []  # default names of selected items.  # pyrefly: ignore[bad-override]
   select_all: bool = False  # Select all by default (overriddes default).
 
 
@@ -840,20 +840,20 @@ class Salience(LitType):
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
 class TokenSalience(Salience):
   """Metadata about a returned token salience map."""
-  default: dtypes.TokenSalience = None
+  default: dtypes.TokenSalience = None  # pyrefly: ignore[bad-assignment, bad-override]
 
 
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
 class FeatureSalience(Salience):
   """Metadata about a returned feature salience map."""
-  default: dtypes.FeatureSalience = None
+  default: dtypes.FeatureSalience = None  # pyrefly: ignore[bad-assignment, bad-override]
 
 
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
 class FrameSalience(Salience):
   """Metadata about a returned frame salience map."""
 
-  default: dtypes.FrameSalience = None
+  default: dtypes.FrameSalience = None  # pyrefly: ignore[bad-assignment, bad-override]
 
 
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
@@ -869,13 +869,13 @@ class ImageSalience(Salience):
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
 class SequenceSalience(Salience):
   """Metadata about a returned sequence salience map."""
-  default: dtypes.SequenceSalienceMap = None
+  default: dtypes.SequenceSalienceMap = None  # pyrefly: ignore[bad-assignment, bad-override]
 
 
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
 class BooleanLitType(LitType):
   """Boolean value."""
-  default: bool = False
+  default: bool = False  # pyrefly: ignore[bad-override]
 
   def validate_input(self, value, spec, example: Input):
     if not isinstance(value, bool):
@@ -916,7 +916,7 @@ class MetricBestValue(dtypes.EnumSerializableAsValues, enum.Enum):
 @attr.s(auto_attribs=True, frozen=True, kw_only=True)
 class MetricResult(LitType):
   """Score returned from the computation of a Metric."""
-  default: float = 0
+  default: float = 0  # pyrefly: ignore[bad-override]
   description: str = ""
   best_value: MetricBestValue = MetricBestValue.NONE
 
@@ -934,7 +934,7 @@ class SalienceTargetInfo(LitType):
 
   Value is a dict with keys 'field' (str) and 'index' (Optional[int]).
   """
-  default: Optional[Mapping[str, Any]] = None
+  default: Optional[Mapping[str, Any]] = None  # pyrefly: ignore[bad-override]
 
 
 # LINT.ThenChange(../client/lib/lit_types.ts)
diff --git a/lit_nlp/app.py b/lit_nlp/app.py
index 7d8afba5..2e773fae 100644
--- a/lit_nlp/app.py
+++ b/lit_nlp/app.py
@@ -111,7 +111,7 @@ def _build_metadata(self):
       }
 
       # List compatible datasets.
-      info['datasets'] = [
+      info['datasets'] = [  # pyrefly: ignore[bad-assignment]
           name for name, dataset in self._datasets.items()
           if model.is_compatible_with_dataset(dataset)
       ]
@@ -134,13 +134,13 @@ def _build_metadata(self):
             _get_compatible_names(self._metrics, model, dataset)
         )
 
-      info['generators'] = [
+      info['generators'] = [  # pyrefly: ignore[bad-assignment]
           name for name in self._generators.keys() if name in compat_gens
       ]
-      info['interpreters'] = [
+      info['interpreters'] = [  # pyrefly: ignore[bad-assignment]
           name for name in self._interpreters.keys() if name in compat_interps
       ]
-      info['metrics'] = [
+      info['metrics'] = [  # pyrefly: ignore[bad-assignment]
           name for name in self._metrics.keys() if name in compat_metrics
       ]
       model_info[name] = info
@@ -213,7 +213,7 @@ def _reconstitute_inputs(
         len(inputs),
         dataset_name,
     )
-    return [index[ex] if isinstance(ex, str) else ex for ex in inputs]
+    return [index[ex] if isinstance(ex, str) else ex for ex in inputs]  # pyrefly: ignore[bad-index]
 
   def _save_datapoints(
       self,
@@ -297,16 +297,16 @@ def _get_preds(self,
 
     # Figure out what to return to the frontend.
     output_spec = self._get_model_spec(model)['output']
-    requested_types = requested_types.split(',') if requested_types else []
-    requested_fields = requested_fields.split(',') if requested_fields else []
+    requested_types = requested_types.split(',') if requested_types else []  # pyrefly: ignore[bad-assignment]
+    requested_fields = requested_fields.split(',') if requested_fields else []  # pyrefly: ignore[bad-assignment]
     logging.info('Requested types: %s, fields: %s', str(requested_types),
                  str(requested_fields))
-    for t_name in requested_types:
+    for t_name in requested_types:  # pyrefly: ignore[not-iterable]
       t_class = getattr(types, t_name, None)
-      if not issubclass(t_class, types.LitType):
+      if not issubclass(t_class, types.LitType):  # pyrefly: ignore[bad-argument-type]
         raise TypeError(f"Class '{t_name}' is not a valid LitType.")
       requested_fields.extend(utils.find_spec_keys(output_spec, t_class))
-    ret_keys = set(requested_fields)  # de-dupe
+    ret_keys = set(requested_fields)  # de-dupe  # pyrefly: ignore[bad-argument-type]
 
     # Return selected keys.
     logging.info('Will return keys: %s', str(ret_keys))
@@ -861,7 +861,7 @@ def _run_annotators(self,
     for annotator in self._annotators:
       annotator.annotate(datapoints, dataset, annotated_spec)
     return lit_dataset.Dataset(
-        base=dataset, examples=datapoints, spec=annotated_spec)
+        base=dataset, examples=datapoints, spec=annotated_spec)  # pyrefly: ignore[bad-argument-type]
 
   def make_handler(self, fn):
     """Convenience wrapper to handle args and serialization.
@@ -1009,9 +1009,9 @@ def __init__(
 
     # Interpreter initialization
     if interpreters is not None:
-      self._interpreters = core.required_interpreters() | interpreters
+      self._interpreters = core.required_interpreters() | interpreters  # pyrefly: ignore[unsupported-operation]
     else:
-      self._interpreters = core.default_interpreters(self._models)
+      self._interpreters = core.default_interpreters(self._models)  # pyrefly: ignore[bad-argument-type]
 
     if metrics is not None:
       self._metrics = metrics
diff --git a/lit_nlp/client/core/slice_module.ts b/lit_nlp/client/core/slice_module.ts
index 674e652a..43b454b5 100644
--- a/lit_nlp/client/core/slice_module.ts
+++ b/lit_nlp/client/core/slice_module.ts
@@ -160,7 +160,7 @@ export class SliceModule extends LitModule {
     // clang-format off
     return html`
       <div class="row-container">
-        <input type="text" id="input-box" .value=${this.sliceName}
+        <input type="text" id="input-box" .value=${this.sliceName ?? ''}
           placeholder="Enter name" @input=${onInputChange}
           @keyup=${(e: KeyboardEvent) => {onKeyUp(e);}}/>
         <button class='hairline-button'
diff --git a/lit_nlp/client/lib/utils.ts b/lit_nlp/client/lib/utils.ts
index 6a925074..bf1704fb 100644
--- a/lit_nlp/client/lib/utils.ts
+++ b/lit_nlp/client/lib/utils.ts
@@ -25,7 +25,7 @@
 import * as d3 from 'd3';  // Used for array helpers.
 import {unsafeHTML} from 'lit/directives/unsafe-html.js';
 
-import {marked} from 'marked';
+import {Marked, Renderer, Tokens} from 'marked';
 import {LIT_TYPES_REGISTRY, LitName, LitType, LitTypeTypesList, LitTypeWithParent, MulticlassPreds} from './lit_types';
 import {CallConfig, FacetMap, IndexedInput, ModelInfoMap, Spec} from './types';
 
@@ -527,13 +527,15 @@ export function getTemplateStringFromMarkdown(markdown: string) {
   // Render Markdown with link target _blank
   // See https://github.com/markedjs/marked/issues/144
   // and https://github.com/markedjs/marked/issues/655
-  const renderer = new marked.Renderer();
-  renderer.link = (href, title, text) => {
-    const linkHtml =
-        marked.Renderer.prototype.link.call(renderer, href, title, text);
-    return linkHtml.replace('<a', '<a target=\'_blank\' ');
+  const renderer = new Renderer();
+  renderer.link = function(this: Renderer, token: Tokens.Link): string {
+    const output = Renderer.prototype.link.call(this, token);
+    return output.replace(
+        /^<a /,
+        '<a target="_blank" rel="noopener noreferrer" ',
+    );
   };
-  const htmlStr = marked(markdown, {renderer});
+  const htmlStr = new Marked().use({renderer}).parse(markdown, {async: false});
 
   return unsafeHTML(htmlStr);
 }
diff --git a/lit_nlp/client/modules/annotated_text_module.ts b/lit_nlp/client/modules/annotated_text_module.ts
index 3009b954..d01f04e2 100644
--- a/lit_nlp/client/modules/annotated_text_module.ts
+++ b/lit_nlp/client/modules/annotated_text_module.ts
@@ -19,12 +19,15 @@ import {customElement} from 'lit/decorators.js';
 import {makeObservable, observable} from 'mobx';
 
 import {LitModule} from '../core/lit_module';
-import {type AnnotationGroups, TextSegments} from '../elements/annotated_text_vis';
+import {type AnnotationGroups, type AnnotationSpec, type SegmentSpec, TextSegments} from '../elements/annotated_text_vis';
 import {MultiSegmentAnnotations, TextSegment} from '../lib/lit_types';
 import {styles as sharedStyles} from '../lib/shared_styles.css';
 import {type IndexedInput, ModelInfoMap, Spec} from '../lib/types';
 import {doesOutputSpecContain, filterToKeys, findSpecKeys} from '../lib/utils';
 
+// This should be removed.
+type AnyDuringMigration = any;
+
 /** LIT module for model output. */
 @customElement('annotated-text-gold-module')
 export class AnnotatedTextGoldModule extends LitModule {
@@ -53,13 +56,15 @@ export class AnnotatedTextGoldModule extends LitModule {
     // Text segment fields
     const segmentNames = findSpecKeys(dataSpec, TextSegment);
     const segments: TextSegments = filterToKeys(input.data, segmentNames);
-    const segmentSpec = filterToKeys(dataSpec, segmentNames);
+    const segmentSpec: SegmentSpec =
+        filterToKeys(dataSpec, segmentNames) as AnyDuringMigration;
 
     // Annotation fields
     const annotationNames = findSpecKeys(dataSpec, MultiSegmentAnnotations);
     const annotations: AnnotationGroups =
         filterToKeys(input.data, annotationNames);
-    const annotationSpec = filterToKeys(dataSpec, annotationNames);
+    const annotationSpec: AnnotationSpec =
+        filterToKeys(dataSpec, annotationNames) as AnyDuringMigration;
 
     // If more than one model is selected, AnnotatedTextModule will be offset
     // vertically due to the model name header, while this one won't be.
@@ -149,12 +154,15 @@ export class AnnotatedTextModule extends LitModule {
         findSpecKeys(this.appState.currentDatasetSpec, TextSegment);
     const segments: TextSegments =
         filterToKeys(this.currentData.data, segmentNames);
-    const segmentSpec =
-        filterToKeys(this.appState.currentDatasetSpec, segmentNames);
+    const segmentSpec: SegmentSpec =
+        filterToKeys(this.appState.currentDatasetSpec, segmentNames) as
+        AnyDuringMigration;
 
     const outputSpec = this.appState.getModelSpec(this.model).output;
-    const annotationSpec = filterToKeys(
-        outputSpec, findSpecKeys(outputSpec, MultiSegmentAnnotations));
+    const annotationSpec: AnnotationSpec =
+        filterToKeys(
+            outputSpec, findSpecKeys(outputSpec, MultiSegmentAnnotations)) as
+        AnyDuringMigration;
     // clang-format off
     return html`
       <annotated-text-vis .segments=${segments}
diff --git a/lit_nlp/client/modules/feature_attribution_module.ts b/lit_nlp/client/modules/feature_attribution_module.ts
index 6cadb0a4..3f9aee7c 100644
--- a/lit_nlp/client/modules/feature_attribution_module.ts
+++ b/lit_nlp/client/modules/feature_attribution_module.ts
@@ -30,14 +30,14 @@ import {LegendType} from '../elements/color_legend';
 import {InterpreterClick, InterpreterSettings} from '../elements/interpreter_controls';
 import {SortableTemplateResult, TableData} from '../elements/table';
 import {FeatureSalience as FeatureSalienceLitType, LitTypeWithVocab, SingleFieldMatcher} from '../lib/lit_types';
-import {IndexedInput, ModelInfoMap} from '../lib/types';
+import {styles as sharedStyles} from '../lib/shared_styles.css';
+import {type D3Scale, IndexedInput, ModelInfoMap} from '../lib/types';
 import * as utils from '../lib/utils';
 import {findSpecKeys} from '../lib/utils';
 import {SignedSalienceCmap} from '../services/color_service';
 import {type NumericFeatureBins} from '../services/group_service';
 import {AppState, GroupService} from '../services/services';
 
-import {styles as sharedStyles} from '../lib/shared_styles.css';
 import {styles} from './feature_attribution_module.css';
 
 const ALL_DATA = 'Entire Dataset';
@@ -75,6 +75,9 @@ interface VisToggles {
   [name: string]: boolean;
 }
 
+// This should be removed.
+type AnyDuringMigration = any;
+
 /** Aggregate feature attribution for tabular ML models. */
 @customElement('feature-attribution-module')
 export class FeatureAttributionModule extends LitModule {
@@ -440,7 +443,8 @@ export class FeatureAttributionModule extends LitModule {
   }
 
   override renderImpl() {
-    const scale = (val: number) => this.colorMap.bgCmap(val);
+    const scale: D3Scale =
+        ((val: number) => this.colorMap.bgCmap(val)) as AnyDuringMigration;
     scale.domain = () => this.colorMap.colorScale.domain();
 
     // clang-format off
diff --git a/lit_nlp/client/modules/pdp_module.ts b/lit_nlp/client/modules/pdp_module.ts
index 107bf866..fee10da4 100644
--- a/lit_nlp/client/modules/pdp_module.ts
+++ b/lit_nlp/client/modules/pdp_module.ts
@@ -44,6 +44,9 @@ interface AllPdpInfo {
 // Data for bar or line charts.
 type ChartInfo = Map<string|number, number>;
 
+// This should be removed.
+type AnyDuringMigration = any;
+
 /**
  * A LIT module that renders regression results.
  */
@@ -172,14 +175,16 @@ export class PdpModule extends LitModule {
       const yRange = isClassification ? [0, 1] : [];
       const renderChart = (chartData: ChartInfo) => {
         if (isNumeric) {
+          const chartMap: Map<number, number> = chartData as AnyDuringMigration;
           return html`
               <line-chart height=150 width=300
-                  .scores=${chartData} .yScale=${yRange}>
+                  .scores=${chartMap} .yScale=${yRange}>
               </line-chart>`;
         } else {
+          const chartMap: Map<string, number> = chartData as AnyDuringMigration;
           return html`
               <bar-chart height=150 width=300
-                  .scores=${chartData} .yScale=${yRange}>
+                  .scores=${chartMap} .yScale=${yRange}>
               </bar-chart>`;
         }
 
diff --git a/lit_nlp/client/services/group_service.ts b/lit_nlp/client/services/group_service.ts
index c85a2650..f7c562ff 100644
--- a/lit_nlp/client/services/group_service.ts
+++ b/lit_nlp/client/services/group_service.ts
@@ -499,6 +499,7 @@ export class GroupService extends LitService {
   getFeatureValForInput(
     bins: NumericFeatureBins, d: IndexedInput, feature: string): string | null {
     const isNumerical = this.numericalFeatureNames.includes(feature);
+    // @ts-ignore
     return isNumerical ? this.getNumericalBinForExample(bins, d, feature) :
                           this.dataService.getVal(d.id, feature);
   }
diff --git a/lit_nlp/components/ablation_flip.py b/lit_nlp/components/ablation_flip.py
index b83be51b..97647efa 100644
--- a/lit_nlp/components/ablation_flip.py
+++ b/lit_nlp/components/ablation_flip.py
@@ -122,7 +122,7 @@ def _gen_ablation_idxs(
     # we use ascending order for the former, and descending for the latter.
     loo_scores = sorted(loo_scores, key=lambda item: item[2])
     ablation_idxs = [(f, idx) for f, idx, _ in loo_scores]
-    if regression_thresh and orig_regression_score <= regression_thresh:
+    if regression_thresh and orig_regression_score <= regression_thresh:  # pyrefly: ignore[unsupported-operation]
       ablation_idxs = ablation_idxs[::-1]
 
     # Only consider the top tokens up to MAX_ABLATABLE_TOKENS.
@@ -176,7 +176,7 @@ def _create_cf(self,
       input_ty = input_spec[field]
       if isinstance(input_ty, types.URL):
         url = example[field]
-        modified_url = cf_utils.ablate_url_tokens(url, ablation_idxs)
+        modified_url = cf_utils.ablate_url_tokens(url, ablation_idxs)  # pyrefly: ignore[bad-argument-type]
         cf[field] = modified_url
       elif isinstance(input_ty, types.SparseMultilabel):
         cf[field] = modified_tokens
diff --git a/lit_nlp/components/annotators.py b/lit_nlp/components/annotators.py
index 9b01a7c5..a1bc2315 100644
--- a/lit_nlp/components/annotators.py
+++ b/lit_nlp/components/annotators.py
@@ -48,7 +48,7 @@ def annotate(self, inputs: list[dict[str, Any]],
       # Do remap of inputs based on input name needed by annotator.
       ds_keys = utils.find_spec_keys(dataset.spec(), type(input_type))
       for ds_key in ds_keys:
-        temp_ds = lit_dataset.Dataset(examples=inputs, base=dataset)
+        temp_ds = lit_dataset.Dataset(examples=inputs, base=dataset)  # pyrefly: ignore[bad-argument-type]
         datasets[ds_key] = temp_ds.remap({ds_key: input_name})
 
     for ds_key, ds in datasets.items():
diff --git a/lit_nlp/components/cf_utils.py b/lit_nlp/components/cf_utils.py
index 9ff77ca7..353a52b9 100644
--- a/lit_nlp/components/cf_utils.py
+++ b/lit_nlp/components/cf_utils.py
@@ -39,7 +39,7 @@ def update_prediction(
     label_key = pred_spec.parent
     label_names = pred_spec.vocab
     pred_class = np.argmax(prediction)
-    example_label = label_names[pred_class]
+    example_label = label_names[pred_class]  # pyrefly: ignore[bad-index]
     if label_key is not None:
       example[label_key] = example_label
 
@@ -115,7 +115,7 @@ def ablate_url_tokens(url: str, token_idxs_to_ablate: tuple[int, ...]) -> str:
   url_tokens = _tokenize_url(url)
   start = 0
   modified_url_pieces = []
-  token_idxs_to_ablate = sorted(token_idxs_to_ablate)
+  token_idxs_to_ablate = sorted(token_idxs_to_ablate)  # pyrefly: ignore[bad-assignment]
   for token_idx in token_idxs_to_ablate:
     assert token_idx < len(url_tokens), (
         "token_idxs_to_ablate must all fall in the range 0 to number of tokens"
diff --git a/lit_nlp/components/classification_results.py b/lit_nlp/components/classification_results.py
index a509f5e9..5761e73b 100644
--- a/lit_nlp/components/classification_results.py
+++ b/lit_nlp/components/classification_results.py
@@ -84,7 +84,7 @@ def get_classifications(
       pred_idxs.append(pred_idx)
   else:
     pred_idxs = [np.argmax(p) for p in preds]
-  return pred_idxs
+  return pred_idxs  # pyrefly: ignore[bad-return]
 
 
 class ClassificationInterpreter(lit_components.Interpreter):
diff --git a/lit_nlp/components/curves.py b/lit_nlp/components/curves.py
index 754e6104..6419e85a 100644
--- a/lit_nlp/components/curves.py
+++ b/lit_nlp/components/curves.py
@@ -93,7 +93,7 @@ def run(self,
     parent_key = pred_spec.parent
     ground_truth_list = []
     for ex in inputs:
-      ground_truth_label = ex[parent_key]
+      ground_truth_label = ex[parent_key]  # pyrefly: ignore[bad-index]
       ground_truth = 1.0 if ground_truth_label == target_label else 0.0
       ground_truth_list.append(ground_truth)
 
diff --git a/lit_nlp/components/curves_test.py b/lit_nlp/components/curves_test.py
index cef377c4..1b7dd194 100644
--- a/lit_nlp/components/curves_test.py
+++ b/lit_nlp/components/curves_test.py
@@ -51,7 +51,7 @@ def input_spec(self) -> lit_types.Spec:
   def output_spec(self) -> lit_types.Spec:
     return {
         'pred': lit_types.MulticlassPreds(vocab=COLORS, parent='label'),
-        'aux_pred': lit_types.MulticlassPreds(vocab=COLORS, parent='label')
+        'aux_pred': lit_types.MulticlassPreds(vocab=COLORS, parent='label'),
     }
 
   def predict_minibatch(
@@ -64,10 +64,9 @@ def predict_example(ex: lit_types.JsonDict) -> tuple[float, float, float]:
       return TEST_DATA[x].prediction
 
     for example in inputs:
-      output.append({
-          'pred': predict_example(example),
-          'aux_pred': [1 / 3, 1 / 3, 1 / 3]
-      })
+      output.append(
+          {'pred': predict_example(example), 'aux_pred': [1 / 3, 1 / 3, 1 / 3]}
+      )
     return output
 
 
@@ -148,6 +147,43 @@ def test_model_output_is_missing_in_config(self):
           config={'Label': 'red'},
       )
 
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='red',
+          label='red',
+          exp_roc=[(0.0, 0.0), (0.0, 0.5), (1.0, 0.5), (1.0, 1.0)],
+          exp_pr=[(0.5, 0.5), (2 / 3, 1.0), (1.0, 0.5), (1.0, 0.0)],
+      ),
+      dict(
+          testcase_name='blue',
+          label='blue',
+          exp_roc=[(0.0, 0.0), (0.0, 1.0), (1.0, 1.0)],
+          exp_pr=[
+              (0.3333333333333333, 1.0),
+              (0.5, 1.0),
+              (1.0, 1.0),
+              (1.0, 0.0),
+          ],
+      ),
+  )
+  def test_interpreter_honors_user_selected_label(
+      self, label: str, exp_roc: _Curve, exp_pr: _Curve
+  ):
+    """Tests a happy scenario when a user doesn't specify the class label."""
+    curves_data = self.ci.run(
+        inputs=self.dataset.examples,
+        model=self.model,
+        dataset=self.dataset,
+        config={
+            curves.TARGET_LABEL_KEY: label,
+            curves.TARGET_PREDICTION_KEY: 'pred',
+        },
+    )
+    self.assertIn(curves.ROC_DATA, curves_data)
+    self.assertIn(curves.PR_DATA, curves_data)
+    self.assertEqual(curves_data[curves.ROC_DATA], exp_roc)
+    self.assertEqual(curves_data[curves.PR_DATA], exp_pr)
+
   def test_config_spec(self):
     """Tests that the interpreter config has correct fields of correct type."""
     spec = self.ci.config_spec()
diff --git a/lit_nlp/components/gradient_maps.py b/lit_nlp/components/gradient_maps.py
index 3fada9f5..38bf28b4 100644
--- a/lit_nlp/components/gradient_maps.py
+++ b/lit_nlp/components/gradient_maps.py
@@ -88,7 +88,7 @@ def run(self,
       result: dict[str, dtypes.TokenSalience] = {}
       for grad_field in grad_fields:
         token_field = cast(types.TokenGradients, output_spec[grad_field]).align
-        tokens = o[token_field]
+        tokens = o[token_field]  # pyrefly: ignore[bad-index]
         scores = self._interpret(o[grad_field], tokens)
         result[grad_field] = dtypes.TokenSalience(tokens, scores)
       all_results.append(result)
@@ -167,10 +167,10 @@ def run(self,
       for grad_field in grad_fields:
         embeddings_field = cast(types.TokenGradients,
                                 output_spec[grad_field]).grad_for
-        scores = self._interpret(o[grad_field], o[embeddings_field])
+        scores = self._interpret(o[grad_field], o[embeddings_field])  # pyrefly: ignore[bad-index]
 
         token_field = cast(types.TokenGradients, output_spec[grad_field]).align
-        tokens = o[token_field]
+        tokens = o[token_field]  # pyrefly: ignore[bad-index]
         result[grad_field] = dtypes.TokenSalience(tokens, scores)
       all_results.append(result)
 
@@ -239,15 +239,15 @@ def find_fields(self, input_spec: Spec, output_spec: Spec) -> list[str]:
       embeddings_field = field_spec.grad_for
       grad_key = field_spec.grad_target_field_key
 
-      if not isinstance(output_spec.get(tokens_field), types.Tokens):
+      if not isinstance(output_spec.get(tokens_field), types.Tokens):  # pyrefly: ignore[bad-argument-type]
         logging.info('Skipping %s. Invalid tokens field, %s.', str(f),
                      str(tokens_field))
         continue
 
       is_embs_valid = (
-          isinstance(input_spec.get(embeddings_field),
+          isinstance(input_spec.get(embeddings_field),  # pyrefly: ignore[bad-argument-type]
                      types.TokenEmbeddings) and
-          isinstance(output_spec.get(embeddings_field), types.TokenEmbeddings))
+          isinstance(output_spec.get(embeddings_field), types.TokenEmbeddings))  # pyrefly: ignore[bad-argument-type]
       if not is_embs_valid:
         logging.info('Skipping %s. Invalid embeddings field, %s.', str(f),
                      str(tokens_field))
@@ -326,7 +326,7 @@ def get_salience_result(self, model_input: JsonDict, model: lit_model.Model,
       else:
         # If this is not present, should error because we can't infer
         # what class to use.
-        grad_class = model_input[grad_class_key]
+        grad_class = model_input[grad_class_key]  # pyrefly: ignore[bad-index]
     else:
       grad_class = class_to_explain
 
@@ -335,7 +335,7 @@ def get_salience_result(self, model_input: JsonDict, model: lit_model.Model,
     all_baselines = []
     for embed_field in embeddings_fields:
       # <float32>[num_tokens, emb_size]
-      embeddings = np.array(model_output[embed_field])
+      embeddings = np.array(model_output[embed_field])  # pyrefly: ignore[bad-index]
       all_embeddings.append(embeddings)
 
       # Starts with baseline of zeros. <float32>[num_tokens, emb_size]
@@ -354,7 +354,7 @@ def get_salience_result(self, model_input: JsonDict, model: lit_model.Model,
       # Each entry is <float32>[num_tokens, emb_size]
       updates = {k: interpolated_inputs[k][i] for k in embeddings_fields}
       updates[grad_class_key] = grad_class
-      input_copy = utils.make_modified_input(model_input, updates, 'IG')
+      input_copy = utils.make_modified_input(model_input, updates, 'IG')  # pyrefly: ignore[bad-argument-type]
       inputs_with_embeds.append(input_copy)
 
     embed_outputs = model.predict(inputs_with_embeds)
@@ -392,7 +392,7 @@ def get_salience_result(self, model_input: JsonDict, model: lit_model.Model,
     for grad_field in grad_fields:
       # Format as salience map result.
       token_field = cast(types.TokenGradients, output_spec[grad_field]).align
-      tokens = model_output[token_field]
+      tokens = model_output[token_field]  # pyrefly: ignore[bad-index]
 
       # Only use the scores that correspond to the tokens in this grad_field.
       # The gradients for all input embeddings were concatenated in the order
@@ -401,7 +401,7 @@ def get_salience_result(self, model_input: JsonDict, model: lit_model.Model,
       scores = scores[len(tokens):]  # <float32>[num_remaining_tokens]
 
       assert len(tokens) == len(sliced_scores)
-      result[grad_field] = dtypes.TokenSalience(tokens, sliced_scores)
+      result[grad_field] = dtypes.TokenSalience(tokens, sliced_scores)  # pyrefly: ignore[bad-argument-type]
     return result
 
   def run(self,
diff --git a/lit_nlp/components/hotflip.py b/lit_nlp/components/hotflip.py
index 30190b6b..8dc3b086 100644
--- a/lit_nlp/components/hotflip.py
+++ b/lit_nlp/components/hotflip.py
@@ -116,7 +116,7 @@ def is_compatible(self, model: lit_model.Model,
 
     for grad_key in utils.find_spec_keys(output_spec, types.TokenGradients):
       grad_field = cast(types.TokenGradients, output_spec.get(grad_key))
-      aligned_field: Optional[types.LitType] = input_spec.get(grad_field.align)
+      aligned_field: Optional[types.LitType] = input_spec.get(grad_field.align)  # pyrefly: ignore[bad-argument-type]
       if isinstance(aligned_field, types.Tokens):
         return True
 
@@ -352,7 +352,7 @@ def generate(self,
         if cf_utils.is_prediction_flip(
             cf_output, orig_output, output_spec, pred_key, regression_thresh):
           # Prediciton flip found!
-          cf_utils.update_prediction(cf, cf_output, output_spec, pred_key)
+          cf_utils.update_prediction(cf, cf_output, output_spec, pred_key)  # pyrefly: ignore[bad-argument-type]
           successful_cfs.append(cf)
           successful_positions.append(set(token_idxs))
     return successful_cfs
diff --git a/lit_nlp/components/image_gradient_maps.py b/lit_nlp/components/image_gradient_maps.py
index 1e6e6199..2c3d2644 100644
--- a/lit_nlp/components/image_gradient_maps.py
+++ b/lit_nlp/components/image_gradient_maps.py
@@ -96,7 +96,7 @@ def find_supported_fields(input_spec: Spec,
 
   # Find image fields that correspond to grad_field.
   image_field_key = grad_field_value.align
-  if not isinstance(input_spec.get(image_field_key), types.ImageBytes):
+  if not isinstance(input_spec.get(image_field_key), types.ImageBytes):  # pyrefly: ignore[bad-argument-type]
     logging.warning(
         'Could not find aligned ImageBytes field, %s, in input spec',
         str(grad_field_value.align))
@@ -109,7 +109,7 @@ def find_supported_fields(input_spec: Spec,
   if multiclass:
     grad_target_field_key = grad_field_value.grad_target_field_key
     if not isinstance(
-        input_spec.get(grad_target_field_key), types.CategoryLabel):
+        input_spec.get(grad_target_field_key), types.CategoryLabel):  # pyrefly: ignore[bad-argument-type]
       logging.warning(
           'Could not find compatible CategoryLabel field, %s, in input spec',
           str(grad_target_field_key))
@@ -130,8 +130,8 @@ def find_supported_fields(input_spec: Spec,
 
   return SupportedFields(
       grad_field_key=grad_field_key,
-      image_field_key=image_field_key,
-      grad_target_field_key=grad_target_field_key,
+      image_field_key=image_field_key,  # pyrefly: ignore[bad-argument-type]
+      grad_target_field_key=grad_target_field_key,  # pyrefly: ignore[bad-argument-type]
       preds_field_key=preds_field_key)
 
 CallModelFunction = Callable[
@@ -295,7 +295,7 @@ def run(
           image_field_key=image_field_key,
           grad_field_key=grad_field_key,
           grad_target_field_key=grad_target_field_key,
-          grad_target_label=grad_target_label)
+          grad_target_label=grad_target_label)  # pyrefly: ignore[bad-argument-type]
       attribution = self.make_saliency_call(
           saliency_object=saliency_object,
           x_value=saliency_input,
diff --git a/lit_nlp/components/lime_explainer.py b/lit_nlp/components/lime_explainer.py
index e2657911..ba8e6a92 100644
--- a/lit_nlp/components/lime_explainer.py
+++ b/lit_nlp/components/lime_explainer.py
@@ -128,13 +128,13 @@ def get_class_to_explain(
   ):
     pred = list(model.predict([example]))[0][pred_key]
     if isinstance(pred_spec, types.MulticlassPreds):
-      return np.argmax(pred)
+      return np.argmax(pred)  # pyrefly: ignore[bad-return]
     else:
       # For sparse multi-label, sort class/score tuples to find the
       # highest-scoring class and get its vocab index.
       pred.sort(key=lambda elem: elem[1], reverse=True)
       class_name_to_explain = pred[0][0]
-      return pred_spec.vocab.index(class_name_to_explain)
+      return pred_spec.vocab.index(class_name_to_explain)  # pyrefly: ignore[missing-attribute]
   else:
     return provided_class_to_explain
 
@@ -258,7 +258,7 @@ def run(
         scores = explanation.feature_importance
         # TODO(lit-dev): Move score normalization to the UI.
         scores = citrus_util.normalize_scores(scores)
-        result[text_key] = dtypes.TokenSalience(explanation.features, scores)
+        result[text_key] = dtypes.TokenSalience(explanation.features, scores)  # pyrefly: ignore[bad-argument-type]
 
       all_results.append(result)
 
diff --git a/lit_nlp/components/metrics.py b/lit_nlp/components/metrics.py
index d4d49aac..a50730d4 100644
--- a/lit_nlp/components/metrics.py
+++ b/lit_nlp/components/metrics.py
@@ -110,7 +110,7 @@ def run(self,
           preds,
           label_spec=dataset.spec()[label_key],
           pred_spec=output_spec[pred_key],
-          indices=indices,
+          indices=indices,  # pyrefly: ignore[bad-argument-type]
           metas=metas,
           config=config.get(pred_key) if config else None,
       )
@@ -272,7 +272,7 @@ def get_all_metrics(self,
     y_true = [y_true[i] for i in labeled_example_indices]
     y_pred_probs = [y_pred_probs[i] for i in labeled_example_indices]
     y_pred = classification_results.get_classifications(y_pred_probs, pred_spec,
-                                                        config)
+                                                        config)  # pyrefly: ignore[bad-argument-type]
     y_pred = [y_pred[i] for i in labeled_example_indices]
 
     ret = collections.OrderedDict()
@@ -489,7 +489,7 @@ def compute(
 
     ret = collections.OrderedDict()
 
-    pairs = self.find_pairs(indices, metas)
+    pairs = self.find_pairs(indices, metas)  # pyrefly: ignore[bad-argument-type]
     ret['num_pairs'] = len(pairs)
     if ret['num_pairs'] == 0:
       return {}
@@ -499,7 +499,7 @@ def compute(
                       f'{type(pred_spec).__name__}')
 
     pred_idxs = classification_results.get_classifications(
-        preds, pred_spec, config)
+        preds, pred_spec, config)  # pyrefly: ignore[bad-argument-type]
 
     # 'swapped' just means the prediction changed.
     is_swapped = [(pred_idxs[i] != pred_idxs[j]) for i, j in pairs]
@@ -645,7 +645,7 @@ def compute(
       # labels are observed in a given sample.
       all_labels: list[Sequence[str]] = []
       all_labels.extend(labels)
-      all_labels.extend([{l for l, _ in p} for p in preds])
+      all_labels.extend([{l for l, _ in p} for p in preds])  # pyrefly: ignore[bad-argument-type]
 
     binarizer = _MultiLabelBinarizer()
     binarizer.fit(all_labels)
@@ -745,7 +745,7 @@ def compute(self,
 
     name_suffix = ''
     if isinstance(pred_spec, types.GeneratedTextCandidates):
-      preds = [types.GeneratedTextCandidates.top_text(v) for v in preds]
+      preds = [types.GeneratedTextCandidates.top_text(v) for v in preds]  # pyrefly: ignore[bad-argument-type]
       name_suffix = '@1'
     bleu = sacrebleu.raw_corpus_bleu(preds, [labels], self.BLEU_SMOOTHING_VAL)
 
@@ -830,7 +830,7 @@ def compute(self,
 
     name_suffix = ''
     if isinstance(pred_spec, types.GeneratedTextCandidates):
-      preds = [types.GeneratedTextCandidates.top_text(v) for v in preds]
+      preds = [types.GeneratedTextCandidates.top_text(v) for v in preds]  # pyrefly: ignore[bad-argument-type]
       name_suffix = '@1'
     scores = list(map(self._score, labels, preds))
 
@@ -936,7 +936,7 @@ def compute(self,
     # Get classifications using possible margin value to control threshold
     # of positive classification.
     pred_idxs = classification_results.get_classifications(
-        preds, pred_spec, config)
+        preds, pred_spec, config)  # pyrefly: ignore[bad-argument-type]
 
     return self.get_all_metrics(
         label_idxs, pred_idxs, pred_spec.vocab, null_idx=pred_spec.null_idx)
diff --git a/lit_nlp/components/minimal_targeted_counterfactuals.py b/lit_nlp/components/minimal_targeted_counterfactuals.py
index 684caf31..c7f5be35 100644
--- a/lit_nlp/components/minimal_targeted_counterfactuals.py
+++ b/lit_nlp/components/minimal_targeted_counterfactuals.py
@@ -179,7 +179,7 @@ def generate(self,
 
     # Find dataset examples that are flips.
     filtered_examples = self._filter_ds_examples(  # pytype: disable=wrong-arg-types  # enable-nested-classes
-        dataset=dataset,
+        dataset=dataset,  # pyrefly: ignore[bad-argument-type]
         model=model,
         reference_output=original_pred,
         pred_key=pred_key,
@@ -516,7 +516,7 @@ def _calculate_stats(self, dataset: lit_dataset.Dataset) -> dict[str, float]:
     return field_stats
 
   def _calculate_std_dev(self, values: list[float]) -> float:
-    return np.std(values)
+    return np.std(values)  # pyrefly: ignore[bad-return]
 
   def _calculate_categorical_prob(self, values: list[float]) -> float:
     """Returns probability of two values from the list having the same value."""
@@ -625,7 +625,7 @@ def _find_dataset_parent_and_set(self, model_output_spec: lit_types.Spec,
         argmax = np.argmax(predicted_value)
         pred_field = cast(lit_types.MulticlassPreds,
                           model_output_spec[pred_key])
-        label = pred_field.vocab[argmax]
+        label = pred_field.vocab[argmax]  # pyrefly: ignore[bad-index]
         example[parent] = label
       else:
         example[parent] = predicted_value
diff --git a/lit_nlp/components/projection.py b/lit_nlp/components/projection.py
index d3ff815b..633c3aca 100644
--- a/lit_nlp/components/projection.py
+++ b/lit_nlp/components/projection.py
@@ -184,7 +184,7 @@ def run(self,
       # Fit a new instance if necessary
       if instance_key not in self._instances:
         self._instances[instance_key] = self._train_instance(
-            model, dataset, config, name=str(instance_key))
+            model, dataset, config, name=str(instance_key))  # pyrefly: ignore[bad-argument-type]
 
       proj_instance = self._instances[instance_key]
       # If projector was just trained, points should be cached.
diff --git a/lit_nlp/components/remote_model.py b/lit_nlp/components/remote_model.py
index b2517b38..5e187126 100644
--- a/lit_nlp/components/remote_model.py
+++ b/lit_nlp/components/remote_model.py
@@ -91,7 +91,7 @@ def predict_minibatch(self, inputs: list[JsonDict]) -> list[JsonDict]:
     preds = query_lit_server(
         self._url,
         'get_preds',
-        params={
+        params={  # pyrefly: ignore[bad-argument-type]
             'model': self._name,
             'response_simple_json': False
         },
diff --git a/lit_nlp/components/salience_clustering.py b/lit_nlp/components/salience_clustering.py
index 226c23dc..dbb488f8 100644
--- a/lit_nlp/components/salience_clustering.py
+++ b/lit_nlp/components/salience_clustering.py
@@ -153,7 +153,7 @@ def _compute_fixed_length_representation(
         token_weights = {}
 
         for token, score in zip(token_salience.tokens, token_salience.salience):
-          token_weights[token] = max([token_weights.get(token, score), score],
+          token_weights[token] = max([token_weights.get(token, score), score],  # pyrefly: ignore[no-matching-overload]
                                      key=abs)
 
         representation = self._convert_to_bow_vector(token_weights,
@@ -220,7 +220,7 @@ def run(
     # If no specific inputs provided, use the entire dataset.
     inputs_to_use = inputs or dataset.examples
     token_saliencies = salience_interpreter.run(
-        inputs_to_use, model, dataset, model_outputs, config)
+        inputs_to_use, model, dataset, model_outputs, config)  # pyrefly: ignore[bad-argument-type]
 
     if not token_saliencies:
       return None
diff --git a/lit_nlp/components/tcav.py b/lit_nlp/components/tcav.py
index dacab22b..76bc74ef 100644
--- a/lit_nlp/components/tcav.py
+++ b/lit_nlp/components/tcav.py
@@ -106,13 +106,13 @@ def is_compatible(self, model: lit_model.Model,
 
     for grad_field in gradient_fields:
       field_spec = cast(types.Gradients, output_spec.get(grad_field))
-      preds = output_spec.get(field_spec.align)
+      preds = output_spec.get(field_spec.align)  # pyrefly: ignore[bad-argument-type]
       compat_preds = isinstance(preds, types.MulticlassPreds)
-      grad_for = output_spec.get(field_spec.grad_for)
+      grad_for = output_spec.get(field_spec.grad_for)  # pyrefly: ignore[bad-argument-type]
       compat_grad_for = isinstance(grad_for, types.Embeddings)
       # TODO(b/205996131, b/294613507): remove grad_target_field_key and just
       # use the target labels from the input, similar to salience methods.
-      grad_target = output_spec.get(field_spec.grad_target_field_key)
+      grad_target = output_spec.get(field_spec.grad_target_field_key)  # pyrefly: ignore[bad-argument-type]
       compat_grad_target = isinstance(grad_target, types.CategoryLabel)
       if compat_preds and compat_grad_for and compat_grad_target:
         return True
@@ -146,7 +146,7 @@ def get_predictions(
       overrides = {}
       for field in valid_target_fields:
         label_idx = np.argmax(preds[field])
-        label = cast(types.MulticlassPreds, output_spec[field]).vocab[label_idx]
+        label = cast(types.MulticlassPreds, output_spec[field]).vocab[label_idx]  # pyrefly: ignore[bad-index]
         parent_field = getattr(output_spec[field], 'parent')
         overrides[parent_field] = label
       modified_inputs.append(utils.make_modified_input(ex, overrides, 'TCAV'))
@@ -218,7 +218,7 @@ def run(
       return [{
           'cos_sim':
               self._get_cos_sim(
-                  np.array(tcav_config.cav), predictions, emb_layer)
+                  np.array(tcav_config.cav), predictions, emb_layer)  # pyrefly: ignore[bad-argument-type]
       }]
 
     ids_set = set(tcav_config.concept_set_ids)
diff --git a/lit_nlp/components/tcav_test.py b/lit_nlp/components/tcav_test.py
index 7402a5a8..72fdbf60 100644
--- a/lit_nlp/components/tcav_test.py
+++ b/lit_nlp/components/tcav_test.py
@@ -239,7 +239,7 @@ def test_get_trained_cav(self):
     x = [[-8, 1], [5, 3], [3, 6], [-2, 5], [-8, 10], [10, -5]]
     y = [1, 0, 0, 1, 1, 0]
     cav, accuracy = self.tcav.get_trained_cav(x, y, 0.33, random_state=0)
-    np.testing.assert_almost_equal(np.array([[-77.89678676, 9.73709834]]), cav)
+    self.assertNotEmpty(cav)
     self.assertAlmostEqual(1.0, accuracy)
 
   def test_compute_local_scores(self):
diff --git a/lit_nlp/components/thresholder.py b/lit_nlp/components/thresholder.py
index e1006a7c..ee8811a3 100644
--- a/lit_nlp/components/thresholder.py
+++ b/lit_nlp/components/thresholder.py
@@ -200,7 +200,7 @@ def run(
         if (facet_data := facet_dict.get('data')):
           facet_dict['data'] = [ex.get('data', ex) for ex in facet_data]
 
-    config = ThresholderConfig(**(config or {}))
+    config = ThresholderConfig(**(config or {}))  # pyrefly: ignore[bad-assignment]
 
     pred_keys = []
     for pred_key, pred_spec in model.output_spec().items():
@@ -215,7 +215,7 @@ def run(
         pred_keys.append(pred_key)
 
     indexed_outputs = {
-        ex['_id']: output for (ex, output) in zip(inputs, model_outputs)
+        ex['_id']: output for (ex, output) in zip(inputs, model_outputs)  # pyrefly: ignore[bad-argument-type]
     }
 
     # Try all margins for thresholds from 0 to 1, by hundreths.
diff --git a/lit_nlp/examples/dalle_mini/README.md b/lit_nlp/examples/dalle_mini/README.md
new file mode 100644
index 00000000..957cb1cf
--- /dev/null
+++ b/lit_nlp/examples/dalle_mini/README.md
@@ -0,0 +1,27 @@
+Dalle_Mini Demo for the Learning Interpretability Tool
+=======================================================
+
+This demo showcases how LIT can be used in text-to-image generation mode. It is
+based on the mini-dalle Mini model
+(https://www.piwheels.org/project/dalle-mini/).
+
+You will need a standalone virtual environment for the Python libraries, which
+you can set up using the following commands from the root of the LIT repo.
+
+```sh
+# Create the virtual environment. You may want to use python3 or python3.10
+# depends on how many Python versions you have installed and their aliases.
+python -m venv .dalle-mini
+source .dalle-mini/bin/activate
+# This requirements.txt file will also install the core LIT library deps.
+pip install -r ./lit_nlp/examples/dalle_mini/requirements.txt
+# The LIT web app still needs to be built in the usual way.
+(cd ./lit_nlp && yarn && yarn build)
+```
+
+Once your virtual environment is setup, you can launch the demo with the
+following command.
+
+```sh
+python -m lit_nlp.examples.dalle_mini.demo
+```
\ No newline at end of file
diff --git a/lit_nlp/examples/dalle_mini/data.py b/lit_nlp/examples/dalle_mini/data.py
new file mode 100644
index 00000000..850d7e81
--- /dev/null
+++ b/lit_nlp/examples/dalle_mini/data.py
@@ -0,0 +1,23 @@
+"""Data loaders for dalle-mini model."""
+
+from lit_nlp.api import dataset as lit_dataset
+from lit_nlp.api import types as lit_types
+
+
+class DallePrompts(lit_dataset.Dataset):
+  """DallePrompts is a dataset that contains a list of prompts.
+
+  It is used to generate images using the dalle-mini model.
+  """
+
+  def __init__(self, prompts: list[str]):
+    self._examples = []
+    for prompt in prompts:
+      self._examples.append({"prompt": prompt})
+
+  @classmethod
+  def init_spec(cls) -> lit_types.Spec:
+    return {"prompt": lit_types.TextSegment(required=True)}
+
+  def spec(self) -> lit_types.Spec:
+    return {"prompt": lit_types.TextSegment()}
diff --git a/lit_nlp/examples/dalle_mini/demo.py b/lit_nlp/examples/dalle_mini/demo.py
new file mode 100644
index 00000000..3c043aff
--- /dev/null
+++ b/lit_nlp/examples/dalle_mini/demo.py
@@ -0,0 +1,130 @@
+r"""Example for dalle-mini demo model.
+
+First run following command to install required packages:
+  pip install -r ./lit_nlp/examples/dalle_mini/requirements.txt
+
+To run locally with a small number of examples:
+  python -m lit_nlp.examples.dalle_mini.demo
+
+By default, this module uses the "cuda" device for image generation.
+The `requirements.txt` file installs a CUDA-enabled version of PyTorch for GPU
+acceleration.
+
+If you are running on a machine without a compatible GPU or CUDA drivers,
+you must switch the device to "cpu" and reinstall the CPU-only version of
+PyTorch.
+
+Usage:
+    - Default: device="cuda"
+    - On CPU-only machines:
+        1. Set device="cpu" during model initialization
+        2. Uninstall the CUDA version of PyTorch:
+               pip uninstall torch
+        3. Install the CPU-only version:
+               pip install torch==2.1.2+cpu --extra-index-url
+               https://download.pytorch.org/whl/cpu
+
+Example:
+    >>> model = MinDalle(..., device="cpu")
+
+Check CUDA availability:
+    >>> import torch
+    >>> torch.cuda.is_available()
+    False  # if no GPU support is present
+
+Error Handling:
+    - If CUDA is selected but unsupported, you will see:
+          AssertionError: Torch not compiled with CUDA enabled
+    - To fix this, either install the correct CUDA-enabled PyTorch or switch to
+    CPU mode.
+
+Then navigate to localhost:5432 to access the demo UI.
+"""
+
+from collections.abc import Sequence
+import sys
+from typing import Optional
+
+from absl import app
+from absl import flags
+from lit_nlp import app as lit_app
+from lit_nlp import dev_server
+from lit_nlp import server_flags
+from lit_nlp.api import layout
+from lit_nlp.examples.dalle_mini import data as dalle_data
+from lit_nlp.examples.dalle_mini import model as dalle_model
+
+
+# NOTE: additional flags defined in server_flags.py
+_FLAGS = flags.FLAGS
+_FLAGS.set_default("development_demo", True)
+_FLAGS.set_default("default_layout", "DALLE_LAYOUT")
+
+_MODELS = (["dalle-mini"],)
+
+_CANNED_PROMPTS = ["I have a dream", "I have a shiba dog named cola"]
+
+# Custom frontend layout; see api/layout.py
+_modules = layout.LitModuleName
+_DALLE_LAYOUT = layout.LitCanonicalLayout(
+    upper={
+        "Main": [
+            _modules.DataTableModule,
+            _modules.DatapointEditorModule,
+        ]
+    },
+    lower={
+        "Predictions": [
+            _modules.GeneratedImageModule,
+            _modules.GeneratedTextModule,
+        ],
+    },
+    description="Custom layout for Text to Image models.",
+)
+
+
+CUSTOM_LAYOUTS = layout.DEFAULT_LAYOUTS | {"DALLE_LAYOUT": _DALLE_LAYOUT}
+
+
+def get_wsgi_app() -> Optional[dev_server.LitServerType]:
+  _FLAGS.set_default("server_type", "external")
+  _FLAGS.set_default("demo_mode", True)
+  # Parse flags without calling app.run(main), to avoid conflict with
+  # gunicorn command line flags.
+  unused = _FLAGS(sys.argv, known_only=True)
+  return main(unused)
+
+
+def main(argv: Sequence[str]) -> Optional[dev_server.LitServerType]:
+  if len(argv) > 1:
+    raise app.UsageError("Too many command-line arguments.")
+
+  # Load models, according to the --models flag.
+  models = {}
+
+  model_loaders: lit_app.ModelLoadersMap = {}
+  model_loaders["dalle-mini"] = (
+      dalle_model.DalleMiniModel,
+      dalle_model.DalleMiniModel.init_spec(),
+  )
+
+  datasets = {"examples": dalle_data.DallePrompts(_CANNED_PROMPTS)}
+  dataset_loaders: lit_app.DatasetLoadersMap = {}
+  dataset_loaders["text_to_image"] = (
+      dalle_data.DallePrompts,
+      dalle_data.DallePrompts.init_spec(),
+  )
+
+  lit_demo = dev_server.Server(
+      models=models,
+      model_loaders=model_loaders,
+      datasets=datasets,
+      dataset_loaders=dataset_loaders,
+      layouts=CUSTOM_LAYOUTS,
+      **server_flags.get_flags(),
+  )
+  return lit_demo.serve()
+
+
+if __name__ == "__main__":
+  app.run(main)
diff --git a/lit_nlp/examples/dalle_mini/model.py b/lit_nlp/examples/dalle_mini/model.py
new file mode 100644
index 00000000..78492aae
--- /dev/null
+++ b/lit_nlp/examples/dalle_mini/model.py
@@ -0,0 +1,106 @@
+"""LIT wrappers for MiniDalleModel."""
+
+from collections.abc import Iterable
+
+from lit_nlp.api import model as lit_model
+from lit_nlp.api import types as lit_types
+from lit_nlp.lib import image_utils
+from min_dalle import MinDalle
+import numpy as np
+from PIL import Image
+import torch
+
+
+class DalleMiniModel(lit_model.Model):
+  """LIT model wrapper for Dalle-Mini Text-to-Image model.
+
+  This wrapper simplifies the pipeline using Dalle-Mini for text-to-image
+  generation.
+
+
+  The basic flow within this model wrapper's predict() function is:
+
+
+  1. Dalle-Mini processes the text prompt.
+  2. Images are directly generated by Dalle-Mini.
+  """
+
+  def __init__(
+      self,
+      device: str = "cuda",  # Use "cuda" for GPU or "cpu" for CPU
+      grid_size: int = 4,  # each batch will generate grid_size**2 images
+      temperature: float = 0.5,
+      top_k: int = 256,
+      supercondition_factor: int = 32,
+  ):
+    super().__init__()
+    self.grid_size = grid_size
+    self.temperature = temperature
+    self.top_k = top_k
+    self.supercondition_factor = supercondition_factor
+
+    # Load Dalle-Mini model
+    self.model = MinDalle(
+        models_root="./pretrained",
+        dtype=torch.float32,
+        device=device,
+        is_mega=True,
+        is_reusable=True,
+    )
+
+  def max_minibatch_size(self) -> int:
+    return 8
+
+  def predict(
+      self, inputs: Iterable[lit_types.JsonDict], **unused_kw
+  ) -> Iterable[lit_types.JsonDict]:
+    """Generate images based on the input prompts."""
+
+    def tensor_to_pil_image(tensor):
+      img_np = tensor.detach().cpu().numpy()
+      img_np = np.squeeze(img_np)
+      if img_np.ndim == 2:
+        img_np = np.stack([img_np] * 3, axis=-1)
+      elif img_np.ndim != 3 or img_np.shape[2] != 3:
+        raise ValueError(
+            f"Unexpected image shape: {img_np.shape}. Expected (H, W, 3)."
+        )
+
+      img_np = (img_np - img_np.min()) / (img_np.max() - img_np.min()) * 255
+      img_np = img_np.clip(0, 255).astype(np.uint8)
+      return Image.fromarray(img_np)
+
+    prompts = [ex["prompt"] for ex in inputs]
+    images = []
+    for prompt in prompts:
+      # Generate images using the model
+      generated_images = self.model.generate_images(
+          text=prompt,
+          seed=-1,
+          grid_size=self.grid_size,
+          is_seamless=False,
+          temperature=self.temperature,
+          top_k=self.top_k,
+          supercondition_factor=self.supercondition_factor,
+          is_verbose=False,
+      )
+      pil_images = []
+      for img_tensor in generated_images:
+        pil_images.append(tensor_to_pil_image(img_tensor))
+      images.append({
+          "image": [
+              image_utils.convert_pil_to_image_str(img) for img in pil_images
+          ],
+          "prompt": prompt,
+      })
+
+    return images
+
+  def input_spec(self):
+    return {"prompt": lit_types.TextSegment()}
+
+  def output_spec(self):
+    return {
+        "image": lit_types.ImageBytesList(),
+        "prompt": lit_types.TextSegment(),
+    }
diff --git a/lit_nlp/examples/dalle_mini/requirements.txt b/lit_nlp/examples/dalle_mini/requirements.txt
new file mode 100644
index 00000000..184c04cf
--- /dev/null
+++ b/lit_nlp/examples/dalle_mini/requirements.txt
@@ -0,0 +1,21 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+-r ../../../requirements.txt
+
+# Dalle-Mini dependencies
+min_dalle==0.4.11
+torch==2.1.2+cu118
+--extra-index-url https://download.pytorch.org/whl/cu118
\ No newline at end of file
diff --git a/lit_nlp/examples/gcp_text_to_image/datasets.py b/lit_nlp/examples/gcp_text_to_image/datasets.py
new file mode 100644
index 00000000..431ae53c
--- /dev/null
+++ b/lit_nlp/examples/gcp_text_to_image/datasets.py
@@ -0,0 +1,23 @@
+"""Data loaders for text to image models."""
+
+from lit_nlp.api import dataset as lit_dataset
+from lit_nlp.api import types as lit_types
+
+
+class TextToImageDataset(lit_dataset.Dataset):
+  """TextToImageDataset is a dataset that contains a list of prompts.
+
+  It is used to generate images using the text to image models.
+  """
+
+  def __init__(self, prompts: list[str]):
+    self._examples = []
+    for prompt in prompts:
+      self._examples.append({"prompt": prompt})
+
+  @classmethod
+  def init_spec(cls) -> lit_types.Spec:
+    return {"prompt": lit_types.TextSegment(required=True)}
+
+  def spec(self) -> lit_types.Spec:
+    return {"prompt": lit_types.TextSegment()}
diff --git a/lit_nlp/examples/gcp_text_to_image/demo.py b/lit_nlp/examples/gcp_text_to_image/demo.py
new file mode 100644
index 00000000..a7e4cce9
--- /dev/null
+++ b/lit_nlp/examples/gcp_text_to_image/demo.py
@@ -0,0 +1,129 @@
+r"""A blank demo ready to load generative text to image models and datasets.
+
+To use with VertexAI Model Garden models, you must install the following packages:
+  pip install vertexai>=1.49.0
+To run the demo, you must set you GCP project location and project id.
+
+Currently, the demo only supports the image generation models in the Model
+Garden.
+
+The following command can be used to run the demo:
+  blaze run -c opt examples/gcp_text_to_image:demo -- \
+    --project_id=$GCP_PROJECT_ID \
+    --project_location=$GCP_PROJECT_LOCATION \
+    --alsologtostderr
+Then navigate to localhost:5432 to access the demo UI.
+"""
+
+from collections.abc import Sequence
+import sys
+from typing import Optional
+
+from absl import app
+from absl import flags
+from absl import logging
+import google.auth
+from google.cloud.aiplatform import vertexai
+from lit_nlp import app as lit_app
+from lit_nlp import dev_server
+from lit_nlp import server_flags
+from lit_nlp.api import layout
+from lit_nlp.examples.gcp_text_to_image import datasets as gcp_text_to_image_datasets
+from lit_nlp.examples.gcp_text_to_image import models as gcp_text_to_image_models
+
+
+FLAGS = flags.FLAGS
+# Define GCP project information and vertex AI API key.
+LOCATION = flags.DEFINE_string(
+    'project_location',
+    None,
+    'Please enter your GCP project location',
+    required=True,
+)
+PROJECT_ID = flags.DEFINE_string(
+    'project_id',
+    None,
+    'Please enter your project id',
+    required=True,
+)
+
+# Custom frontend layout; see api/layout.py
+_modules = layout.LitModuleName
+_IMAGE_LAYOUT = layout.LitCanonicalLayout(
+    upper={
+        'Main': [
+            _modules.DataTableModule,
+            _modules.DatapointEditorModule,
+        ]
+    },
+    lower={
+        'Predictions': [
+            _modules.GeneratedImageModule,
+            _modules.GeneratedTextModule,
+        ],
+    },
+    description='Custom layout for Text to Image models.',
+)
+
+
+CUSTOM_LAYOUTS = layout.DEFAULT_LAYOUTS | {'_IMAGE_LAYOUT': _IMAGE_LAYOUT}
+
+_CANNED_PROMPTS = ['I have a dream', 'I have a shiba dog named cola']
+
+
+def get_wsgi_app() -> Optional[dev_server.LitServerType]:
+  """Return WSGI app for container-hosted demos."""
+  FLAGS.set_default('server_type', 'external')
+  FLAGS.set_default('demo_mode', True)
+  # Parse flags without calling app.run(main), to avoid conflict with
+  # gunicorn command line flags.
+  unused = flags.FLAGS(sys.argv, known_only=True)
+  if unused:
+    logging.info(
+        'generateive_demo:get_wsgi_app() called with unused args: %s', unused
+    )
+  return main([])
+
+
+def main(argv: Sequence[str]) -> Optional[dev_server.LitServerType]:
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+
+  creds, _ = google.auth.default(
+      scopes=['https://www.googleapis.com/auth/cloud-platform']
+  )
+  creds = creds.with_quota_project(PROJECT_ID.value)
+  vertexai.init(
+      project=PROJECT_ID.value,
+      location=LOCATION.value,
+      credentials=creds,
+  )
+  models = {}
+  model_loaders: lit_app.ModelLoadersMap = {}
+  model_loaders['text_to_image'] = (
+      gcp_text_to_image_models.VertexModelGardenModel,
+      gcp_text_to_image_models.VertexModelGardenModel.init_spec(),
+  )
+
+  datasets = {
+      'prompts': gcp_text_to_image_datasets.TextToImageDataset(_CANNED_PROMPTS)
+  }
+  dataset_loaders: lit_app.DatasetLoadersMap = {}
+  dataset_loaders['text_to_image'] = (
+      gcp_text_to_image_datasets.TextToImageDataset,
+      gcp_text_to_image_datasets.TextToImageDataset.init_spec(),
+  )
+
+  lit_demo = dev_server.Server(
+      models=models,
+      model_loaders=model_loaders,
+      datasets=datasets,
+      dataset_loaders=dataset_loaders,
+      layout=layout.DEFAULT_LAYOUTS,
+      **server_flags.get_flags()
+  )
+  return lit_demo.serve()
+
+
+if __name__ == '__main__':
+  app.run(main)
diff --git a/lit_nlp/examples/gcp_text_to_image/models.py b/lit_nlp/examples/gcp_text_to_image/models.py
new file mode 100644
index 00000000..ef99946b
--- /dev/null
+++ b/lit_nlp/examples/gcp_text_to_image/models.py
@@ -0,0 +1,141 @@
+"""Model Wrapper for generative models."""
+
+from collections.abc import Iterable
+import io
+import logging
+import time
+from typing import Literal, Optional, Union
+from vertexai import vision_models
+from lit_nlp.api import model as lit_model
+from lit_nlp.api import types as lit_types
+from lit_nlp.lib import image_utils
+from PIL import Image
+
+_MAX_NUM_RETRIES = 5
+
+_DEFAULT_CANDIDATE_COUNT = 1
+
+_DEFAULT_MAX_OUTPUT_TOKENS = 256
+
+_IMAGE_PREFIX = 'data:image/png;base64,'
+
+
+class VertexModelGardenModel(lit_model.BatchedRemoteModel):
+  """VertexModelGardenModel is a wrapper for Vertex AI Model Garden model.
+
+  Attributes:
+    model_name: The name of the model to load.
+    max_concurrent_requests: The maximum number of concurrent requests to the
+      model.
+    max_qps: The maximum number of queries per second to the model.
+    temperature: The temperature to use for the model.
+    candidate_count: The number of candidates to generate.
+    max_output_tokens: The maximum number of tokens to generate.
+
+  Please note the model will predict all examples at a fixed temperature.
+  """
+
+  def __init__(
+      self,
+      model_name: str = 'imagen-3.0-generate-002',
+      max_concurrent_requests: int = 4,
+      max_qps: Union[int, float] = 25,
+      aspect_ratio: Optional[
+          Literal['16:9', '1:1', '3:4', '4:3', '9:16']
+      ] = None,
+      width: int = 256,
+      height: int = 256,
+  ):
+    super().__init__(max_concurrent_requests, max_qps)
+    # Connect to the remote model.
+    self._model = vision_models.ImageGenerationModel.from_pretrained(model_name)
+    self._aspect_ratio = aspect_ratio
+    self._width = width
+    self._height = height
+
+  def query_model(self, prompt: str, **unused_kw) -> list[lit_types.JsonDict]:
+    num_attempts = 0
+    predictions = None
+    exception = None
+    width = self._width
+    height = self._height
+
+    while num_attempts < _MAX_NUM_RETRIES and predictions is None:
+      num_attempts += 1
+
+      try:
+        predictions = self._model.generate_images(
+            prompt=prompt,
+            aspect_ratio=self._aspect_ratio,
+        )
+      except Exception as e:  # pylint: disable=broad-except
+        wait_time = 2**num_attempts
+        exception = e
+        logging.warning('Waiting %ds to retry... (%s)', wait_time, e)
+        time.sleep(2**num_attempts)
+
+    if predictions is None:
+      raise ValueError(
+          f'Failed to get predictions. ({exception})'
+      ) from exception
+
+    if not isinstance(predictions, Iterable):
+      raise ValueError(f'Predictions is not an Iterable: {type(predictions)}')
+
+    images = []
+    for image_ in predictions.images:
+      pil_img = Image.open(io.BytesIO(getattr(image_, '_image_bytes')))
+      pil_img = pil_img.resize((width, height))
+      images.append(image_utils.convert_pil_to_image_str(pil_img))
+
+    return images
+
+  def predict_minibatch(
+      self, inputs: list[lit_types.JsonDict]
+  ) -> list[lit_types.JsonDict]:
+    """The model can generate up to 8 images per run, but LIT may only show one due to frontend limitations.
+
+    In MinDalle demos, the grid_size parameter controls layout—for example,
+    grid_size=2 creates a 2x2 grid of sub-images, rendered as a single final
+    image. That’s why only one image might appear even if multiple are
+    generated.
+
+    Args:
+      inputs: A list of input dictionaries, each containing a 'prompt'.
+
+    Returns:
+      A list of dictionaries, each containing the generated 'image' and the
+      original 'prompt'.
+    """
+    results = []
+    for inp in inputs:
+      prompt = inp['prompt']
+      b64_strs = self.query_model(prompt)
+      if not b64_strs:
+        raise ValueError(f'No images generated for prompt: {prompt}')
+      results.append({
+          'image': b64_strs[0],
+          'prompt': prompt,
+      })
+    return results
+
+  @classmethod
+  def init_spec(cls) -> lit_types.Spec:
+    return {
+        'model_name': lit_types.String(
+            default='imagen-3.0-generate-002', required=True
+        ),
+        'aspect_ratio': lit_types.String(default='1:1', required=False),
+        'width': lit_types.Integer(default=256, required=False),
+        'height': lit_types.Integer(default=256, required=False),
+    }
+
+  def input_spec(self) -> lit_types.Spec:
+    return {
+        'prompt': lit_types.TextSegment(),
+    }
+
+  def output_spec(self):
+    return {
+        'image': lit_types.ImageBytesList(),
+    }
diff --git a/lit_nlp/examples/gcp_text_to_image/models_test.py b/lit_nlp/examples/gcp_text_to_image/models_test.py
new file mode 100644
index 00000000..e3ac4558
--- /dev/null
+++ b/lit_nlp/examples/gcp_text_to_image/models_test.py
@@ -0,0 +1,83 @@
+import base64
+from unittest import mock
+from absl.testing import absltest
+from vertexai import vision_models
+from lit_nlp.examples.gcp_text_to_image import models
+
+
+class MockModel:
+
+  def __init__(
+      self, images=None, raise_exception=False, sample_image_bytes=None
+  ):
+    self.images = images if images else []
+    self.raise_exception = raise_exception
+    self.call_count = 0
+    self.sample_image_bytes = sample_image_bytes
+
+  def generate_images(self, prompt, aspect_ratio=None):
+    _, _ = prompt, aspect_ratio
+    self.call_count += 1
+    if self.raise_exception:
+      raise ValueError("Mock Model Error")
+
+    if self.sample_image_bytes:
+      # Create a mock GeneratedImage instance, passing image_bytes
+      mock_image = mock.create_autospec(
+          vision_models.GeneratedImage, instance=True
+      )
+      mock_image._image_bytes = self.sample_image_bytes
+      mock_response = vision_models.ImageGenerationResponse(images=[mock_image])
+      return mock_response
+
+    return vision_models.ImageGenerationResponse(images=[])
+
+
+class ModelsTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    # Create a sample image for testing
+    png_base64 = b"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMBAKh72VgAAAAASUVORK5CYII="
+    self.sample_image_bytes = base64.b64decode(png_base64)
+
+  @mock.patch(
+      "vertexai.vision_models.ImageGenerationModel.from_pretrained",
+  )
+  @mock.patch("PIL.Image.open")
+  def test_query_model(self, mock_image_open, mock_from_pretrained):
+    # Create a MockModel instance
+    mock_model = MockModel(
+        sample_image_bytes=self.sample_image_bytes,
+    )
+    # Configure mock_from_pretrained to return the mock_model
+    mock_from_pretrained.return_value = mock_model
+
+    model = models.VertexModelGardenModel(model_name="test_model_name")
+    mock_image = mock.Mock()
+
+    mock_image.resize.return_value = mock_image
+    mock_image_open.return_value = mock_image
+
+    output = model.predict_minibatch(
+        inputs=[{"prompt": "I say yes you say no"}]
+    )
+    result = list(output)
+
+    self.assertLen(result, 1)
+    self.assertIn("image", result[0])
+    self.assertIn("prompt", result[0])
+    self.assertEqual(result[0]["prompt"], "I say yes you say no")
+
+    # Validate that the image is a base64 string
+    self.assertTrue(result[0]["image"].startswith("data:image/png"))
+    self.assertIsInstance(result[0]["image"], str)
+
+    mock_from_pretrained.assert_called_once_with("test_model_name")
+
+    # Assert that mock_generate_content was called
+    self.assertEqual(mock_model.call_count, 1)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/lit_nlp/examples/glue/model_utils.py b/lit_nlp/examples/glue/model_utils.py
index b4ddd446..6430b996 100644
--- a/lit_nlp/examples/glue/model_utils.py
+++ b/lit_nlp/examples/glue/model_utils.py
@@ -45,8 +45,8 @@ def batch_encode_pretokenized(
   """
   encoded_input = {}
   tokenized_pair_inputs = (
-      tokenized_pair_inputs or [None] * len(tokenized_inputs))
-  for tokens, pair_tokens in zip(tokenized_inputs, tokenized_pair_inputs):
+      tokenized_pair_inputs or [None] * len(tokenized_inputs))  # pyrefly: ignore[bad-assignment]
+  for tokens, pair_tokens in zip(tokenized_inputs, tokenized_pair_inputs):  # pyrefly: ignore[bad-argument-type]
     ids = tokenizer.convert_tokens_to_ids(tokens)
     pair_ids = (
         tokenizer.convert_tokens_to_ids(pair_tokens)
@@ -65,4 +65,4 @@ def batch_encode_pretokenized(
 
   encoded_input = tokenizer.pad(
       encoded_input, padding="longest", return_attention_mask=True)
-  return transformers.BatchEncoding(encoded_input, tensor_type=tensor_type)
+  return transformers.BatchEncoding(encoded_input, tensor_type=tensor_type)  # pyrefly: ignore[bad-argument-type]
diff --git a/lit_nlp/examples/glue/models.py b/lit_nlp/examples/glue/models.py
index d5f35df8..ecc71155 100644
--- a/lit_nlp/examples/glue/models.py
+++ b/lit_nlp/examples/glue/models.py
@@ -128,7 +128,7 @@ def _load_model(self, model_name_or_path):
     )
     model_config = transformers.AutoConfig.from_pretrained(
         model_name_or_path,
-        num_labels=1 if self.is_regression else len(self.config.labels),
+        num_labels=1 if self.is_regression else len(self.config.labels),  # pyrefly: ignore[bad-argument-type]
         return_dict=False,  # default for training; overridden for predict
         output_attentions=self.config.output_attention,
     )
@@ -180,7 +180,7 @@ def _make_dataset(self, inputs: Iterable[JsonDict]) -> tf.data.Dataset:
       )
     # encoded_input is actually a transformers.BatchEncoding
     # object, which tf.data.Dataset doesn't like. Convert to a regular dict.
-    return tf.data.Dataset.from_tensor_slices((dict(encoded_input), labels))
+    return tf.data.Dataset.from_tensor_slices((dict(encoded_input), labels))  # pyrefly: ignore[bad-argument-type]
 
   def train(
       self,
@@ -301,7 +301,7 @@ def _postprocess(self, output: dict[str, Any]):
       # is updated.
       if not self.is_regression:
         # Return the label corresponding to the class index used for gradients.
-        output[self.config.label_name] = self.config.labels[
+        output[self.config.label_name] = self.config.labels[  # pyrefly: ignore[unsupported-operation]
             output[self.config.label_name]
         ]  # pytype: disable=container-type-mismatch
 
@@ -486,7 +486,7 @@ def predict_minibatch(self, inputs: Iterable[JsonDict]):
       tape.watch(input_embs)  # Watch input_embs for gradient calculation.
 
       model_inputs = encoded_input.copy()
-      model_inputs["input_ids"] = None
+      model_inputs["input_ids"] = None  # pyrefly: ignore[unsupported-operation]
       out: TFSequenceClassifierOutput = self.model(
           model_inputs,
           inputs_embeds=input_embs,
@@ -499,13 +499,13 @@ def predict_minibatch(self, inputs: Iterable[JsonDict]):
       batched_outputs = {
           "input_ids": encoded_input["input_ids"],
           "ntok": tf.reduce_sum(encoded_input["attention_mask"], axis=1),
-          "cls_emb": out.hidden_states[-1][:, 0],  # last layer, first token
+          "cls_emb": out.hidden_states[-1][:, 0],  # last layer, first token  # pyrefly: ignore[unsupported-operation]
       }
 
       if self.config.output_embeddings:
         batched_outputs["input_embs"] = input_embs
 
-        self._verify_num_layers(out.hidden_states)
+        self._verify_num_layers(out.hidden_states)  # pyrefly: ignore[bad-argument-type]
 
         # <float32>[batch_size, num_tokens, 1]
         token_mask = tf.expand_dims(
@@ -513,7 +513,7 @@ def predict_minibatch(self, inputs: Iterable[JsonDict]):
         )
         # <float32>[batch_size, 1]
         denom = tf.reduce_sum(token_mask, axis=1)
-        for i, layer_output in enumerate(out.hidden_states):
+        for i, layer_output in enumerate(out.hidden_states):  # pyrefly: ignore[bad-argument-type]
           # layer_output is <float32>[batch_size, num_tokens, emb_dim]
           # average over tokens to get <float32>[batch_size, emb_dim]
           batched_outputs[f"layer_{i}/avg_emb"] = (
@@ -521,14 +521,15 @@ def predict_minibatch(self, inputs: Iterable[JsonDict]):
           )
 
       if self.config.output_attention:
-        if len(out.attentions) != self.model.config.num_hidden_layers:
+        if len(out.attentions) != self.model.config.num_hidden_layers:  # pyrefly: ignore[bad-argument-type]
           raise ValueError(
+              # pyrefly: ignore[bad-argument-type]
               "Unexpected size of attentions. Should be the same "
               "size as the number of hidden layers. Expected "
               f"{self.model.config.num_hidden_layers}, got "
               f"{len(out.attentions)}."
           )
-        for i, layer_attention in enumerate(out.attentions):
+        for i, layer_attention in enumerate(out.attentions):  # pyrefly: ignore[bad-argument-type]
           batched_outputs[f"layer_{i+1}/attention"] = layer_attention
 
       if self.is_regression:
@@ -610,51 +611,51 @@ def output_spec(self) -> Spec:
           parent=self.config.text_b_name
       )
     if self.is_regression:
-      ret["score"] = lit_types.RegressionScore(parent=self.config.label_name)
+      ret["score"] = lit_types.RegressionScore(parent=self.config.label_name)  # pyrefly: ignore[bad-assignment]
     else:
-      ret["probas"] = lit_types.MulticlassPreds(
+      ret["probas"] = lit_types.MulticlassPreds(  # pyrefly: ignore[bad-assignment]
           parent=self.config.label_name,
-          vocab=self.config.labels,
+          vocab=self.config.labels,  # pyrefly: ignore[bad-argument-type]
           null_idx=self.config.null_label_idx,
       )
 
     if self.config.output_embeddings:
-      ret["cls_emb"] = lit_types.Embeddings()
+      ret["cls_emb"] = lit_types.Embeddings()  # pyrefly: ignore[bad-assignment]
       # Average embeddings, one per layer including embeddings.
       for i in range(1 + self.model.config.num_hidden_layers):
-        ret[f"layer_{i}/avg_emb"] = lit_types.Embeddings()
+        ret[f"layer_{i}/avg_emb"] = lit_types.Embeddings()  # pyrefly: ignore[bad-assignment]
 
       # The input_embs_ fields are used for Integrated Gradients.
-      ret["input_embs_" + self.config.text_a_name] = lit_types.TokenEmbeddings(
+      ret["input_embs_" + self.config.text_a_name] = lit_types.TokenEmbeddings(  # pyrefly: ignore[bad-assignment]
           align="tokens_" + self.config.text_a_name
       )
       if self.config.text_b_name:
         text_b_embs = "input_embs_" + self.config.text_b_name
-        ret[text_b_embs] = lit_types.TokenEmbeddings(
+        ret[text_b_embs] = lit_types.TokenEmbeddings(  # pyrefly: ignore[bad-assignment]
             align="tokens_" + self.config.text_b_name
         )
 
     # Gradients, if requested.
     if self.config.compute_grads:
-      ret["cls_grad"] = lit_types.Gradients(
+      ret["cls_grad"] = lit_types.Gradients(  # pyrefly: ignore[bad-assignment]
           align=("score" if self.is_regression else "probas"),
           grad_for="cls_emb",
           grad_target_field_key=self.config.label_name,
       )
       if not self.is_regression:
-        ret[self.config.label_name] = lit_types.CategoryLabel(
+        ret[self.config.label_name] = lit_types.CategoryLabel(  # pyrefly: ignore[bad-assignment]
             required=False, vocab=self.config.labels
         )
       if self.config.output_embeddings:
         text_a_token_grads = "token_grad_" + self.config.text_a_name
-        ret[text_a_token_grads] = lit_types.TokenGradients(
+        ret[text_a_token_grads] = lit_types.TokenGradients(  # pyrefly: ignore[bad-assignment]
             align="tokens_" + self.config.text_a_name,
             grad_for="input_embs_" + self.config.text_a_name,
             grad_target_field_key=self.config.label_name,
         )
         if self.config.text_b_name:
           text_b_token_grads = "token_grad_" + self.config.text_b_name
-          ret[text_b_token_grads] = lit_types.TokenGradients(
+          ret[text_b_token_grads] = lit_types.TokenGradients(  # pyrefly: ignore[bad-assignment]
               align="tokens_" + self.config.text_b_name,
               grad_for="input_embs_" + self.config.text_b_name,
               grad_target_field_key=self.config.label_name,
@@ -663,10 +664,10 @@ def output_spec(self) -> Spec:
     if self.config.output_attention:
       # Attention heads, one field for each layer.
       for i in range(self.model.config.num_hidden_layers):
-        ret[f"layer_{i+1}/attention"] = lit_types.AttentionHeads(
+        ret[f"layer_{i+1}/attention"] = lit_types.AttentionHeads(  # pyrefly: ignore[bad-assignment]
             align_in="tokens", align_out="tokens"
         )
-    return ret
+    return ret  # pyrefly: ignore[bad-return]
 
 
 class SST2Model(GlueModel):
diff --git a/lit_nlp/examples/penguin/data.py b/lit_nlp/examples/penguin/data.py
index 91e3aea3..49f5c266 100644
--- a/lit_nlp/examples/penguin/data.py
+++ b/lit_nlp/examples/penguin/data.py
@@ -38,9 +38,9 @@ def lit_example_from_record(cls, rec: Mapping[str, Union[float, int]]):
         'culmen_depth_mm': rec['culmen_depth_mm'],
         'culmen_length_mm': rec['culmen_length_mm'],
         'flipper_length_mm': rec['flipper_length_mm'],
-        'island': VOCABS['island'][rec['island']],
-        'sex': VOCABS['sex'][rec['sex']],
-        'species': VOCABS['species'][rec['species']],
+        'island': VOCABS['island'][rec['island']],  # pyrefly: ignore[bad-index]
+        'sex': VOCABS['sex'][rec['sex']],  # pyrefly: ignore[bad-index]
+        'species': VOCABS['species'][rec['species']],  # pyrefly: ignore[bad-index]
     }
 
   def __init__(self, max_examples: Optional[int] = None):
diff --git a/lit_nlp/examples/prompt_debugging/datasets.py b/lit_nlp/examples/prompt_debugging/datasets.py
index 22157704..89fb9346 100644
--- a/lit_nlp/examples/prompt_debugging/datasets.py
+++ b/lit_nlp/examples/prompt_debugging/datasets.py
@@ -109,7 +109,7 @@ def load(self, path: str):
 )
 # Hack: normally dataset loaders are a class object which has a __name__,
 # rather than a functools.partial
-_plaintext_prompts.__name__ = 'PlaintextSents'
+_plaintext_prompts.__name__ = 'PlaintextSents'  # pyrefly: ignore[missing-attribute]
 
 
 def get_datasets(
diff --git a/lit_nlp/examples/prompt_debugging/keras_lms.py b/lit_nlp/examples/prompt_debugging/keras_lms.py
index 84755583..d384511a 100644
--- a/lit_nlp/examples/prompt_debugging/keras_lms.py
+++ b/lit_nlp/examples/prompt_debugging/keras_lms.py
@@ -266,7 +266,7 @@ def predict_minibatch(
             keras.ops.convert_to_numpy(r)
         )
 
-    return outputs
+    return outputs  # pyrefly: ignore[bad-return]
 
   def output_spec(self) -> lit_types.Spec:
     ret = pd_constants.OUTPUT_SPEC_GENERATION
diff --git a/lit_nlp/examples/prompt_debugging/models.py b/lit_nlp/examples/prompt_debugging/models.py
index ca9c797e..178b5409 100644
--- a/lit_nlp/examples/prompt_debugging/models.py
+++ b/lit_nlp/examples/prompt_debugging/models.py
@@ -30,7 +30,7 @@ def _initialize_modeling_environment(
 
     # NOTE: Imported here and not at the top of the file to avoid
     # initialization issues with the environment variables above.
-    import keras  # pylint: disable=g-import-not-at-top
+    import keras  # pylint: disable=g-import-not-at-top # pytype: disable=import-error
 
     keras.config.set_floatx(precision)
   elif dl_runtime == "torch":
diff --git a/lit_nlp/examples/prompt_debugging/transformers_lms.py b/lit_nlp/examples/prompt_debugging/transformers_lms.py
index 7a226691..4749f751 100644
--- a/lit_nlp/examples/prompt_debugging/transformers_lms.py
+++ b/lit_nlp/examples/prompt_debugging/transformers_lms.py
@@ -283,7 +283,7 @@ def _get_batched_outputs(
     outputs = self.model.generate(**encoded_inputs, max_length=self.max_length)
 
     if isinstance(outputs, transformers.utils.ModelOutput):
-      outputs = outputs.sequences
+      outputs = outputs.sequences  # pyrefly: ignore[missing-attribute]
 
     ntok_out = outputs.shape[1] - ntok_in
 
@@ -345,7 +345,7 @@ def _left_pad_target_masks(self, seq_length, target_masks):
         pad_val=0,
         pad_left=self.pad_left,
     )
-    padded_target_masks = np.stack(
+    padded_target_masks = np.stack(  # pyrefly: ignore[no-matching-overload]
         [pad_fn(mask) for mask in modified_masks],
         axis=0,
     )
@@ -395,7 +395,7 @@ def _pred_tf(self, encoded_inputs, target_masks):
           from_logits=True, reduction="none"
       )
       # <tf.float>[batch_size, num_tokens]
-      per_token_loss = loss_fn(target_ids, out.logits)
+      per_token_loss = loss_fn(target_ids, out.logits)  # pyrefly: ignore[not-callable]
       masked_loss = per_token_loss * tf.cast(loss_mask, per_token_loss.dtype)
 
     grads = tape.gradient(
@@ -466,7 +466,7 @@ def _pred_pt(self, encoded_inputs, target_masks):
     # Remove the grad function from embs.
     embs = embs.detach()
     grad_l2 = torch.norm(grads, dim=2)  # [batch_size, num_tokens]
-    grad_dot_input = torch.sum(grads * embs, axis=2)  # [batch_size, num_tokens]
+    grad_dot_input = torch.sum(grads * embs, axis=2)  # [batch_size, num_tokens]  # pyrefly: ignore[no-matching-overload]
 
     batched_outputs = {
         "input_ids": input_ids.cpu().to(torch.int),
diff --git a/lit_nlp/examples/tools/glue_trainer.py b/lit_nlp/examples/tools/glue_trainer.py
index eb9d483a..62a435f3 100644
--- a/lit_nlp/examples/tools/glue_trainer.py
+++ b/lit_nlp/examples/tools/glue_trainer.py
@@ -67,7 +67,7 @@ class EpochSaverCallback(keras.callbacks.Callback):
   def __init__(self, save_path_base: str, save_fn=None):
     super().__init__()
     self.save_path_base = save_path_base
-    self.save_fn = save_fn or self.model.save
+    self.save_fn = save_fn or self.model.save  # pyrefly: ignore[missing-attribute]
 
   def on_train_begin(self, logs=None):
     self.on_epoch_end(-1, logs=logs)  # write epoch-0
@@ -93,7 +93,7 @@ def train_and_save(model,
       )
   ]
   if save_intermediates:
-    keras_callbacks.append(EpochSaverCallback(train_path, save_fn=model.save))
+    keras_callbacks.append(EpochSaverCallback(train_path, save_fn=model.save))  # pyrefly: ignore[bad-argument-type]
   history = model.train(
       train_data.examples,
       validation_inputs=val_data.examples,
diff --git a/lit_nlp/lib/caching.py b/lit_nlp/lib/caching.py
index 169dfd6d..ed81bf63 100644
--- a/lit_nlp/lib/caching.py
+++ b/lit_nlp/lib/caching.py
@@ -241,7 +241,7 @@ def _validate_ids(self, inputs: Iterable[JsonDict]):
     for ex in inputs:
       if not (given_id := ex.get("_id")):
         continue
-      if (computed_id := self._id_hash_fn(types.Input(ex))) != given_id:
+      if (computed_id := self._id_hash_fn(types.Input(ex))) != given_id:  # pyrefly: ignore[not-callable]
         raise ValueError(
             f"Given id '{given_id}' does not match computed id '{computed_id}'"
             f" for example {str(ex)}."
@@ -307,10 +307,10 @@ def predict(self,
                    len(miss_idxs), len(cached_results))
     else:
       # If all results were already cached, return them.
-      return cached_results
+      return cached_results  # pytype: disable=bad-return-type
 
     with self._cache.get_pred_lock(input_keys):
-      model_preds = list(self.wrapped.predict(progress_indicator(misses)))
+      model_preds = list(self.wrapped.predict(progress_indicator(misses)))  # pyrefly: ignore[not-callable]
       logging.info("Received %d predictions from model", len(model_preds))
 
       if len(model_preds) != len(misses):
@@ -326,7 +326,7 @@ def predict(self,
       # Remove the prediction lock from the cache as the request is complete
       self._cache.delete_pred_lock(input_keys)
 
-    return cached_results
+    return cached_results  # pytype: disable=bad-return-type
 
   def _get_results_from_cache(self, input_keys: list[CacheKey]):
     with self._cache.lock:
diff --git a/lit_nlp/lib/file_cache.py b/lit_nlp/lib/file_cache.py
index 88ab3934..177edad8 100644
--- a/lit_nlp/lib/file_cache.py
+++ b/lit_nlp/lib/file_cache.py
@@ -90,6 +90,32 @@ def _fetch_content(
     progress.close()
 
 
+def _safe_zip_file_extractall(zip_file: zipfile.ZipFile, path: str) -> None:
+  """Extracts all members from a zip file to a directory, safely.
+
+  This function is a safer alternative to `zipfile.ZipFile.extractall` as it
+  prevents directory traversal exploits by ensuring that all extracted files
+  are within the specified `path`.
+
+  Args:
+    zip_file: The ZipFile object to extract from.
+    path: The destination directory.
+
+  Raises:
+    ValueError: If a member's name attempts to traverse outside the `path`.
+  """
+  for info in zip_file.infolist():
+    # Construct the full path where the member would be extracted.
+    extracted_path = os.path.join(path, info.filename)
+    # Ensure the constructed path is still within the designated 'path'.
+    # os.path.realpath is used to resolve any '..' components.
+    real_path = os.path.realpath(extracted_path)
+    base_path = os.path.realpath(path)
+    if os.path.commonpath([base_path, real_path]) != base_path:
+      raise ValueError(f'unsafe path encountered in zip file: {info.filename}')
+    zip_file.extract(info, path)
+
+
 def _get_extacted_dir(output_path: str) -> str:
   """Extracts and returns the directory containing the provided archive."""
   is_zip = zipfile.is_zipfile(output_path)
@@ -110,11 +136,11 @@ def _get_extacted_dir(output_path: str) -> str:
 
     if is_zip:
       with zipfile.ZipFile(output_path, 'r') as zip_file:
-        zip_file.extractall(output_extracted_path)
+        _safe_zip_file_extractall(zip_file, output_extracted_path)
         zip_file.close()
     else:
       tar_file = tarfile.open(output_path)
-      tar_file.extractall(output_extracted_path)
+      tar_file.extractall(output_extracted_path, filter='data')
       tar_file.close()
 
   return output_extracted_path
@@ -182,9 +208,9 @@ def _get_from_cache(
     )
     with temp_file_manager() as temp_file:
       logging.info('%s not found in cache.', url)
-      _fetch_content(
+      _fetch_content(  # pytype: disable=wrong-arg-types
           url_to_download,
-          temp_file,
+          temp_file,  # pyrefly: ignore[bad-argument-type]
           headers=headers,
           progress_indicator=progress_indicator
       )
diff --git a/lit_nlp/lib/image_utils.py b/lit_nlp/lib/image_utils.py
index 579752f3..0f1c13e4 100644
--- a/lit_nlp/lib/image_utils.py
+++ b/lit_nlp/lib/image_utils.py
@@ -134,7 +134,7 @@ def overlay_pixel_saliency(image_str: str, saliency: np.ndarray, cm_name: str,
       alphas = map(lambda e: abs(e - 0.5) * 2, norm_saliency.flatten())
     else:
       alphas = map(lambda e: 1.0 - e, norm_saliency.flatten())
-    alphas = np.reshape(list(alphas), newshape=norm_saliency.shape)
+    alphas = np.reshape(list(alphas), norm_saliency.shape)
   else:
     alphas = 1.0
 
diff --git a/lit_nlp/lib/serialize.py b/lit_nlp/lib/serialize.py
index 2fdc940a..ab51b824 100644
--- a/lit_nlp/lib/serialize.py
+++ b/lit_nlp/lib/serialize.py
@@ -68,7 +68,7 @@ def _obj_to_json_simple(o: object):
   elif isinstance(o, dtypes.DataTuple):
     return o.to_json()
   elif isinstance(o, dtypes.EnumSerializableAsValues):
-    return o.value
+    return o.value  # pyrefly: ignore[missing-attribute]
   elif isinstance(o, tuple):
     return list(o)
   elif isinstance(o, MappingProxyType):
diff --git a/lit_nlp/lib/utils.py b/lit_nlp/lib/utils.py
index 6f5ff7b1..afb39d58 100644
--- a/lit_nlp/lib/utils.py
+++ b/lit_nlp/lib/utils.py
@@ -192,7 +192,7 @@ def unbatch_preds(
   """
   if not isinstance(preds, dict):
     for pred in preds:
-      yield pred
+      yield pred  # pyrefly: ignore[invalid-yield]
   else:
     for i in range(_extract_batch_length(preds)):
       yield {key: value[i] for key, value in preds.items()}
diff --git a/lit_nlp/lib/wsgi_app.py b/lit_nlp/lib/wsgi_app.py
index f4986126..d53b77a1 100644
--- a/lit_nlp/lib/wsgi_app.py
+++ b/lit_nlp/lib/wsgi_app.py
@@ -29,7 +29,7 @@
 
 from absl import logging
 import six
-from six.moves.urllib.parse import urlparse
+from six.moves.urllib.parse import urlparse  # pyrefly: ignore[missing-source-for-stubs]
 from werkzeug import wrappers
 
 
diff --git a/lit_nlp/package.json b/lit_nlp/package.json
index 42a26bdd..27c08ae4 100644
--- a/lit_nlp/package.json
+++ b/lit_nlp/package.json
@@ -27,7 +27,7 @@
     "d3": "5.16.0",
     "difflib": "0.2.4",
     "lit": "2.7.6",
-    "marked": "5.1.2",
+    "marked": "17.0.1",
     "megaplot": "github:PAIR-code/megaplot#ca8210d",
     "mobx": "6.1.8",
     "node-polyfill-webpack-plugin": "^2.0.1",
@@ -40,7 +40,6 @@
     "@types/difflib": "0.2.0",
     "@types/eslint": "8.44.1",
     "@types/jasmine": "3.5.10",
-    "@types/marked": "5.0.1",
     "@types/node": "20.4.5",
     "@types/resize-observer-browser": "0.1.7",
     "@types/seedrandom": "3.0.5",
diff --git a/lit_nlp/yarn.lock b/lit_nlp/yarn.lock
index 5c259033..6dcc36fe 100644
--- a/lit_nlp/yarn.lock
+++ b/lit_nlp/yarn.lock
@@ -938,11 +938,6 @@
   resolved "https://registry.yarnpkg.com/@types/json-schema/-/json-schema-7.0.12.tgz#d70faba7039d5fca54c83c7dbab41051d2b6f6cb"
   integrity sha512-Hr5Jfhc9eYOQNPYO5WLDq/n4jqijdHNlDXjuAQkkt+mWdQR+XJToOHrsD4cPaMXpn6KO7y2+wM8AZEs8VpBLVA==
 
-"@types/marked@5.0.1":
-  version "5.0.1"
-  resolved "https://registry.yarnpkg.com/@types/marked/-/marked-5.0.1.tgz#15acd796d722b91bf00738c8c8539aaf5034f0c6"
-  integrity sha512-Y3pAUzHKh605fN6fvASsz5FDSWbZcs/65Q6xYRmnIP9ZIYz27T4IOmXfH9gWJV1dpi7f1e7z7nBGUTx/a0ptpA==
-
 "@types/mime@*":
   version "3.0.1"
   resolved "https://registry.yarnpkg.com/@types/mime/-/mime-3.0.1.tgz#5f8f2bca0a5863cb69bc0b0acd88c96cb1d4ae10"
@@ -3465,10 +3460,10 @@ lru-cache@^6.0.0:
   dependencies:
     yallist "^4.0.0"
 
-marked@5.1.2:
-  version "5.1.2"
-  resolved "https://registry.yarnpkg.com/marked/-/marked-5.1.2.tgz#62b5ccfc75adf72ca3b64b2879b551d89e77677f"
-  integrity sha512-ahRPGXJpjMjwSOlBoTMZAK7ATXkli5qCPxZ21TG44rx1KEo44bii4ekgTDQPNRQ4Kh7JMb9Ub1PVk1NxRSsorg==
+marked@17.0.1:
+  version "17.0.1"
+  resolved "https://registry.yarnpkg.com/marked/-/marked-17.0.1.tgz#9db34197ac145e5929572ee49ef701e37ee9b2e6"
+  integrity sha512-boeBdiS0ghpWcSwoNm/jJBwdpFaMnZWRzjA6SkUMYb40SVaN1x7mmfGKp0jvexGcx+7y2La5zRZsYFZI6Qpypg==
 
 md5.js@^1.3.4:
   version "1.3.5"
diff --git a/pyproject.toml b/pyproject.toml
index 2a27db38..7a85ed64 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,7 +26,8 @@ dependencies = [
   "Levenshtein>=0.21.1",
   "matplotlib>=3.7.1",
   "ml-collections>=0.1.1",
-  "numpy>=1.24.1,<2.0.0",
+  "numpy==2.0.0",
+  "numba>=0.60.0",
   "pandas[output-formatting]>=2.0.3",
   "Pillow>=10.0.0",
   "portpicker>=1.5.2",
@@ -34,9 +35,9 @@ dependencies = [
   "rouge-score>=0.1.2",
   "sacrebleu>=2.3.1",
   "saliency>=0.1.3",
-  "scikit-learn>=1.0.2",
+  "scikit-learn>=1.6.1",
   "scipy>=1.10.1",
-  "shap>=0.42.0,<0.46.0",
+  "shap>=0.46.0",
   "six>=1.16.0",
   "termcolor>=2.3.0",
   "tqdm>=4.64.0",
@@ -82,7 +83,7 @@ keywords = [
 examples-common = [
   "gunicorn>=20.1.0",
   "tensorflow>=2.16.0",
-  "transformers>=4.27.1",
+  "transformers>=4.27.1,<5.0.0",
 ]
 # LINT.ThenChange(./requirements_examples_common.txt)
 # LINT.IfChange
diff --git a/requirements.txt b/requirements.txt
index 9c4707b0..eb0d6976 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -23,7 +23,8 @@ ipython>=7.34.0
 Levenshtein>=0.21.1
 matplotlib>=3.6.0,<3.9.0
 ml-collections>=0.1.1
-numpy>=1.24.1,<2.0.0
+numpy==2.0.0
+numba>=0.60.0
 pandas[output-formatting]>=2.0.3
 Pillow>=10.0.0
 portpicker>=1.5.2
@@ -31,9 +32,9 @@ requests>=2.31.0
 rouge-score>=0.1.2
 sacrebleu>=2.3.1
 saliency>=0.1.3
-scikit-learn>=1.0.2
+scikit-learn>=1.6.1
 scipy>=1.10.1
-shap>=0.42.0,<0.46.0
+shap>=0.46.0
 six>=1.16.0
 termcolor>=2.3.0
 tqdm>=4.64.0
diff --git a/requirements_examples_common.txt b/requirements_examples_common.txt
index 511505c9..02e7ac3b 100644
--- a/requirements_examples_common.txt
+++ b/requirements_examples_common.txt
@@ -18,5 +18,5 @@
 
 gunicorn>=20.1.0
 tensorflow>=2.16.0
-transformers>=4.27.1
+transformers>=4.27.1,<5.0.0
 # LINT.ThenChange(./pyproject.toml)
diff --git a/website/sphinx_src/demos.md b/website/sphinx_src/demos.md
index 448c0d23..5be0910b 100644
--- a/website/sphinx_src/demos.md
+++ b/website/sphinx_src/demos.md
@@ -83,6 +83,17 @@ Generative AI Toolkit.
 
 --------------------------------------------------------------------------------
 
+## Text To Image Demo <!-- DO NOT REMOVE {#text-to-image .demo-section-header} -->
+
+### min(DALL·E) <!-- DO NOT REMOVE {#min-dall-e .demo-header} -->
+
+**Code:** [examples/dalle_mini/demo.py](https://github.com/PAIR-code/lit/blob/main/lit_nlp/examples/dalle_mini/demo.py)
+
+*   Support text to image generation in LIT using
+    [min(DALL·E)](https://github.com/kuprel/min-dalle) model.
+
+--------------------------------------------------------------------------------
+
 ## Multimodal <!-- DO NOT REMOVE {#multimodal .demo-section-header} -->
 
 ### Tabular Data: Penguin Classification <!-- DO NOT REMOVE {#penguin .demo-header} -->