186 lines
6.3 KiB
Python
186 lines
6.3 KiB
Python
import dotenv
|
|
dotenv.load_dotenv()
|
|
|
|
import os
|
|
import sys
|
|
import torch
|
|
import datasets
|
|
import diffusers
|
|
import dotenv
|
|
import transformers
|
|
import argparse
|
|
from .hyperparams import MOONDREAM_REVISION
|
|
|
|
print(f"HF_HOME set to {os.getenv('HF_HOME')}")
|
|
|
|
# DATASET_SIZE = 10000
|
|
# ROWS_PER_DS = 1250
|
|
BATCH_SIZE = 4
|
|
PARQUET_BATCH_SIZE = 200
|
|
SKIP_PARQUET_BATCH = 203
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("-m", "--message")
|
|
|
|
args = parser.parse_args()
|
|
auth_token = os.getenv("HF_ACCESS_TOKEN")
|
|
if not auth_token:
|
|
print("huggingface access token not provided! please use the HF_ACCESS_TOKEN env var.")
|
|
sys.exit(1)
|
|
else:
|
|
print("huggingface access token loaded!")
|
|
|
|
tokenizer = transformers.AutoTokenizer.from_pretrained("vikhyatk/moondream2")
|
|
moondream = transformers.AutoModelForCausalLM.from_pretrained(
|
|
"vikhyatk/moondream2",
|
|
revision=MOONDREAM_REVISION,
|
|
trust_remote_code=True,
|
|
torch_dtype=torch.bfloat16,
|
|
device_map={"": "cuda"},
|
|
)
|
|
|
|
pipe = diffusers.StableDiffusion3Pipeline.from_pretrained(
|
|
"stabilityai/stable-diffusion-3.5-large",
|
|
torch_dtype=torch.bfloat16,
|
|
token=auth_token,
|
|
device_map="balanced",
|
|
)
|
|
|
|
def collate(batch):
|
|
images = []
|
|
keywords = []
|
|
|
|
for sample in batch:
|
|
images.append(sample["image"])
|
|
keywords.append([""])
|
|
|
|
return images, keywords
|
|
|
|
# flickr_dataset = datasets.load_dataset("nlphuji/flickr30k", split="test", streaming=True)\
|
|
# .select_columns(["image"])\
|
|
|
|
wiki_art_dataset = datasets.load_dataset("huggan/wikiart", split="train", streaming=True)\
|
|
.select_columns(["image"])
|
|
|
|
# anime_dataset_ft = datasets.Features({"image": datasets.Image(decode=True)})
|
|
# anime_dataset = datasets.load_dataset("animelover/danbooru2022", "1-full", trust_remote_code=True, split="train", streaming=True, features=anime_dataset_ft)\
|
|
# .select_columns(["image"])\
|
|
# .take(ROWS_PER_DS)\
|
|
# .add_column("question", ["Describe this image in one sentence. Include the word anime in the sentence."] * ROWS_PER_DS)\
|
|
# .add_column("keywords", [["anime"]] * ROWS_PER_DS)
|
|
|
|
# coco_dataset = datasets.load_dataset("detection-datasets/coco", split="train", streaming=True)\
|
|
# .select_columns(["image"])\
|
|
# .take(ROWS_PER_DS)\
|
|
# .add_column("question", ["Describe this image in one sentence."] * ROWS_PER_DS)\
|
|
# .add_column("keywords", [[""]] * ROWS_PER_DS)
|
|
|
|
# movie_poster_dataset = datasets.load_dataset("skvarre/movie_posters-100k", split="train", streaming=True)\
|
|
# .select_columns(["image"])\
|
|
# .take(ROWS_PER_DS)\
|
|
# .add_column("question", ["Describe this image in one sentence."] * ROWS_PER_DS)\
|
|
# .add_column("keywords", [[""]] * ROWS_PER_DS)
|
|
|
|
# cars_dataset = datasets.load_dataset("tanganke/stanford_cars", split="train", streaming=True)\
|
|
# .select_columns(["image"])\
|
|
# .take(ROWS_PER_DS)\
|
|
# .add_column("question", ["Describe this image in one sentence."] * ROWS_PER_DS)\
|
|
# .add_column("keywords", [[""]] * ROWS_PER_DS)
|
|
|
|
# website_dataset = datasets.load_dataset("silatus/1k_Website_Screenshots_and_Metadata", split="train", streaming=True)\
|
|
# .select_columns(["image"])\
|
|
# .take(ROWS_PER_DS)\
|
|
# .add_column("question", ["Describe this image in one sentence."] * ROWS_PER_DS)\
|
|
# .add_column("keywords", [[""]] * ROWS_PER_DS)
|
|
|
|
# movie_scene_dataset = datasets.load_dataset("unography/movie-scenes-resized-captioned", split="train", streaming=True)\
|
|
# .select_columns(["image"])\
|
|
# .take(ROWS_PER_DS)\
|
|
# .add_column("question", ["Describe this image in one sentence."] * ROWS_PER_DS)\
|
|
# .add_column("keywords", [[""]] * ROWS_PER_DS)
|
|
|
|
# ds = datasets.concatenate_datasets([
|
|
# flickr_dataset,
|
|
# wiki_art_dataset,
|
|
# anime_dataset,
|
|
# coco_dataset,
|
|
# movie_poster_dataset,
|
|
# cars_dataset,
|
|
# website_dataset,
|
|
# movie_scene_dataset,
|
|
# ]).cast_column("image", datasets.Image(decode=True)).skip(SKIP_PARQUET_BATCH * PARQUET_BATCH_SIZE)
|
|
|
|
ds = wiki_art_dataset.cast_column("image", datasets.Image(decode=True))
|
|
|
|
data_loader = torch.utils.data.DataLoader(
|
|
ds,
|
|
batch_size=BATCH_SIZE,
|
|
collate_fn=collate
|
|
)
|
|
|
|
temp_ds = {
|
|
"image": [],
|
|
"keywords": [],
|
|
"caption": [],
|
|
"generated_image": []
|
|
}
|
|
temp_ds_size = 0
|
|
|
|
ds_features = datasets.Features({
|
|
"image": datasets.Image(),
|
|
"keywords": datasets.Sequence(datasets.Value(dtype="string")),
|
|
"caption": datasets.Value(dtype="string"),
|
|
"generated_image": datasets.Image(),
|
|
})
|
|
|
|
generator = torch.Generator(device="cpu").manual_seed(12321313)
|
|
|
|
batch_count = SKIP_PARQUET_BATCH
|
|
|
|
for batch_index, batch in enumerate(data_loader):
|
|
images, keywords = batch
|
|
|
|
prompts = []
|
|
for i, img in enumerate(images):
|
|
caption = moondream.caption(img, length="normal")["caption"]
|
|
|
|
add_keywords = len(keywords[i]) > 0 and keywords[i][0] != ""
|
|
for k in keywords[i]:
|
|
if k and k in caption:
|
|
add_keywords = False
|
|
break
|
|
|
|
prompt = caption
|
|
if add_keywords:
|
|
prompt = f"{', '.join(keywords[i])}, {caption}"
|
|
|
|
prompts.append(prompt)
|
|
|
|
gen_imgs = pipe(
|
|
prompts,
|
|
num_inference_steps=28,
|
|
guidance_scale=3.5,
|
|
generator=generator,
|
|
max_sequence_length=512,
|
|
).images
|
|
|
|
temp_ds["image"].extend(images)
|
|
temp_ds["caption"].extend(prompts)
|
|
temp_ds["keywords"].extend(keywords)
|
|
temp_ds["generated_image"].extend(gen_imgs)
|
|
|
|
temp_ds_size += BATCH_SIZE
|
|
|
|
if temp_ds_size == PARQUET_BATCH_SIZE:
|
|
batch_ds = datasets.Dataset.from_dict(temp_ds, features=ds_features)
|
|
batch_ds.to_parquet(
|
|
f"data/batch_{batch_count}.parquet",
|
|
)
|
|
temp_ds_size = 0
|
|
temp_ds["image"].clear()
|
|
temp_ds["caption"].clear()
|
|
temp_ds["keywords"].clear()
|
|
temp_ds["generated_image"].clear()
|
|
|
|
batch_count += 1
|