awacke1 commited on
Commit
a9c30d8
·
verified ·
1 Parent(s): d7c78a2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -198,7 +198,7 @@ class DiffusionBuilder:
198
  optimizer = torch.optim.AdamW(self.pipeline.unet.parameters(), lr=1e-5)
199
  self.pipeline.unet.train()
200
  for epoch in range(epochs):
201
- with—for st.spinner(f"Training diffusion epoch {epoch + 1}/{epochs}... ⚙️"):
202
  total_loss = 0
203
  for batch in dataloader:
204
  optimizer.zero_grad()
@@ -371,7 +371,8 @@ async def process_ocr(image, output_file):
371
  start_time = time.time()
372
  status = st.empty()
373
  status.text("Processing GOT-OCR2_0... (0s)")
374
- tokenizer, model = load_ocr_got()
 
375
  result = model.chat(tokenizer, image, ocr_type='ocr')
376
  elapsed = int(time.time() - start_time)
377
  status.text(f"GOT-OCR2_0 completed in {elapsed}s!")
 
198
  optimizer = torch.optim.AdamW(self.pipeline.unet.parameters(), lr=1e-5)
199
  self.pipeline.unet.train()
200
  for epoch in range(epochs):
201
+ with st.spinner(f"Training diffusion epoch {epoch + 1}/{epochs}... ⚙️"):
202
  total_loss = 0
203
  for batch in dataloader:
204
  optimizer.zero_grad()
 
371
  start_time = time.time()
372
  status = st.empty()
373
  status.text("Processing GOT-OCR2_0... (0s)")
374
+ tokenizer = AutoTokenizer.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True)
375
+ model = AutoModel.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True, torch_dtype=torch.float32).to("cpu").eval()
376
  result = model.chat(tokenizer, image, ocr_type='ocr')
377
  elapsed = int(time.time() - start_time)
378
  status.text(f"GOT-OCR2_0 completed in {elapsed}s!")