fal 1.2.1__py3-none-any.whl → 1.7.2__py3-none-any.whl

fal/toolkit/file/providers/r2.py

@@ -6,8 +6,10 @@ import posixpath
 import uuid
 from dataclasses import dataclass
 from io import BytesIO
+from typing import Optional
 
 from fal.toolkit.file.types import FileData, FileRepository
+from fal.toolkit.utils.retry import retry
 
 DEFAULT_URL_TIMEOUT = 60 * 15  # 15 minutes
 
@@ -67,7 +69,12 @@ class R2Repository(FileRepository):
 
         return self._bucket
 
-    def save(self, data: FileData) -> str:
+    @retry(max_retries=3, base_delay=1, backoff_type="exponential", jitter=True)
+    def save(
+        self,
+        data: FileData,
+        object_lifecycle_preference: Optional[dict[str, str]] = None,
+    ) -> str:
         destination_path = posixpath.join(
             self.key,
             f"{uuid.uuid4().hex}_{data.file_name}",

fal/toolkit/file/providers/s3.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import os
+import posixpath
+import uuid
+from dataclasses import dataclass
+from io import BytesIO
+from typing import Optional
+
+from fal.toolkit.file.types import FileData, FileRepository
+from fal.toolkit.utils.retry import retry
+
+DEFAULT_URL_TIMEOUT = 60 * 15  # 15 minutes
+
+
+@dataclass
+class S3Repository(FileRepository):
+    bucket_name: str = "fal_file_storage"
+    url_expiration: int = DEFAULT_URL_TIMEOUT
+    aws_access_key_id: str | None = None
+    aws_secret_access_key: str | None = None
+
+    _s3_client = None
+
+    def __post_init__(self):
+        try:
+            import boto3
+            from botocore.client import Config
+        except ImportError:
+            raise Exception("boto3 is not installed")
+
+        if self.aws_access_key_id is None:
+            self.aws_access_key_id = os.environ.get("AWS_ACCESS_KEY_ID")
+            if self.aws_access_key_id is None:
+                raise Exception("AWS_ACCESS_KEY_ID environment variable is not set")
+
+        if self.aws_secret_access_key is None:
+            self.aws_secret_access_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
+            if self.aws_secret_access_key is None:
+                raise Exception("AWS_SECRET_ACCESS_KEY environment variable is not set")
+
+        self._s3_client = boto3.client(
+            "s3",
+            aws_access_key_id=self.aws_access_key_id,
+            aws_secret_access_key=self.aws_secret_access_key,
+            config=Config(signature_version="s3v4"),
+        )
+
+    @property
+    def storage_client(self):
+        if self._s3_client is None:
+            raise Exception("S3 client is not initialized")
+
+        return self._s3_client
+
+    @retry(max_retries=3, base_delay=1, backoff_type="exponential", jitter=True)
+    def save(
+        self,
+        data: FileData,
+        object_lifecycle_preference: Optional[dict[str, str]] = None,
+        key: Optional[str] = None,
+    ) -> str:
+        destination_path = posixpath.join(
+            key or "",
+            f"{uuid.uuid4().hex}_{data.file_name}",
+        )
+
+        self.storage_client.upload_fileobj(
+            BytesIO(data.data),
+            self.bucket_name,
+            destination_path,
+            ExtraArgs={"ContentType": data.content_type},
+        )
+
+        public_url = self.storage_client.generate_presigned_url(
+            ClientMethod="get_object",
+            Params={"Bucket": self.bucket_name, "Key": destination_path},
+            ExpiresIn=self.url_expiration,
+        )
+        return public_url

fal/toolkit/file/types.py

@@ -2,7 +2,8 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 from mimetypes import guess_extension, guess_type
-from typing import Literal
+from pathlib import Path
+from typing import Literal, Optional
 from uuid import uuid4
 
 
@@ -28,10 +29,34 @@ class FileData:
             self.file_name = file_name
 
 
-RepositoryId = Literal["fal", "fal_v2", "in_memory", "gcp_storage", "r2", "cdn"]
+RepositoryId = Literal[
+    "fal", "fal_v2", "fal_v3", "in_memory", "gcp_storage", "r2", "cdn"
+]
 
 
 @dataclass
 class FileRepository:
-    def save(self, data: FileData) -> str:
+    def save(
+        self,
+        data: FileData,
+        object_lifecycle_preference: Optional[dict[str, str]] = None,
+    ) -> str:
         raise NotImplementedError()
+
+    def save_file(
+        self,
+        file_path: str | Path,
+        content_type: str,
+        multipart: bool | None = None,
+        multipart_threshold: int | None = None,
+        multipart_chunk_size: int | None = None,
+        multipart_max_concurrency: int | None = None,
+        object_lifecycle_preference: Optional[dict[str, str]] = None,
+    ) -> tuple[str, FileData | None]:
+        if multipart:
+            raise NotImplementedError()
+
+        with open(file_path, "rb") as fobj:
+            data = FileData(fobj.read(), content_type, Path(file_path).name)
+
+        return self.save(data, object_lifecycle_preference), data

fal/toolkit/image/__init__.py

@@ -1,3 +1,74 @@
 from __future__ import annotations
 
+import urllib.request
+from functools import lru_cache
+from typing import TYPE_CHECKING
+
 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 = urllib.request.Request(url, headers=TEMP_HEADERS)
+        response = urllib.request.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/toolkit/image/image.py

@@ -4,9 +4,10 @@ import io
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING, Literal, Optional, Union
 
+from fastapi import Request
 from pydantic import BaseModel, Field
 
-from fal.toolkit.file.file import DEFAULT_REPOSITORY, File
+from fal.toolkit.file.file import DEFAULT_REPOSITORY, FALLBACK_REPOSITORY, File
 from fal.toolkit.file.types import FileRepository, RepositoryId
 from fal.toolkit.utils.download_utils import _download_file_python
 
@@ -79,12 +80,18 @@ class Image(File):
         size: ImageSize | None = None,
         file_name: str | None = None,
         repository: FileRepository | RepositoryId = DEFAULT_REPOSITORY,
+        fallback_repository: Optional[
+            FileRepository | RepositoryId
+        ] = FALLBACK_REPOSITORY,
+        request: Optional[Request] = None,
     ) -> Image:
         obj = super().from_bytes(
             data,
             content_type=f"image/{format}",
             file_name=file_name,
             repository=repository,
+            fallback_repository=fallback_repository,
+            request=request,
         )
         obj.width = size.width if size else None
         obj.height = size.height if size else None
