Upload 2 files

env.py

@@ -40,7 +40,8 @@ load_diffusers_format_model = [
     'rubbrband/realcartoonRealistic_v14',
     'KBlueLeaf/Kohaku-XL-Epsilon-rev2',
     'KBlueLeaf/Kohaku-XL-Epsilon-rev3',
-    'kayfahaarukku/UrangDiffusion-1.1',
+    'KBlueLeaf/Kohaku-XL-Zeta',
+    'kayfahaarukku/UrangDiffusion-1.2',
     'Eugeoter/artiwaifu-diffusion-1.0',
     'Raelina/Rae-Diffusion-XL-V2',
     'Raelina/Raemu-XL-V4',

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