Issues with ZeroGPU

app.py

@@ -1,9 +1,7 @@
 import gradio as gr
-from mosaic import Mosaic  # adjust import as needed
+from mosaic import Mosaic 
 import spaces
 import traceback
-from transformers import AutoModelForCausalLM
-import torch
 
 # Maximum number of model textboxes
 MAX_MODELS = 10
@@ -28,9 +26,7 @@ def update_textboxes(n_visible):
     if n_visible < MAX_MODELS:
         n_visible += 1
     tb_updates = [gr.update(visible=(i < n_visible)) for i in range(MAX_MODELS)]
-    btn_updates = [gr.update(visible=(i < n_visible)) for i in range(MAX_MODELS)]
-    status_updates = [gr.update(visible=(i < n_visible)) for i in range(MAX_MODELS)]
-    return (n_visible, *tb_updates, *btn_updates, *status_updates)
+    return (n_visible, *tb_updates)
 
 # Decrease model slots and clear removed entries
 def remove_textboxes(n_visible):
@@ -41,17 +37,13 @@ def remove_textboxes(n_visible):
     # Remove cached models for slots now hidden
     for idx in range(new, old):
         LOADED_MODELS.pop(idx+1, None)
-    tb_updates, btn_updates, status_updates = [], [], []
+    tb_updates = []
     for i in range(MAX_MODELS):
         if i < n_visible:
             tb_updates.append(gr.update(visible=True))
-            btn_updates.append(gr.update(visible=True))
-            status_updates.append(gr.update(visible=True))
         else:
             tb_updates.append(gr.update(visible=False, value=""))
-            btn_updates.append(gr.update(visible=False))
-            status_updates.append(gr.update(visible=False, value="Not loaded"))
-    return (n_visible, *tb_updates, *btn_updates, *status_updates)
+    return (n_visible, *tb_updates)
 
 def apply_config1():
     """
@@ -61,23 +53,19 @@ def apply_config1():
       - new visibility for each Load button & status box
     """
     n_vis = len(GPT_CONFIG_MODELS)
-    tb_updates, btn_updates, status_updates = [], [], []
+    tb_updates = []
 
     for i in range(MAX_MODELS):
         if i < n_vis:
             # show this slot, set its value from CONFIG_MODELS
             tb_updates.append(gr.update(visible=True, value=GPT_CONFIG_MODELS[i]))
-            btn_updates.append(gr.update(visible=True))
-            status_updates.append(gr.update(visible=True, value="Not loaded"))
         else:
             # hide all others
             tb_updates.append(gr.update(visible=False, value=""))
-            btn_updates.append(gr.update(visible=False))
-            status_updates.append(gr.update(visible=False, value="Not loaded"))
 
     # Return in the same shape as your update_textboxes/remove_textboxes:
     #   (n_models_state, *all textboxes, *all load buttons, *all status boxes)
-    return (n_vis, *tb_updates, *btn_updates, *status_updates)
+    return (n_vis, *tb_updates)
 
 def apply_config2():
     """
@@ -87,53 +75,20 @@ def apply_config2():
       - new visibility for each Load button & status box
     """
     n_vis = len(Falcon_CONFIG_MODELS)
-    tb_updates, btn_updates, status_updates = [], [], []
+    tb_updates = []
 
     for i in range(MAX_MODELS):
         if i < n_vis:
             # show this slot, set its value from CONFIG_MODELS
             tb_updates.append(gr.update(visible=True, value=Falcon_CONFIG_MODELS[i]))
-            btn_updates.append(gr.update(visible=True))
-            status_updates.append(gr.update(visible=True, value="Not loaded"))
         else:
             # hide all others
             tb_updates.append(gr.update(visible=False, value=""))
-            btn_updates.append(gr.update(visible=False))
-            status_updates.append(gr.update(visible=False, value="Not loaded"))
 
     # Return in the same shape as your update_textboxes/remove_textboxes:
     #   (n_models_state, *all textboxes, *all load buttons, *all status boxes)
-    return (n_vis, *tb_updates, *btn_updates, *status_updates)
+    return (n_vis, *tb_updates)
 
-# Load a single model and report status
-@spaces.GPU()
-def load_single_model(model_path, use_bfloat16=True):
-    try:
-        repo = model_path
-        if not repo:
-            return "Error: No path provided"
-        if repo in LOADED_MODELS:
-            return "Loaded"
-        # actual load; may raise
-        model = AutoModelForCausalLM.from_pretrained(
-                    repo,
-                    device_map="auto",
-                    trust_remote_code=True,
-                    torch_dtype=torch.bfloat16 if use_bfloat16 else torch.float32,
-                )
-        model.eval()
-        LOADED_MODELS[repo] = model
-        return "Loaded"
-    except Exception as e:
-        return f"Error loading model: {e}"
-
-# Determine interactive state for Run button
-def check_all_loaded(n_visible, *status_texts):
-    # status_texts are strings: "Loaded" indicates success
-    needed = status_texts[:n_visible]
-    if all(s == "Loaded" for s in needed):
-        return gr.update(interactive=True)
-    return gr.update(interactive=False)
 
 @spaces.GPU()
 def run_scoring(input_text, *args):
@@ -148,7 +103,7 @@ def run_scoring(input_text, *args):
         if len(models) < 2:
             return "Please enter at least two model paths.", None, None
         threshold = 0.0 if threshold_choice == "default" else custom_threshold
-        mosaic_instance = Mosaic(model_name_or_paths=models, one_model_mode=False, loaded_models=LOADED_MODELS)
+        mosaic_instance = Mosaic(model_name_or_paths=models, one_model_mode=False)
         final_score = mosaic_instance.compute_end_score(input_text)
         msg = "This text was probably generated." if final_score < threshold else "This text is likely human-written."
         return msg, final_score, threshold
