caption-flow 0.1.0__py3-none-any.whl

caption_flow/utils/auth.py

@@ -0,0 +1,67 @@
+"""Authentication management."""
+
+from typing import Dict, Any, Optional
+from dataclasses import dataclass
+
+
+@dataclass
+class WorkerAuthenticationDetails:
+    """Details for worker authentication."""
+
+    name: str
+    token: str
+    role: str
+
+
+class AuthManager:
+    """Manages authentication tokens."""
+
+    def __init__(self, config: Dict[str, Any]):
+        self.reload_config(config=config)
+
+    def authenticate(self, token: str) -> Optional[str]:
+        """Authenticate token and return role."""
+        role = None
+        for worker_token in self.worker_tokens:
+            if token == worker_token:
+                role = "worker"
+                break
+        if role is None:
+            for admin_token in self.admin_tokens:
+                if token == admin_token:
+                    role = "admin"
+                    break
+        if role is None:
+            for monitor_token in self.monitor_tokens:
+                if token == monitor_token:
+                    role = "monitor"
+                    break
+
+        worker_auth_details = WorkerAuthenticationDetails(
+            role=role, name=self.worker_tokens.get(token, f"Anonymous {role}"), token=token
+        )
+        return worker_auth_details
+
+    def reload_config(self, config: dict) -> None:
+        """Reload configuration from file."""
+        self.worker_tokens = {}
+        self.admin_tokens = {}
+        self.monitor_tokens = {}
+
+        # Load worker tokens
+        for worker in config.get("worker_tokens", []):
+            worker_name = worker.get("name", None)
+            assert worker_name is not None, "Worker token must have a name"
+            self.worker_tokens[worker["token"]] = worker_name
+
+        # Load admin tokens
+        for admin in config.get("admin_tokens", []):
+            admin_name = admin.get("name", None)
+            assert admin_name is not None, "Admin token must have a name"
+            self.admin_tokens[admin["token"]] = admin_name
+
+        # Load monitor tokens
+        for monitor in config.get("monitor_tokens", []):
+            monitor_name = monitor.get("name", None)
+            assert monitor_name is not None, "Monitor token must have a name"
+            self.monitor_tokens[monitor["token"]] = monitor_name

caption_flow/utils/caption_utils.py

@@ -0,0 +1,172 @@
+"""Caption processing utilities from the original vLLM script."""
+
+from typing import List, Dict
+
+
+class CaptionUtils:
+    """Utilities for cleaning and combining captions."""
+
+    @staticmethod
+    def clean_caption(c: str) -> str:
+        """Clean a single caption by removing generic phrases and formatting."""
+        if not c:
+            return ""
+
+        generic = [
+            "in this image we can see ",
+            "this image shows ",
+            "the image depicts ",
+            "the image features ",
+            "this is an image of ",
+            "the image contains ",
+            "the picture shows ",
+            "we can see ",
+            "there is ",
+            "there are ",
+        ]
+
+        low = c.lower()
+        for p in generic:
+            if low.startswith(p):
+                c = c[len(p) :]
+                if c:
+                    c = c[0].upper() + c[1:]
+                break
+
+        # Remove leading articles if the rest isn't capitalized
+        if c.lower().startswith(("a ", "an ")):
+            parts = c.split(maxsplit=1)
+            if len(parts) > 1 and not parts[1][0].isupper():
+                c = parts[1]
+                c = c[0].upper() + c[1:]
+
+        # Clean whitespace
+        c = " ".join(c.split())
+
+        # Add period if missing
+        if c and c[-1] not in ".!?":
+            c += "."
+
+        return c
+
+    @classmethod
+    def combine(cls, descs: List[str]) -> str:
+        """Combine multiple descriptions into a rich, multi-line caption."""
+        if not descs:
+            return ""
+
+        filtered = []
+        heads = [
+            "in this image we can see",
+            "this image shows",
+            "the image depicts",
+            "a cartoon",
+            "a drawing",
+            "an illustration",
+        ]
+
+        # Filter out short generic descriptions
+        for d in descs:
+            if not d:
+                continue
+            dl = d.lower().strip()
+            if any(dl.startswith(h) and len(dl.split()) < 8 for h in heads):
+                continue
+            if len(d) > 10:
+                filtered.append(d)
+
+        if not filtered:
+            filtered = [max(descs, key=len, default="")]
+
+        # Use the longest as the main description
+        main = cls.clean_caption(max(filtered, key=len))
+        parts = [main]
+        seen = set(main.lower().split())
+
+        # Categorize additional descriptions
+        buckets = {
+            "characters": [
+                "character",
+                "person",
+                "animal",
+                "anthro",
+                "wearing",
+                "dressed",
+            ],
+            "actions": ["doing", "action", "playing", "running", "sitting", "standing"],
+            "settings": [
+                "room",
+                "outdoor",
+                "indoor",
+                "setting",
+                "background",
+                "environment",
+            ],
+            "styles": ["style", "art", "drawn", "sketch", "painted", "digital"],
+            "moods": [
+                "mood",
+                "emotion",
+                "feeling",
+                "atmosphere",
+                "happy",
+                "sad",
+                "angry",
+            ],
+        }
+
+        def categorize(text: str) -> str:
+            """Categorize a description based on keywords."""
+            text_lower = text.lower()
+            for category, keywords in buckets.items():
+                if any(keyword in text_lower for keyword in keywords):
+                    return category
+            return "details"
+
+        # Group descriptions by category
+        by_bucket: Dict[str, List[str]] = {}
+        for desc in filtered:
+            category = categorize(desc)
+            by_bucket.setdefault(category, []).append(desc)
+
+        # Add descriptions from each category
+        for category in ["characters", "actions", "settings", "moods", "styles", "details"]:
+            if category in by_bucket and by_bucket[category]:
+                desc = by_bucket[category][0]
+                words = desc.lower().split()
+
+                # Check if this adds enough new information
+                new_words = [w for w in words if w not in seen and len(w) > 3]
+                if len(new_words) > 3:
+                    clean = cls.clean_caption(desc)
+                    if clean and clean not in parts:
+                        parts.append(clean)
+                        seen.update(words)
+
+        # Return each part as a separate line for rich captions
+        return "\n".join(parts)
+
+    @staticmethod
+    def validate_caption(caption: str, min_length: int = 20) -> bool:
+        """Validate if a caption meets quality standards."""
+        if not caption or len(caption) < min_length:
+            return False
+
+        # Check for refusal patterns
+        refusal_patterns = [
+            "i'm sorry",
+            "i cannot",
+            "i apologize",
+            "inappropriate",
+            "unable to",
+            "refuse to",
+        ]
+
+        caption_lower = caption.lower()
+        if any(pattern in caption_lower for pattern in refusal_patterns):
+            return False
+
+        # Check for too generic
+        if caption_lower in ["image", "picture", "photo", "illustration"]:
+            return False
+
+        return True

