Test and fix errors in dalle_mini examples. Now it should works fine locally.

PiperOrigin-RevId: 758762358

Author: llcourage

lit_nlp/examples/dalle_mini/data.py

@@ -5,14 +5,19 @@
 
 
 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 = []
+    self._examples = []
     for prompt in prompts:
-      self.examples.append({"prompt": prompt})
+      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()}
-
-  def __iter__(self):
-    return iter(self.examples)

lit_nlp/examples/dalle_mini/demo.py

@@ -1,8 +1,42 @@
 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.
 """
@@ -26,8 +60,6 @@
 _FLAGS.set_default("development_demo", True)
 _FLAGS.set_default("default_layout", "DALLE_LAYOUT")
 
-_FLAGS.DEFINE_integer("grid_size", 4, "The grid size to use for the model.")
-
 _MODELS = (["dalle-mini"],)
 
 _CANNED_PROMPTS = ["I have a dream", "I have a shiba dog named cola"]

lit_nlp/examples/dalle_mini/model.py

@@ -97,12 +97,7 @@ def tensor_to_pil_image(tensor):
     return images
 
   def input_spec(self):
-    return {
-        "grid_size": lit_types.Scalar(),
-        "temperature": lit_types.Scalar(),
-        "top_k": lit_types.Scalar(),
-        "supercondition_factor": lit_types.Scalar(),
-    }
+    return {"prompt": lit_types.TextSegment()}
 
   def output_spec(self):
     return {

lit_nlp/examples/dalle_mini/requirements.txt

@@ -17,3 +17,5 @@
 
 # Dalle-Mini dependencies
 min_dalle==0.4.11
+torch==2.1.2+cu118
+--extra-index-url https://download.pytorch.org/whl/cu118