Modify to use Hugging Face ZeRO GPU

app.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 from huggingface_hub import list_models
 from typing import List
@@ -22,9 +23,7 @@ dataset = None
 def load_merit_dataset():
     global dataset
     if dataset is None:
-        dataset = load_dataset(
-            "de-Rodrigo/merit", name="en-digital-seq", split="train", num_proc=8
-        )
+        dataset = load_dataset("de-Rodrigo/merit", name="en-digital-seq", split="train")
     return dataset
 
 
@@ -42,6 +41,7 @@ def get_collection_models(tag: str) -> List[str]:
     return [model.modelId for model in models if tag in model.tags]
 
 
+@spaces.GPU
 def get_donut():
     global donut_model, donut_processor
     if donut_model is None or donut_processor is None:
@@ -50,28 +50,26 @@ def get_donut():
                 "de-Rodrigo/donut-merit"
             )
             donut_processor = DonutProcessor.from_pretrained("de-Rodrigo/donut-merit")
-            if torch.cuda.is_available():
-                donut_model = donut_model.to("cuda")
-            logger.info("Donut model loaded successfully")
+            donut_model = donut_model.to("cuda")
+            logger.info("Donut model loaded successfully on GPU")
         except Exception as e:
             logger.error(f"Error loading Donut model: {str(e)}")
             raise
     return donut_model, donut_processor
 
 
+@spaces.GPU
 def process_image_donut(model, processor, image):
     try:
         if not isinstance(image, Image.Image):
             image = Image.fromarray(image)
 
-        pixel_values = processor(image, return_tensors="pt").pixel_values
-        if torch.cuda.is_available():
-            pixel_values = pixel_values.to("cuda")
+        pixel_values = processor(image, return_tensors="pt").pixel_values.to("cuda")
 
         task_prompt = "<s_cord-v2>"
         decoder_input_ids = processor.tokenizer(
             task_prompt, add_special_tokens=False, return_tensors="pt"
-        )["input_ids"]
+        )["input_ids"].to("cuda")
 
         outputs = model.generate(
             pixel_values,
@@ -99,6 +97,7 @@ def process_image_donut(model, processor, image):
         return f"Error: {str(e)}"
 
 
+@spaces.GPU
 def process_image(model_name, image=None, dataset_image_index=None):
     if dataset_image_index is not None:
         image = get_image_from_dataset(dataset_image_index)

requirements.txt

@@ -1,7 +1,10 @@
-gradio
+gradio>=4.0.0
 transformers
 huggingface_hub
-torch
+torch>=2.0.0,!=2.3.0,<2.5.0
 numpy
 Pillow
-datasets
+datasets
+deepspeed
+huggingface-hub[cli,torch]
+huggingface_hub[spaces]