caption_flow/utils/certificates.py

@@ -0,0 +1,140 @@
+"""SSL certificate management."""
+
+import subprocess
+from pathlib import Path
+from typing import Optional
+from cryptography import x509
+from cryptography.x509.oid import NameOID
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from datetime import datetime, timedelta
+
+
+class CertificateManager:
+    """Manages SSL certificate generation."""
+
+    def generate_self_signed(
+        self, output_dir: Path, domain: str = "localhost"
+    ) -> tuple[Path, Path]:
+        """Generate self-signed certificate for development."""
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        # Generate private key
+        key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+
+        # Generate certificate
+        subject = issuer = x509.Name(
+            [
+                x509.NameAttribute(NameOID.COMMON_NAME, domain),
+            ]
+        )
+
+        cert = (
+            x509.CertificateBuilder()
+            .subject_name(subject)
+            .issuer_name(issuer)
+            .public_key(key.public_key())
+            .serial_number(x509.random_serial_number())
+            .not_valid_before(datetime.utcnow())
+            .not_valid_after(datetime.utcnow() + timedelta(days=365))
+            .add_extension(
+                x509.SubjectAlternativeName(
+                    [
+                        x509.DNSName(domain),
+                        x509.DNSName("localhost"),
+                        x509.DNSName("127.0.0.1"),
+                    ]
+                ),
+                critical=False,
+            )
+            .sign(key, hashes.SHA256())
+        )
+
+        # Write files
+        cert_path = output_dir / "cert.pem"
+        key_path = output_dir / "key.pem"
+
+        with open(cert_path, "wb") as f:
+            f.write(cert.public_bytes(serialization.Encoding.PEM))
+
+        with open(key_path, "wb") as f:
+            f.write(
+                key.private_bytes(
+                    encoding=serialization.Encoding.PEM,
+                    format=serialization.PrivateFormat.TraditionalOpenSSL,
+                    encryption_algorithm=serialization.NoEncryption(),
+                )
+            )
+
+        return cert_path, key_path
+
+    def generate_letsencrypt(
+        self, domain: str, email: str, output_dir: Optional[Path] = None, staging: bool = False
+    ) -> tuple[Path, Path]:
+        """
+        Generate Let's Encrypt certificate.
+
+        Args:
+            domain: Domain name for certificate
+            email: Email for Let's Encrypt account
+            output_dir: Custom output directory (uses /etc/letsencrypt by default)
+            staging: Use Let's Encrypt staging server for testing
+        """
+        cmd = [
+            "certbot",
+            "certonly",
+            "--standalone",
+            "--non-interactive",
+            "--agree-tos",
+            "--email",
+            email,
+            "-d",
+            domain,
+        ]
+
+        if staging:
+            cmd.append("--staging")
+
+        if output_dir:
+            # Use custom config and work directories
+            output_dir = Path(output_dir)
+            output_dir.mkdir(parents=True, exist_ok=True)
+            cmd.extend(
+                [
+                    "--config-dir",
+                    str(output_dir),
+                    "--work-dir",
+                    str(output_dir / "work"),
+                    "--logs-dir",
+                    str(output_dir / "logs"),
+                ]
+            )
+            cert_base = output_dir / "live" / domain
+        else:
+            cert_base = Path(f"/etc/letsencrypt/live/{domain}")
+
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        if result.returncode != 0:
+            raise RuntimeError(f"Certbot failed: {result.stderr}")
+
+        cert_path = cert_base / "fullchain.pem"
+        key_path = cert_base / "privkey.pem"
+
+        if not cert_path.exists() or not key_path.exists():
+            raise RuntimeError(f"Certificate files not found at {cert_base}")
+
+        return cert_path, key_path
+
+    def get_cert_info(self, cert_path: Path) -> dict:
+        """Get information about an existing certificate."""
+        with open(cert_path, "rb") as f:
+            cert = x509.load_pem_x509_certificate(f.read())
+
+        return {
+            "subject": cert.subject.rfc4514_string(),
+            "issuer": cert.issuer.rfc4514_string(),
+            "not_before": cert.not_valid_before,
+            "not_after": cert.not_valid_after,
+            "serial_number": cert.serial_number,
+            "is_self_signed": cert.issuer == cert.subject,
+        }