Mirrors
/
EveryDream2trainer
mirror of https://github.com/victorchall/EveryDream2trainer.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

remove errant print from kosmos script

This commit is contained in:
Victor Hall 2023-11-03 13:50:47 -04:00
parent 4c3de29d43
commit 2512981956
1 changed files with 1 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
caption_kosmos2.py
View File

@ -24,10 +24,9 @@ import torch
from PIL import Image from PIL import Image
from pynvml import * from pynvml import *
from transformers import AutoProcessor, AutoModelForVision2Seq from transformers import AutoProcessor, AutoModelForVision2Seq
import colorama
GROUNDING = "<grounding>" GROUNDING = "<grounding>"
SUPPORTED_EXT = [".jpg", ".png", ".jpeg", ".bmp", ".jfif", ".webp"] SUPPORTED_EXT = [".jpg", ".png", ".jpeg", ".bmp", ".jfif", ".webp"]
def get_gpu_memory_map(): def get_gpu_memory_map():
@ -38,8 +37,6 @@ def get_gpu_memory_map():
return info.used/1024/1024 return info.used/1024/1024
def remove_starting_string(a, b): def remove_starting_string(a, b):
print(a)
print(b)
if b.startswith(a): if b.startswith(a):
return b[len(a):] # Remove string A from the beginning of string B return b[len(a):] # Remove string A from the beginning of string B
elif b.strip().startswith(a.strip()): elif b.strip().startswith(a.strip()):
@ -54,7 +51,6 @@ def main(args):
dtype=torch.float32 dtype=torch.float32
if not args.cpu: if not args.cpu:
#move to cuda and use float16
if args.dtype == "fp16": if args.dtype == "fp16":
dtype=torch.float16 dtype=torch.float16
elif args.dtype == "bf16": elif args.dtype == "bf16":