@@ -166,20 +121,11 @@ with demo:
         gr.Markdown("**⚠️ Please make sure all models have the same tokenizer or it won’t work.**")
         gr.Markdown("### Model Paths (at least 2 required)")
         n_models_state = gr.State(4)
-        model_inputs, load_buttons, status_boxes = [], [], []
+        model_inputs = []
         for i in range(1, MAX_MODELS+1):
             with gr.Row():
                 tb = gr.Textbox(label=f"Model {i} Path", value="" if i > 4 else None, visible=(i <= 4))
-                btn = gr.Button("Load", elem_id=f"load_{i}", visible=(i <= 4))
-                status = gr.Textbox(label="Loading status", value="Not loaded", interactive=False, visible=(i <= 4))
-                btn.click(
-                    fn=load_single_model,
-                    inputs=[tb, gr.State(i)],
-                    outputs=status
-                )
                 model_inputs.append(tb)
-                load_buttons.append(btn)
-                status_boxes.append(status)
         with gr.Row():
             plus = gr.Button("Add model slot", elem_id="plus_button")
             minus = gr.Button("Remove model slot", elem_id="minus_button")
@@ -188,31 +134,27 @@ with demo:
         plus.click(
             fn=update_textboxes,
             inputs=n_models_state,
-            outputs=[n_models_state, *model_inputs, *load_buttons, *status_boxes]
+            outputs=[n_models_state, *model_inputs]
         )
         minus.click(
             fn=remove_textboxes,
             inputs=n_models_state,
-            outputs=[n_models_state, *model_inputs, *load_buttons, *status_boxes]
+            outputs=[n_models_state, *model_inputs]
         )
         config1_btn.click(
             fn=apply_config1,
-            inputs=None,                       # no inputs needed
-            outputs=[                          # must match order:
-                n_models_state,                # 1️⃣ the new visible‑count State
-                *model_inputs,                 # 2️⃣ your list of 10 Textboxes
-                *load_buttons,                 # 3️⃣ your list of 10 Load Buttons
-                *status_boxes                  # 4️⃣ your list of 10 Status Textboxes
+            inputs=None,                   
+            outputs=[                        
+                n_models_state,              
+                *model_inputs                 
             ]
         )
         config2_btn.click(
             fn=apply_config2,
-            inputs=None,                       # no inputs needed
-            outputs=[                          # must match order:
-                n_models_state,                # 1️⃣ the new visible‑count State
-                *model_inputs,                 # 2️⃣ your list of 10 Textboxes
-                *load_buttons,                 # 3️⃣ your list of 10 Load Buttons
-                *status_boxes                  # 4️⃣ your list of 10 Status Textboxes
+            inputs=None,                      
+            outputs=[                         
+                n_models_state,               
+                *model_inputs          
             ]
         )
     with gr.Row():
@@ -222,20 +164,8 @@ with demo:
         output_message = gr.Textbox(label="Result Message")
         output_score = gr.Number(label="Final Score")
         output_threshold = gr.Number(label="Threshold Used")
-    gr.Markdown("**⚠️ All models need to be loaded before scoring.**")
-    run_button = gr.Button("Run Scoring", interactive=False)
-    # Enable Run button when all statuses reflect "Loaded"
-    for status in status_boxes:
-        status.change(
-            fn=check_all_loaded,
-            inputs=[n_models_state, *status_boxes],
-            outputs=run_button
-        )
-    n_models_state.change(
-        fn=check_all_loaded,
-        inputs=[n_models_state, *status_boxes],
-        outputs=run_button
-    )
+    gr.Markdown("**⚠️ All models need to be loaded for scoring, this can take time**")
+    run_button = gr.Button("Run Scoring")
     run_button.click(
         fn=run_scoring,
         inputs=[input_text, *model_inputs, threshold_choice, custom_threshold],

mosaic.py

@@ -57,41 +57,28 @@ class Mosaic(object):
         unigram: Optional[str] = None,
         custom_config: Optional[List[bool]] = None,
         stupid_mode: bool = False,
-        one_model_mode: bool = False,
-        # new optional argument: preloaded models dict
-        loaded_models: Optional[Dict[str, AutoModelForCausalLM]] = None,
-    ) -> None:
+        one_model_mode: bool = False
+        ) -> None:
         """
         If `loaded_models` is provided, re-use any entries matching
         model_name_or_paths; otherwise load and optionally register
         into that dict.
         """
         self.models = []
-        # ensure we have a dict to cache into if passed
-        cache = loaded_models if loaded_models is not None else {}
 
         for model_name_or_path in model_name_or_paths:
-            # reuse if already loaded
-            if loaded_models is not None and model_name_or_path in cache:
-                model = cache[model_name_or_path]
-            else:
-                print("Reloading a model was necessary, you probably messed up.")
-                # load from pre-trained hub or path
-                model = AutoModelForCausalLM.from_pretrained(
-                    model_name_or_path,
-                    device_map="auto",
-                    trust_remote_code=True,
-                    torch_dtype=torch.bfloat16 if use_bfloat16 else torch.float32,
-                )
-                model.eval()
-                # cache for reuse
-                if loaded_models is not None:
-                    cache[model_name_or_path] = model
+            # load from pre-trained hub or path
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name_or_path,
+                device_map="auto",
+                trust_remote_code=True,
+                torch_dtype=torch.bfloat16 if use_bfloat16 else torch.float32,
+            )
+            model.eval()
+
             self.models.append(model)
             print(f"Loaded model: {model_name_or_path}")
 
-        # store optional references
-        self.loaded_models = cache
         self.one_model_mode = one_model_mode
         
         if stupid_mode: