abetlen commited on
Commit
f78095a
1 Parent(s): a9dabe2
Files changed (1) hide show
  1. models.py +15 -7
models.py CHANGED
@@ -9,6 +9,8 @@ import PIL.Image
9
  # pylint: disable=g-bad-import-order
10
  import gradio_helpers
11
 
 
 
12
 
13
  ORGANIZATION = 'abetlen'
14
  BASE_MODELS = [
@@ -69,10 +71,16 @@ def generate(
69
  model_name: str, sampler: str, image: PIL.Image.Image, prompt: str
70
  ) -> str:
71
  """Generates output with specified `model_name`, `sampler`."""
72
- model, params_cpu = get_cached_model(model_name)
73
- batch = model.shard_batch(model.prepare_batch([image], [prompt]))
74
- with gradio_helpers.timed('sharding'):
75
- params = model.shard_params(params_cpu)
76
- with gradio_helpers.timed('computation', start_message=True):
77
- tokens = model.predict(params, batch, sampler=sampler)
78
- return model.tokenizer.to_str(tokens[0])
 
 
 
 
 
 
 
9
  # pylint: disable=g-bad-import-order
10
  import gradio_helpers
11
 
12
+ import llama_cpp
13
+
14
 
15
  ORGANIZATION = 'abetlen'
16
  BASE_MODELS = [
 
71
  model_name: str, sampler: str, image: PIL.Image.Image, prompt: str
72
  ) -> str:
73
  """Generates output with specified `model_name`, `sampler`."""
74
+ # model, params_cpu = get_cached_model(model_name)
75
+ # batch = model.shard_batch(model.prepare_batch([image], [prompt]))
76
+ # with gradio_helpers.timed('sharding'):
77
+ # params = model.shard_params(params_cpu)
78
+ # with gradio_helpers.timed('computation', start_message=True):
79
+ # tokens = model.predict(params, batch, sampler=sampler)
80
+
81
+ model_path = gradio_helpers.get_paths()[model_name]
82
+ model = llama_cpp.Llama(
83
+ model_path,
84
+ )
85
+ # return model.tokenizer.to_str(tokens[0])
86
+ return "output"