Update app.py

app.py

@@ -163,6 +163,11 @@ def model_main(args, master_port, rank):
 
 @torch.no_grad()
 def inference(args, infer_args, text_encoder, tokenizer, vae, model):
+    dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[
+        args.precision
+    ]
+    train_args = torch.load(os.path.join(args.ckpt, "model_args.pth"))
+    torch.cuda.set_device(0)
     with torch.autocast("cuda", dtype):
         while True:
             (