Upload joycaption.py

joycaption.py

@@ -1,4 +1,13 @@
-import spaces
+import os
+if os.environ.get("SPACES_ZERO_GPU") is not None:
+    import spaces
+else:
+    class spaces:
+        @staticmethod
+        def GPU(func):
+            def wrapper(*args, **kwargs):
+                return func(*args, **kwargs)
+            return wrapper
 import gradio as gr
 from huggingface_hub import InferenceClient
 from torch import nn
@@ -7,11 +16,13 @@ from pathlib import Path
 import torch
 import torch.amp.autocast_mode
 from PIL import Image
-import os
 import torchvision.transforms.functional as TVF
 import gc
 from peft import PeftConfig
 
+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"
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 use_inference_client = False
@@ -119,6 +130,8 @@ class ImageAdapter(nn.Module):
 # https://huggingface.co/blog/4bit-transformers-bitsandbytes
 # https://huggingface.co/docs/transformers/main/en/peft
 # https://huggingface.co/docs/transformers/main/en/peft#enable-and-disable-adapters
+# https://huggingface.co/docs/transformers/main/quantization/bitsandbytes?bnb=4-bit
+# https://huggingface.co/lllyasviel/flux1-dev-bnb-nf4
 tokenizer = None
 text_model_client = None
 text_model = None
@@ -171,14 +184,12 @@ load_text_model.zerogpu = True
 print("Loading CLIP")
 clip_processor = AutoProcessor.from_pretrained(CLIP_PATH)
 clip_model = AutoModel.from_pretrained(CLIP_PATH).vision_model
-
 if (CHECKPOINT_PATH / "clip_model.pt").exists():
     print("Loading VLM's custom vision model")
     checkpoint = torch.load(CHECKPOINT_PATH / "clip_model.pt", map_location='cpu')
     checkpoint = {k.replace("_orig_mod.module.", ""): v for k, v in checkpoint.items()}
     clip_model.load_state_dict(checkpoint)
     del checkpoint
-
 clip_model.eval().requires_grad_(False).to(device)
 
 # Tokenizer
@@ -376,6 +387,19 @@ def is_repo_exists(repo_id):
         return True # for safe
 
 
+def is_valid_repo(repo_id):
+    from huggingface_hub import HfApi
+    import re
+    try:
+        if not re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', repo_id): return False
+        api = HfApi()
+        if api.repo_exists(repo_id=repo_id): return True
+        else: return False
+    except Exception as e:
+        print(f"Failed to connect {repo_id}. {e}")
+        return False
+
+
 def get_text_model():
     return list(llm_models.keys())