update

Dockerfile

@@ -80,4 +80,4 @@ EXPOSE 5003
 CMD ["bash", "demo/demo.sh"]
 
 # DOCKER_BUILDKIT=1 docker build --ssh default --network=host -t unidisc .
-# docker run --network=host -it -p 5003:5003 unidisc
+# docker run --gpus all --network=host -it -p 5003:5003 unidisc

demo/client.py

@@ -1,21 +1,37 @@
+print(f"Import 00")
 from fasthtml.common import *
 from fasthtml.svg import *
 from monsterui.all import *
+print(f"Import 01")
 from pathlib import Path
 import requests
 import base64
 from PIL import Image
+print(f"Import 02")
 import numpy as np
 import io
 import json
 
-print(f"Finished import 0")
+
 
 SHOW_DEV_BUTTONS = True
 DEMO_DIR = Path(__file__).parent
 ADD_DEV_FORM = True
 
-print(f"Finished import 1. Demo dir: {DEMO_DIR}")
+def list_ckpt_files():
+    """List all files in the ckpts directory."""
+    ckpt_dir = Path("ckpts").absolute()
+    files = []
+    
+    if ckpt_dir.exists():
+        for file_path in ckpt_dir.glob("**/*"):
+            if file_path.is_file() or file_path.is_symlink():
+                files.append(str(file_path))
+    
+    return sorted(files)
+
+print(f"CKPT files: {list_ckpt_files()}")
+
 
 DEMOS = [
     {

demo/demo.sh

@@ -1,6 +1,11 @@
 #!/bin/bash
 
-UNIDISC_FORCE_CUDNN_SPDA_CONTEXT=1 uv run python demo/server.py experiments='[large_scale_train,large_scale_train_high_res_interleaved,eval_unified,large_scale_high_res_interleaved_inference]' \
-trainer.load_from_state_dict="/home/appuser/app/ckpts/unidisc_interleaved/unidisc_interleaved.pt" &
+ls -la /home/appuser/app/ckpts
+ls -la /home/appuser/app/ckpts/unidisc_interleaved
 
-uv run python demo/client.py
+PYTHONUNBUFFERED=1 uv run python demo/client.py &
+
+PYTHONUNBUFFERED=1 UNIDISC_FORCE_CUDNN_SPDA_CONTEXT=1 uv run python demo/server.py experiments='[large_scale_train,large_scale_train_high_res_interleaved,eval_unified,large_scale_high_res_interleaved_inference]' \
+trainer.load_from_state_dict="/home/appuser/app/ckpts/unidisc_interleaved/unidisc_interleaved.pt" model.use_custom_vae_ckpt="/home/appuser/app/ckpts/unidisc_interleaved/vq_ds16_t2i.pt" &
+
+wait