@@ -97,6 +104,10 @@ class Image(File):
         format: ImageFormat | None = None,
         file_name: str | None = None,
         repository: FileRepository | RepositoryId = DEFAULT_REPOSITORY,
+        fallback_repository: Optional[
+            FileRepository | RepositoryId
+        ] = FALLBACK_REPOSITORY,
+        request: Optional[Request] = None,
     ) -> Image:
         size = ImageSize(width=pil_image.width, height=pil_image.height)
         if format is None:
@@ -110,12 +121,24 @@ class Image(File):
             # enough result quickly to utilize the underlying resources
             # efficiently.
             saving_options["compress_level"] = 1
+        elif format == "jpeg":
+            # JPEG quality is set to 95 by default, which is a good balance
+            # between file size and image quality.
+            saving_options["quality"] = 95
 
         with io.BytesIO() as f:
             pil_image.save(f, format=format, **saving_options)
             raw_image = f.getvalue()
 
-        return cls.from_bytes(raw_image, format, size, file_name, repository)
+        return cls.from_bytes(
+            raw_image,
+            format,
+            size,
+            file_name,
+            repository,
+            fallback_repository=fallback_repository,
+            request=request,
+        )
 
     def to_pil(self, mode: str = "RGB") -> PILImage.Image:
         try:

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/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/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/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/toolkit/image/nsfw_filter/requirements.txt

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

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/toolkit/types.py

