Enhance draft model loading and response handling

- Update load_draft function to check for local directory and config.json before downloading.
- Refactor response handling in _do_generate to simplify text accumulation and ensure proper chunking.

This improves the efficiency of loading draft models and enhances the streaming response mechanism.

Author: Vect0rM

dflash/model_mlx.py

@@ -130,7 +130,11 @@ def load(model_id: str):
 
 
 def load_draft(draft_id: str) -> DFlashDraftModel:
-    path = Path(snapshot_download(draft_id, allow_patterns=["*.safetensors", "*.json"]))
+    local = Path(draft_id)
+    if local.is_dir() and (local / "config.json").exists():
+        path = local
+    else:
+        path = Path(snapshot_download(draft_id, allow_patterns=["*.safetensors", "*.json"]))
     cfg = json.loads((path / "config.json").read_text())
     config = DFlashConfig(
         hidden_size=cfg["hidden_size"],

dflash/server.py

@@ -151,12 +151,11 @@ def _sync_worker() -> None:
                     ):
                         if _cancel_event.is_set():
                             break
-                        new_text = resp.text
-                        if new_text and new_text != text_so_far:
-                            delta = new_text[len(text_so_far) :]
-                            text_so_far = new_text
+                        delta_text = resp.text
+                        if delta_text:
+                            text_so_far += delta_text
                             gen_tokens.extend(resp.tokens)
-                            _put({"type": "chunk", "content": delta})
+                            _put({"type": "chunk", "content": delta_text})
 
                         if resp.finish_reason is not None:
                             elapsed = time.perf_counter() - tic
@@ -422,10 +421,10 @@ def _load_models(args: argparse.Namespace) -> None:
     log.info("Target model loaded successfully")
 
     if args.draft_model:
-        from dflash.model_mlx import DFlashDraftModel
+        from dflash.model_mlx import load_draft
 
         log.info("Loading DFlash draft model: %s", args.draft_model)
-        _draft = DFlashDraftModel.from_pretrained(args.draft_model)
+        _draft = load_draft(args.draft_model)
         log.info("DFlash draft model loaded (block_size=%d)", _draft.config.block_size)
 
     _model_id = args.model_id or model_path.stem