minor

app.py

@@ -14,6 +14,7 @@ import argparse
 from tokenizer_image.vq_model import VQ_models
 # from models.generate import generate
 from serve.llm import LLM 
+from serve.sampler import Sampler
 
 device = "cuda"
 
@@ -51,6 +52,7 @@ def load_model(args):
 
 
 def infer(cfg_scale, top_k, top_p, temperature, class_label, seed):
+    llm.llm_engine.model_executor.driver_worker.model_runner.model.sampler = Sampler(cfg_scale)
     args.cfg_scale = cfg_scale
     n = 4
     latent_size = image_size // args.downsample_size