@@ -0,0 +1,140 @@
+import re
+import tempfile
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, Dict, Generator, Union
+
+import pydantic
+from pydantic.utils import update_not_none
+
+from fal.toolkit.image import read_image_from_url
+from fal.toolkit.utils.download_utils import download_file
+
+# https://github.com/pydantic/pydantic/pull/2573
+if not hasattr(pydantic, "__version__") or pydantic.__version__.startswith("1."):
+    IS_PYDANTIC_V2 = False
+else:
+    IS_PYDANTIC_V2 = True
+
+MAX_DATA_URI_LENGTH = 10 * 1024 * 1024
+MAX_HTTPS_URL_LENGTH = 2048
+
+HTTP_URL_REGEX = (
+    r"^https:\/\/(?:[a-zA-Z0-9-]+\.)+[a-zA-Z]{2,}(?::\d{1,5})?(?:\/[^\s]*)?$"
+)
+
+
+class DownloadFileMixin:
+    @contextmanager
+    def as_temp_file(self) -> Generator[Path, None, None]:
+        with tempfile.TemporaryDirectory() as temp_dir:
+            yield download_file(str(self), temp_dir)
+
+
+class DownloadImageMixin:
+    def to_pil(self):
+        return read_image_from_url(str(self))
+
+
+class DataUri(DownloadFileMixin, str):
+    if IS_PYDANTIC_V2:
+
+        @classmethod
+        def __get_pydantic_core_schema__(cls, source_type: Any, handler) -> Any:
+            return {
+                "type": "str",
+                "pattern": "^data:",
+                "max_length": MAX_DATA_URI_LENGTH,
+                "strip_whitespace": True,
+            }
+
+        def __get_pydantic_json_schema__(cls, core_schema, handler) -> Dict[str, Any]:
+            json_schema = handler(core_schema)
+            json_schema.update(format="data-uri")
+            return json_schema
+    else:
+
+        @classmethod
+        def __get_validators__(cls):
+            yield cls.validate
+
+        @classmethod
+        def validate(cls, value: Any) -> "DataUri":
+            from pydantic.validators import str_validator
+
+            value = str_validator(value)
+            value = value.strip()
+
+            if not value.startswith("data:"):
+                raise ValueError("Data URI must start with 'data:'")
+
+            if len(value) > MAX_DATA_URI_LENGTH:
+                raise ValueError(
+                    f"Data URI is too long. Max length is {MAX_DATA_URI_LENGTH} bytes."
+                )
+
+            return cls(value)
+
+        @classmethod
+        def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None:
+            update_not_none(field_schema, format="data-uri")
+
+
+class HttpsUrl(DownloadFileMixin, str):
+    if IS_PYDANTIC_V2:
+
+        @classmethod
+        def __get_pydantic_core_schema__(cls, source_type: Any, handler) -> Any:
+            return {
+                "type": "str",
+                "pattern": HTTP_URL_REGEX,
+                "max_length": MAX_HTTPS_URL_LENGTH,
+                "strip_whitespace": True,
+            }
+
+        def __get_pydantic_json_schema__(cls, core_schema, handler) -> Dict[str, Any]:
+            json_schema = handler(core_schema)
+            json_schema.update(format="https-url")
+            return json_schema
+
+    else:
+
+        @classmethod
+        def __get_validators__(cls):
+            yield cls.validate
+
+        @classmethod
+        def validate(cls, value: Any) -> "HttpsUrl":
+            from pydantic.validators import str_validator
+
+            value = str_validator(value)
+            value = value.strip()
+
+            if not re.match(HTTP_URL_REGEX, value):
+                raise ValueError(
+                    "URL must start with 'https://' and follow the correct format."
+                )
+
+            if len(value) > MAX_HTTPS_URL_LENGTH:
+                raise ValueError(
+                    f"HTTPS URL is too long. Max length is "
+                    f"{MAX_HTTPS_URL_LENGTH} characters."
+                )
+
+            return cls(value)
+
+        @classmethod
+        def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None:
+            update_not_none(field_schema, format="https-url")
+
+
+class ImageHttpsUrl(DownloadImageMixin, HttpsUrl):
+    pass
+
+
+class ImageDataUri(DownloadImageMixin, DataUri):
+    pass
+
+
+FileInput = Union[HttpsUrl, DataUri]
+ImageInput = Union[ImageHttpsUrl, ImageDataUri]

fal/toolkit/utils/download_utils.py

@@ -84,6 +84,9 @@ def _get_remote_file_properties(
             url_path = parsed_url.path
             file_name = Path(url_path).name or _hash_url(url)
 
+    # file name can still contain a forward slash if the server returns a relative path
+    file_name = Path(file_name).name
+
     return file_name, content_length
 
 
@@ -159,6 +162,7 @@ def download_file(
     try:
         file_name = _get_remote_file_properties(url, request_headers)[0]
     except Exception as e:
+        print(f"GOt error: {e}")
         raise DownloadError(f"Failed to get remote file properties for {url}") from e
 
     if "/" in file_name: