Upload 8 files

tagger/fl2sd3longcap.py

@@ -8,9 +8,13 @@ import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-fl_model = AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True).to(device).eval()
-fl_processor = AutoProcessor.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True)
 
+try:
+    fl_model = AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True).to("cpu").eval()
+    fl_processor = AutoProcessor.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True)
+except Exception as e:
+    print(e)
+    fl_model = fl_processor = None
 
 def fl_modify_caption(caption: str) -> str:
     """
@@ -41,7 +45,7 @@ def fl_modify_caption(caption: str) -> str:
     return modified_caption if modified_caption != caption else caption
 
 
-@spaces.GPU
+@spaces.GPU(duration=30)
 def fl_run_example(image):
     task_prompt = "<DESCRIPTION>"
     prompt = task_prompt + "Describe this image in great detail."
@@ -50,6 +54,7 @@ def fl_run_example(image):
     if image.mode != "RGB":
         image = image.convert("RGB")
 
+    fl_model.to(device)
     inputs = fl_processor(text=prompt, images=image, return_tensors="pt").to(device)
     generated_ids = fl_model.generate(
         input_ids=inputs["input_ids"],
@@ -57,6 +62,7 @@ def fl_run_example(image):
         max_new_tokens=1024,
         num_beams=3
     )
+    fl_model.to("cpu")
     generated_text = fl_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
     parsed_answer = fl_processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))
     return fl_modify_caption(parsed_answer["<DESCRIPTION>"])

tagger/tagger.py

@@ -12,10 +12,15 @@ from pathlib import Path
 WD_MODEL_NAMES = ["p1atdev/wd-swinv2-tagger-v3-hf"]
 WD_MODEL_NAME = WD_MODEL_NAMES[0]
 
-wd_model = AutoModelForImageClassification.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
-wd_model.to("cuda" if torch.cuda.is_available() else "cpu")
-wd_processor = AutoImageProcessor.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+default_device = device
 
+try:
+    wd_model = AutoModelForImageClassification.from_pretrained(WD_MODEL_NAME, trust_remote_code=True).to(default_device).eval()
+    wd_processor = AutoImageProcessor.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
+except Exception as e:
+    print(e)
+    wd_model = wd_processor = None
 
 def _people_tag(noun: str, minimum: int = 1, maximum: int = 5):
     return (
@@ -506,7 +511,7 @@ def gen_prompt(rating: list[str], character: list[str], general: list[str]):
     return ", ".join(all_tags)
 
 
-@spaces.GPU()
+@spaces.GPU(duration=30)
 def predict_tags(image: Image.Image, general_threshold: float = 0.3, character_threshold: float = 0.8):
     inputs = wd_processor.preprocess(image, return_tensors="pt")
 
@@ -514,9 +519,11 @@ def predict_tags(image: Image.Image, general_threshold: float = 0.3, character_t
     logits = torch.sigmoid(outputs.logits[0])  # take the first logits
 
     # get probabilities
+    if device != default_device: wd_model.to(device=device)
     results = {
         wd_model.config.id2label[i]: float(logit.float()) for i, logit in enumerate(logits)
     }
+    if device != default_device: wd_model.to(device=default_device)
     # rating, character, general
     rating, character, general = postprocess_results(
         results, general_threshold, character_threshold