Visual indication when loading models

Author: jmbejar

src/chat/mod.rs

@@ -10,6 +10,7 @@ pub mod delete_chat_modal;
 pub mod model_info;
 pub mod model_selector;
 pub mod model_selector_list;
+pub mod model_selector_loading;
 pub mod shared;
 
 use makepad_widgets::Cx;
@@ -25,6 +26,7 @@ pub fn live_design(cx: &mut Cx) {
     model_info::live_design(cx);
     model_selector_list::live_design(cx);
     model_selector::live_design(cx);
+    model_selector_loading::live_design(cx);
     shared::live_design(cx);
     delete_chat_modal::live_design(cx);
     chat_history_card_options::live_design(cx);

src/chat/model_selector.rs

@@ -4,7 +4,7 @@ use crate::{
 };
 use makepad_widgets::*;
 
-use super::model_selector_list::{ModelSelectorAction, ModelSelectorListWidgetExt};
+use super::{model_selector_list::{ModelSelectorAction, ModelSelectorListWidgetExt}, model_selector_loading::ModelSelectorLoadingWidgetExt};
 
 live_design! {
     import makepad_widgets::base::*;
@@ -14,13 +14,15 @@ live_design! {
 
     import crate::chat::model_info::ModelInfo;
     import crate::chat::model_selector_list::ModelSelectorList;
+    import crate::chat::model_selector_loading::ModelSelectorLoading;
 
     ModelSelectorButton = <RoundedView> {
         width: Fill,
         height: 54,
+        flow: Overlay,
 
         align: {x: 0.0, y: 0.5},
-        padding: 16,
+        padding: 0,
 
         draw_bg: {
             instance radius: 3.0,
@@ -29,11 +31,18 @@ live_design! {
 
         cursor: Hand,
 
+        loading = <ModelSelectorLoading> {
+            width: Fill,
+            height: Fill,
+            visible: false,
+        }
+
         choose = <View> {
             width: Fill,
             height: Fit,
 
             align: {x: 0.5, y: 0.5},
+            padding: 16,
 
             label = <Label> {
                 draw_text:{
@@ -49,6 +58,8 @@ live_design! {
             show_bg: false,
             visible: false,
 
+            padding: 16,
+
             label = {
                 draw_text: {
                     text_style: <BOLD_FONT>{font_size: 11},
@@ -139,6 +150,41 @@ impl Widget for ModelSelector {
         self.view.handle_event(cx, event, scope);
         self.widget_match_event(cx, event, scope);
 
+        let store = scope.data.get::<Store>().unwrap();
+
+        if let Hit::FingerDown(fd) =
+            event.hits_with_capture_overload(cx, self.view(id!(button)).area(), true)
+        {
+            if !options_to_display(store) {
+                return;
+            };
+            if fd.tap_count == 1 {
+                self.open = !self.open;
+
+                if self.open {
+                    let list = self.model_selector_list(id!(options.list_container.list));
+                    let height = list.get_height();
+                    if height > MAX_OPTIONS_HEIGHT {
+                        self.options_list_height = Some(MAX_OPTIONS_HEIGHT);
+                    } else {
+                        self.options_list_height = Some(height);
+                    }
+
+                    self.view(id!(options)).apply_over(
+                        cx,
+                        live! {
+                            height: Fit,
+                        },
+                    );
+
+                    self.animator_play(cx, id!(open.show));
+                } else {
+                    self.hide_animation_timer = cx.start_timeout(0.3);
+                    self.animator_play(cx, id!(open.hide));
+                }
+            }
+        }
+
         if self.hide_animation_timer.is_event(event).is_some() {
             // When closing animation is done, hide the wrapper element
             self.view(id!(options)).apply_over(cx, live! { height: 0 });
@@ -181,6 +227,7 @@ impl Widget for ModelSelector {
                     }
                 },
             );
+            self.model_selector_loading(id!(loading)).hide(cx);
         } else if no_active_model(store) {
             choose_label.set_text("Choose a Model");
             let color = vec3(0.0, 0.0, 0.0);
@@ -192,6 +239,7 @@ impl Widget for ModelSelector {
                     }
                 },
             );
+            self.model_selector_loading(id!(loading)).hide(cx);
         } else {
             self.update_selected_model_info(cx, store);
         }
@@ -260,6 +308,25 @@ impl ModelSelector {
     }
 
     fn update_selected_model_info(&mut self, cx: &mut Cx, store: &Store) {
+        if let Some(loader) = &store.chats.model_loader {
+            if !loader.complete {
+                let caption = format!("Loading {}", loader.file.name);
+                self.model_selector_loading(id!(loading)).show_and_animate(cx);
+                self.view(id!(selected)).apply_over(
+                    cx,
+                    live! {
+                        visible: true
+                        label = { text: (caption) }
+                        architecture_tag = { visible: false }
+                        params_size_tag = { visible: false }
+                        file_size_tag = { visible: false }
+                    },
+                );
+                self.redraw(cx);
+                return;
+            }
+        }
+
         let Some(downloaded_file) = store.get_loaded_downloaded_file() else {
             return;
         };
@@ -281,6 +348,7 @@ impl ModelSelector {
         let size = format_model_size(&downloaded_file.file.size).unwrap_or("".to_string());
         let size_visible = !size.trim().is_empty();
 
+        self.model_selector_loading(id!(loading)).hide(cx);
         self.view(id!(selected)).apply_over(
             cx,
             live! {

src/chat/model_selector_loading.rs

@@ -0,0 +1,133 @@
+use makepad_widgets::*;
+
+live_design! {
+    import makepad_widgets::base::*;
+    import makepad_widgets::theme_desktop_dark::*;
+    import makepad_draw::shader::std::*;
+
+    import crate::shared::styles::*;
+    import crate::shared::widgets::*;
+    import crate::landing::model_card::ModelCard;
+
+    ANIMATION_SPEED = 1.5;
+
+    Bar = <RoundedView> {
+        width: Fill,
+        height: 8,
+        show_bg: true,
+        draw_bg: {
+            instance radius: 1.0,
+            instance dither: 0.9
+
+            fn get_color(self) -> vec4 {
+                return mix(
+                    #F3FFA2,
+                    #E3FBFF,
+                    self.pos.x + self.dither
+                )
+            }
+
+            fn pixel(self) -> vec4 {
+                let sdf = Sdf2d::viewport(self.pos * self.rect_size)
+                sdf.box(
+                    self.inset.x + self.border_width,
+                    self.inset.y + self.border_width,
+                    self.rect_size.x - (self.inset.x + self.inset.z + self.border_width * 2.0),
+                    self.rect_size.y - (self.inset.y + self.inset.w + self.border_width * 2.0),
+                    max(1.0, self.radius)
+                )
+                sdf.fill_keep(self.get_color())
+                if self.border_width > 0.0 {
+                    sdf.stroke(self.get_border_color(), self.border_width)
+                }
+                return sdf.result;
+            }
+        }
+    }
+
+    ModelSelectorLoading = {{ModelSelectorLoading}} {
+        width: Fill,
+        height: Fill,
+        align: {x: 0, y: 1},
+
+        line = <Bar> {}
+
+        animator: {
+            line = {
+                default: restart,
+                restart = {
+                    redraw: true,
+                    from: {all: Forward {duration: (ANIMATION_SPEED)}}
+                    apply: {line = { draw_bg: {dither: 0.6} }}
+                }
+                run = {
+                    redraw: true,
+                    from: {all: Forward {duration: (ANIMATION_SPEED)}}
+                    apply: {line = { draw_bg: {dither: -0.3} }}
+                }
+            }
+        }
+    }
+}
+
+#[derive(Live, LiveHook, Widget)]
+pub struct ModelSelectorLoading {
+    #[deref]
+    view: View,
+
+    #[animator]
+    animator: Animator,
+
+    #[rust]
+    timer: Timer,
+}
+
+impl Widget for ModelSelectorLoading {
+    fn handle_event(&mut self, cx: &mut Cx, event: &Event, scope: &mut Scope) {
+        if self.timer.is_event(event).is_some() {
+            self.update_animation(cx);
+        }
+        if self.animator_handle_event(cx, event).must_redraw() {
+            self.redraw(cx);
+        }
+
+        self.view.handle_event(cx, event, scope);
+    }
+
+    fn draw_walk(&mut self, cx: &mut Cx2d, scope: &mut Scope, walk: Walk) -> DrawStep {
+        self.view.draw_walk(cx, scope, walk)
+    }
+}
+
+impl ModelSelectorLoading {
+    pub fn update_animation(&mut self, cx: &mut Cx) {
+        if self.animator_in_state(cx, id!(line.restart)) {
+            self.visible = true;
+            self.animator_play(cx, id!(line.run));
+            self.timer = cx.start_timeout(1.5);
+        } else {
+            self.visible = true;
+            self.animator_play(cx, id!(line.restart));
+            self.timer = cx.start_timeout(1.5);
+        }
+    }
+}
+
+impl ModelSelectorLoadingRef {
+    pub fn show_and_animate(&mut self, cx: &mut Cx) {
+        let Some(mut inner) = self.borrow_mut() else {
+            return;
+        };
+        if inner.timer.is_empty() {
+            inner.timer = cx.start_timeout(0.2);
+        }
+    }
+
+    pub fn hide(&mut self, cx: &mut Cx) {
+        let Some(mut inner) = self.borrow_mut() else {
+            return;
+        };
+        inner.view.visible = false;
+        inner.timer = Timer::default();
+    }
+}

src/data/chats/mod.rs

@@ -1,19 +1,24 @@
 pub mod chat;
+pub mod model_loader;
 
 use anyhow::{Context, Result};
 use chat::{Chat, ChatID};
+use model_loader::ModelLoader;
 use moxin_backend::Backend;
-use moxin_protocol::protocol::{Command, LoadModelResponse};
-use moxin_protocol::{data::*, protocol::LoadModelOptions};
+use moxin_protocol::data::*;
+use moxin_protocol::protocol::Command;
 use std::fs;
-use std::{cell::RefCell, path::PathBuf, rc::Rc, sync::mpsc::channel};
+use std::sync::mpsc::channel;
+use std::{cell::RefCell, path::PathBuf, rc::Rc};
 
 use super::filesystem::setup_chats_folder;
 
 pub struct Chats {
     pub backend: Rc<Backend>,
     pub saved_chats: Vec<RefCell<Chat>>,
+
     pub loaded_model: Option<File>,
+    pub model_loader: Option<ModelLoader>,
 
     current_chat_id: Option<ChatID>,
     chats_dir: PathBuf,
@@ -26,6 +31,7 @@ impl Chats {
             saved_chats: Vec::new(),
             current_chat_id: None,
             loaded_model: None,
+            model_loader: None,
             chats_dir: setup_chats_folder(),
         }
     }
@@ -51,41 +57,28 @@ impl Chats {
             .map(|c| c.borrow().id)
     }
 
-    pub fn load_model(&mut self, file: &File) -> Result<()> {
-        let (tx, rx) = channel();
-        let cmd = Command::LoadModel(
-            file.id.clone(),
-            LoadModelOptions {
-                prompt_template: None,
-                gpu_layers: moxin_protocol::protocol::GPULayers::Max,
-                use_mlock: false,
-                rope_freq_scale: 0.0,
-                rope_freq_base: 0.0,
-                context_overflow_policy:
-                    moxin_protocol::protocol::ContextOverflowPolicy::StopAtLimit,
-            },
-            tx,
-        );
-
-        self.backend.as_ref().command_sender.send(cmd).unwrap();
-
-        if let Ok(response) = rx.recv() {
-            match response {
-                Ok(LoadModelResponse::Completed(_)) => {
-                    self.loaded_model = Some(file.clone());
-                    Ok(())
-                }
-                Ok(_) => {
-                    eprintln!("Error loading model: Unexpected response");
-                    Err(anyhow::anyhow!("Error loading model: Unexpected response"))
-                }
-                Err(err) => {
-                    eprintln!("Error loading model: {:?}", err);
-                    Err(err)
-                }
+    pub fn load_model(&mut self, file: &File) {
+        if let Some(loader) = &self.model_loader {
+            if !loader.complete {
+                return;
+            }
+        }
+
+        let loader = ModelLoader::new(file.clone());
+        loader.load_model(self.backend.as_ref());
+        self.model_loader = Some(loader);
+    }
+
+    pub fn update_load_model(&mut self) {
+        let loader = self.model_loader.as_mut();
+        if let Some(loader) = loader {
+            if loader.check_load_response().is_ok() {
+                self.loaded_model = Some(loader.file.clone());
+            }
+
+            if loader.complete {
+                self.model_loader = None;
             }
-        } else {
-            Err(anyhow::anyhow!("Error loading model"))
         }
     }