First working demo

app.py

@@ -9,6 +9,10 @@ import PIL.Image
 import gradio as gr
 import app_novic
 
+# Sample images
+IMAGE_EXTS = ('jpg', 'jpeg', 'png', 'webp')
+SAMPLE_IMAGES = sorted(image_path for image_ext in IMAGE_EXTS for image_path in glob.glob(os.path.join('sample_images', f'*.{image_ext}')))
+
 # Model checkpoints
 MODEL_CHECKPOINTS = {
 	'NOVIC SigLIP B/16 FT2': os.path.join('ovod_20240610_105233', 'ovod_chunk0899_20240612_005748.train'),
@@ -16,17 +20,20 @@ MODEL_CHECKPOINTS = {
 	'NOVIC DFN-5B H/14-378 FT2': os.path.join('ovod_20240620_162925', 'ovod_chunk0899_20240621_202727.train'),
 	'NOVIC DFN-5B H/14-378 FT0': os.path.join('ovod_20240628_142131', 'ovod_chunk0433_20240630_235415.train'),
 }
+DEFAULT_MODEL = 'NOVIC DFN-5B H/14-378 FT0'
 
-# Sample images
-IMAGE_EXTS = ('jpg', 'jpeg', 'png', 'webp')
-SAMPLE_IMAGES = sorted(image_path for image_ext in IMAGE_EXTS for image_path in glob.glob(os.path.join('sample_images', f'*.{image_ext}')))
+# Get a checkpoint path
+def get_checkpoint(model: str) -> str:
+	return os.path.join('checkpoints', MODEL_CHECKPOINTS[model])
+
+# Ensure the default model is preloaded
+app_novic.get_model(checkpoint=get_checkpoint(model=DEFAULT_MODEL))
 
 # Classify an image
 def classify_image(image: Optional[PIL.Image.Image], model: Optional[str]) -> dict[str, float]:
 	if image is None or model is None:
 		return {}
-	checkpoint = os.path.join('checkpoints', MODEL_CHECKPOINTS[model])
-	return app_novic.classify_image(image=image, checkpoint=checkpoint)
+	return app_novic.classify_image(image=image, checkpoint=get_checkpoint(model=model))
 
 # Gradio UI
 with gr.Blocks(
@@ -43,7 +50,7 @@ with gr.Blocks(
 		with gr.Column(scale=1):
 			input_model = gr.Dropdown(
 				choices=list(MODEL_CHECKPOINTS),
-				value='NOVIC DFN-5B H/14-378 FT0',
+				value=DEFAULT_MODEL,
 				type='value',
 				multiselect=False,
 				allow_custom_value=False,

app_novic.py

@@ -2,11 +2,10 @@
 
 # Imports
 import sys
-import random  # TODO: TEMP
 import pathlib
 import subprocess
 import PIL.Image
-# TODO: Then start coding infer.py (ANY additional import has to be clarified with novic_spaces/requirements.txt (pip install -r requirements.txt) and novic/requirements-infer.txt
+import torch
 
 # Ensure access to NOVIC code
 REPO_ROOT = pathlib.Path(__file__).parent.resolve()
@@ -20,16 +19,38 @@ if not NOVIC_TEST.exists():
 if str(NOVIC_DIR) not in sys.path:
 	sys.path.insert(0, str(NOVIC_DIR))
 
-# NOVIC imports
-import infer  # TODO: TEMP
-print(infer)  # TODO: TEMP
+# Import NOVIC inference
+import infer  # noqa
+
+# General inference setup
+infer.utils.allow_tf32(enable=True)
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+MODELS: dict[str, infer.NOVICModel] = {}
+
+# Retrieve a NOVIC model
+def get_model(checkpoint: str) -> infer.NOVICModel:
+
+	if (model := MODELS.get(checkpoint, None)) is not None:
+		return model
+
+	model = infer.NOVICModel(
+		checkpoint=checkpoint,
+		gencfg='beam_k10_vnone_gp_t1_a0',
+		guide_targets=None,
+		torch_compile=False,
+		batch_size=1,
+		device=DEVICE,
+		cfg_flat_override=None,
+		embedder_override=None,
+	)
+
+	model.__enter__()
+	MODELS[checkpoint] = model
+	return model
 
 # Classify an image
 def classify_image(image: PIL.Image.Image, checkpoint: str) -> dict[str, float]:
-	# TODO: TEMP
-	print(image, checkpoint)
-	labels = ('thing', 'item', 'stuff', 'object', 'entity')
-	probs = [random.random() for _ in labels]
-	prob_sum = sum(probs)
-	return {label: prob / prob_sum for label, prob in zip(labels, probs)}
+	model = get_model(checkpoint=checkpoint)
+	output = model.classify_image(image=image)
+	return dict(zip(output.preds[0], output.probs[0], strict=True))
 # EOF

novic

@@ -1 +1 @@
-Subproject commit cc56882720618202830c334a2d93d8d536cfd412
+Subproject commit 23ab1652b174b0f45bf34139c4d85f45fd28dead