fal 1.2.2__tar.gz → 1.2.4__tar.gz

{fal-1.2.2 → fal-1.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fal
-Version: 1.2.2
+Version: 1.2.4
 Summary: fal is an easy-to-use Serverless Python Framework
 Author: Features & Labels <support@fal.ai>
 Requires-Python: >=3.8

{fal-1.2.2 → fal-1.2.4}/fal.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fal
-Version: 1.2.2
+Version: 1.2.4
 Summary: fal is an easy-to-use Serverless Python Framework
 Author: Features & Labels <support@fal.ai>
 Requires-Python: >=3.8

{fal-1.2.2 → fal-1.2.4}/fal.egg-info/SOURCES.txt

@@ -126,6 +126,12 @@ src/fal/toolkit/file/providers/gcp.py
 src/fal/toolkit/file/providers/r2.py
 src/fal/toolkit/image/__init__.py
 src/fal/toolkit/image/image.py
+src/fal/toolkit/image/safety_checker.py
+src/fal/toolkit/image/nsfw_filter/__init__.py
+src/fal/toolkit/image/nsfw_filter/env.py
+src/fal/toolkit/image/nsfw_filter/inference.py
+src/fal/toolkit/image/nsfw_filter/model.py
+src/fal/toolkit/image/nsfw_filter/requirements.txt
 src/fal/toolkit/utils/__init__.py
 src/fal/toolkit/utils/download_utils.py
 tests/__init__.py

{fal-1.2.2 → fal-1.2.4}/src/fal/_fal_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.2.2'
-__version_tuple__ = version_tuple = (1, 2, 2)
+__version__ = version = '1.2.4'
+__version_tuple__ = version_tuple = (1, 2, 4)

{fal-1.2.2 → fal-1.2.4}/src/fal/app.py

@@ -3,7 +3,9 @@ from __future__ import annotations
 import inspect
 import json
 import os
+import queue
 import re
+import threading
 import time
 import typing
 from contextlib import asynccontextmanager, contextmanager
@@ -72,17 +74,22 @@ def wrap_app(cls: type[App], **kwargs) -> fal.api.IsolatedFunction:
 
 
 class EndpointClient:
-    def __init__(self, url, endpoint, signature):
+    def __init__(self, url, endpoint, signature, timeout: int | None = None):
         self.url = url
         self.endpoint = endpoint
         self.signature = signature
+        self.timeout = timeout
 
         annotations = endpoint.__annotations__ or {}
         self.return_type = annotations.get("return") or None
 
     def __call__(self, data):
         with httpx.Client() as client:
-            resp = client.post(self.url + self.signature.path, json=dict(data))
+            resp = client.post(
+                self.url + self.signature.path,
+                json=data.dict() if hasattr(data, "dict") else dict(data),
+                timeout=self.timeout,
+            )
             resp.raise_for_status()
             resp_dict = resp.json()
 
@@ -93,7 +100,12 @@ class EndpointClient:
 
 
 class AppClient:
-    def __init__(self, cls, url):
+    def __init__(
+        self,
+        cls,
+        url,
+        timeout: int | None = None,
+    ):
         self.url = url
         self.cls = cls
 
@@ -101,19 +113,38 @@ class AppClient:
             signature = getattr(endpoint, "route_signature", None)
             if signature is None:
                 continue
-
-            setattr(self, name, EndpointClient(self.url, endpoint, signature))
+            endpoint_client = EndpointClient(
+                self.url,
+                endpoint,
+                signature,
+                timeout=timeout,
+            )
+            setattr(self, name, endpoint_client)
 
     @classmethod
     @contextmanager
     def connect(cls, app_cls):
         app = wrap_app(app_cls)
         info = app.spawn()
+        _shutdown_event = threading.Event()
+
+        def _print_logs():
+            while not _shutdown_event.is_set():
+                try:
+                    log = info.logs.get(timeout=0.1)
+                except queue.Empty:
+                    continue
+                print(log)
+
+        _log_printer = threading.Thread(target=_print_logs, daemon=True)
+        _log_printer.start()
+
         try:
             with httpx.Client() as client:
                 retries = 100
                 while retries:
                     resp = client.get(info.url + "/health")
+
                     if resp.is_success:
                         break
                     elif resp.status_code != 500:
@@ -121,9 +152,12 @@ class AppClient:
                     time.sleep(0.1)
                     retries -= 1
 
-            yield cls(app_cls, info.url)
+            client = cls(app_cls, info.url)
+            yield client
         finally:
             info.stream.cancel()
+            _shutdown_event.set()
+            _log_printer.join()
 
     def health(self):
         with httpx.Client() as client:
@@ -158,7 +192,7 @@ class App(fal.api.BaseServable):
         app_name = kwargs.pop("name", None) or _to_fal_app_name(cls.__name__)
         parent_settings = getattr(cls, "host_kwargs", {})
         cls.host_kwargs = {**parent_settings, **kwargs}
-        cls.app_name = app_name
+        cls.app_name = getattr(cls, "app_name", app_name)
 
         if cls.__init__ is not App.__init__:
             raise ValueError(

fal-1.2.4/src/fal/toolkit/image/__init__.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import TYPE_CHECKING
+from urllib.request import Request, urlopen
+
+from .image import *  # noqa: F403
+
+if TYPE_CHECKING:
+    from PIL.Image import Image as PILImage
+
+
+def filter_by(
+    has_nsfw_concepts: list[bool],
+    images: list[PILImage],
+) -> list[PILImage]:
+    from PIL import Image as PILImage
+
+    return [
+        (
+            PILImage.new("RGB", (image.width, image.height), (0, 0, 0))
+            if has_nsfw
+            else image
+        )
+        for image, has_nsfw in zip(images, has_nsfw_concepts)
+    ]
+
+
+def preprocess_image(image_pil, convert_to_rgb=True, fix_orientation=True):
+    from PIL import ImageOps, ImageSequence
+
+    # For MPO (multi picture object) format images, we only need the first image
+    images = []
+    for image in ImageSequence.Iterator(image_pil):
+        img = image
+
+        if convert_to_rgb:
+            img = img.convert("RGB")
+
+        if fix_orientation:
+            img = ImageOps.exif_transpose(img)
+
+        images.append(img)
+
+        break
+
+    return images[0]
+
+
+@lru_cache(maxsize=64)
+def read_image_from_url(
+    url: str, convert_to_rgb: bool = True, fix_orientation: bool = True
+):
+    from fastapi import HTTPException
+    from PIL import Image
+
+    TEMP_HEADERS = {
+        "User-Agent": (
+            "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.8; rv:21.0) "
+            "Gecko/20100101 Firefox/21.0"
+        ),
+    }
+
+    try:
+        request = Request(url, headers=TEMP_HEADERS)
+        response = urlopen(request)
+        image_pil = Image.open(response)
+    except Exception:
+        import traceback
+
+        traceback.print_exc()
+        raise HTTPException(422, f"Could not load image from url: {url}")
+
+    return preprocess_image(image_pil, convert_to_rgb, fix_orientation)

fal-1.2.4/src/fal/toolkit/image/nsfw_filter/__init__.py

@@ -0,0 +1,11 @@
+from .inference import (
+    NSFWImageDetectionInput,
+    NSFWImageDetectionOutput,
+    run_nsfw_estimation,
+)
+
+__all__ = [
+    "NSFWImageDetectionInput",
+    "NSFWImageDetectionOutput",
+    "run_nsfw_estimation",
+]

fal-1.2.4/src/fal/toolkit/image/nsfw_filter/env.py

@@ -0,0 +1,9 @@
+from pathlib import Path
+
+CURR_DIR = Path(__file__).resolve().parent
+
+
+def get_requirements():
+    with open(CURR_DIR / "requirements.txt") as fp:
+        requirements = fp.read().split()
+    return requirements

fal-1.2.4/src/fal/toolkit/image/nsfw_filter/inference.py

@@ -0,0 +1,77 @@
+from pydantic import BaseModel, Field
+
+import fal
+from fal.toolkit.image import read_image_from_url
+
+from .env import get_requirements
+from .model import get_model
+
+
+class NSFWImageDetectionInput(BaseModel):
+    image_url: str = Field(
+        description="Input image url.",
+        examples=[
+            "https://storage.googleapis.com/falserverless/model_tests/remove_background/elephant.jpg",
+        ],
+    )
+
+
+class NSFWImageDetectionOutput(BaseModel):
+    nsfw_probability: float = Field(
+        description="The probability of the image being NSFW.",
+    )
+
+
+def check_nsfw_content(pil_image: object):
+    import torch
+
+    model, processor = get_model()
+
+    with torch.no_grad():
+        inputs = processor(images=pil_image, return_tensors="pt")
+        outputs = model(**inputs)
+        logits = outputs.logits.squeeze()  # Remove batch dimension to simplify indexing
+
+    # Apply softmax to convert logits to probabilities
+    probabilities = torch.softmax(logits, dim=0)
+
+    nsfw_class_index = model.config.label2id.get(
+        "nsfw", None
+    )  # Replace "NSFW" with the exact class name if different
+
+    # Validate that NSFW class index is found
+    if nsfw_class_index is not None:
+        nsfw_probability = probabilities[int(nsfw_class_index)].item()
+        return nsfw_probability
+    else:
+        raise ValueError("NSFW class not found in model output.")
+
+
+def run_nsfw_estimation(
+    input: NSFWImageDetectionInput,
+) -> NSFWImageDetectionOutput:
+    img = read_image_from_url(input.image_url)
+    nsfw_probability = check_nsfw_content(img)
+
+    return NSFWImageDetectionOutput(nsfw_probability=nsfw_probability)
+
+
+@fal.function(
+    requirements=get_requirements(),
+    machine_type="GPU-A6000",
+    serve=True,
+)
+def run_nsfw_estimation_on_fal(
+    input: NSFWImageDetectionInput,
+) -> NSFWImageDetectionOutput:
+    return run_nsfw_estimation(input)
+
+
+if __name__ == "__main__":
+    local = run_nsfw_estimation_on_fal.on(serve=False)
+    result = local(
+        NSFWImageDetectionInput(
+            image_url="https://storage.googleapis.com/falserverless/model_tests/remove_background/elephant.jpg",
+        )
+    )
+    print(result)

fal-1.2.4/src/fal/toolkit/image/nsfw_filter/model.py

@@ -0,0 +1,18 @@
+import fal
+
+
+@fal.cached
+def get_model():
+    import os
+
+    from transformers import AutoModelForImageClassification, ViTImageProcessor
+
+    os.environ["TRANSFORMERS_CACHE"] = "/data/models"
+    os.environ["HF_HOME"] = "/data/models"
+
+    model = AutoModelForImageClassification.from_pretrained(
+        "Falconsai/nsfw_image_detection"
+    )
+    processor = ViTImageProcessor.from_pretrained("Falconsai/nsfw_image_detection")
+
+    return model, processor

fal-1.2.4/src/fal/toolkit/image/nsfw_filter/requirements.txt

@@ -0,0 +1,4 @@
+accelerate
+Pillow
+torch
+transformers

fal-1.2.4/src/fal/toolkit/image/safety_checker.py

@@ -0,0 +1,107 @@
+from typing import Any
+
+import fal
+
+from . import filter_by
+from .nsfw_filter.model import get_model
+
+
+@fal.cached
+def load_safety_checker():
+    import torch
+    from diffusers.pipelines.stable_diffusion.safety_checker import (
+        StableDiffusionSafetyChecker,
+    )
+    from transformers import AutoFeatureExtractor
+
+    feature_extractor = AutoFeatureExtractor.from_pretrained(
+        "CompVis/stable-diffusion-safety-checker",
+        torch_dtype="float16",
+    )
+    safety_checker = StableDiffusionSafetyChecker.from_pretrained(
+        "CompVis/stable-diffusion-safety-checker",
+        torch_dtype=torch.float16,
+    ).to("cuda")
+
+    return feature_extractor, safety_checker
+
+
+def run_safety_checker(
+    pil_images: list[object],
+) -> list[bool]:
+    import numpy as np
+    import torch
+
+    feature_extractor, safety_checker = load_safety_checker()
+
+    safety_checker_input = feature_extractor(pil_images, return_tensors="pt").to("cuda")
+
+    np_image = [np.array(val) for val in pil_images]
+
+    _, has_nsfw_concept = safety_checker(
+        images=np_image,
+        clip_input=safety_checker_input.pixel_values.to(torch.float16),
+    )
+
+    return has_nsfw_concept
+
+
+def run_safety_checker_v2(pil_images: list, nsfw_threshold: float = 0.5) -> list[bool]:
+    import torch
+
+    model, processor = get_model()
+
+    has_nsfw_concept = []
+
+    with torch.no_grad():
+        for pil_image in pil_images:
+            inputs = processor(
+                images=pil_image.convert("RGB"),
+                return_tensors="pt",
+            )
+            outputs = model(**inputs)
+            logits = (
+                outputs.logits.squeeze()
+            )  # Remove batch dimension to simplify indexing
+
+            # Apply softmax to convert logits to probabilities
+            probabilities = torch.softmax(logits, dim=0)
+
+            nsfw_class_index = model.config.label2id.get(
+                "nsfw", None
+            )  # Replace "NSFW" with the exact class name if different
+
+            # Validate that NSFW class index is found
+            if nsfw_class_index is not None:
+                nsfw_probability = probabilities[int(nsfw_class_index)].item()
+                print("NSFW probability:", nsfw_probability)
+                has_nsfw_concept.append(nsfw_probability > nsfw_threshold)
+            else:
+                raise ValueError("NSFW class not found in model output.")
+
+    return has_nsfw_concept
+
+
+def postprocess_images(
+    pil_images: list[object],
+    enable_safety_checker: bool = True,
+    safety_checker_version: int = 2,
+) -> dict[str, Any]:
+    outputs: dict[str, list[Any]] = {
+        "images": pil_images,
+    }
+
+    if enable_safety_checker:
+        safety_checker_fn = (
+            run_safety_checker_v2 if safety_checker_version == 2 else run_safety_checker
+        )
+        outputs["has_nsfw_concepts"] = safety_checker_fn(pil_images)  # type: ignore
+    else:
+        outputs["has_nsfw_concepts"] = [False] * len(pil_images)
+
+    outputs["images"] = filter_by(
+        outputs["has_nsfw_concepts"],
+        outputs["images"],
+    )
+
+    return outputs

fal-1.2.2/src/fal/toolkit/image/__init__.py

@@ -1,3 +0,0 @@
-from __future__ import annotations
-
-from .image import *  # noqa: F403