feat: `fal.App` for multiple endpoints #27

isidentical · 2024-01-09T21:35:10Z

~~mainly for discussions on the idea, not the concrete implementation~~ Let's see if we can get it into a decent state to start using it.

import fal
from fal.toolkit import Image, ImageSizeInput, get_image_size
from pydantic import BaseModel, Field


class InputModel(BaseModel):
    prompt: str
    seed: int = Field(default=42, ge=0, le=2**32 - 1)


class OutputModel(BaseModel):
    images: list[Image]
    seed: int


class Text2ImageInputModel(InputModel):
    image_size: ImageSizeInput = "square_hd"


class Image2ImageInputModel(InputModel):
    image_url: str
    strength: float = Field(default=0.5, ge=0.0, le=1.0)


class InpaintingInputModel(InputModel):
    image_url: str
    mask_url: str
    strength: float = Field(default=0.5, ge=0.0, le=1.0)


class StableDiffusion(fal.App, _scheduler="nomad"):
    machine_type = "GPU"
    requirements = [
        "diffusers==0.25.0",
        "transformers",
        "torch>=2.1",
        "accelerate",
    ]

    def setup(self):
        import torch
        from diffusers import (
            AutoPipelineForText2Image,
            AutoPipelineForImage2Image,
            AutoPipelineForInpainting,
        )

        self.pipeline_text2img = AutoPipelineForText2Image.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            torch_dtype=torch.float16,
            use_safetensors=True,
        ).to("cuda")
        self.pipeline_img2img = AutoPipelineForImage2Image.from_pipe(
            self.pipeline_text2img
        )
        self.pipeline_inpainting = AutoPipelineForInpainting.from_pipe(
            self.pipeline_text2img
        )

    def parse_image_url(self, url: str) -> object:
        from urllib.request import urlopen
        from PIL import Image

        with urlopen(url) as stream:
            return Image.open(stream)

    @fal.endpoint("/text-to-image")
    def text_to_image(self, input: Text2ImageInputModel) -> OutputModel:
        import torch

        image_size = get_image_size(input.image_size)
        result = self.pipeline_text2img(
            prompt=input.prompt,
            generator=torch.Generator("cuda").manual_seed(input.seed),
            width=image_size.width,
            height=image_size.height,
        )
        return OutputModel(
            images=[Image.from_pil(image) for image in result.images],
            seed=input.seed,
        )

    @fal.endpoint("/image-to-image")
    def image_to_image(self, input: Image2ImageInputModel) -> OutputModel:
        import torch

        result = self.pipeline_img2img(
            prompt=input.prompt,
            image=self.parse_image_url(input.image_url),
            generator=torch.Generator("cuda").manual_seed(input.seed),
            strength=input.strength,
        )
        return OutputModel(
            images=[Image.from_pil(image) for image in result.images],
            seed=input.seed,
        )

    @fal.endpoint("/inpainting")
    def inpainting(self, input: InpaintingInputModel) -> OutputModel:
        import torch

        result = self.pipeline_inpainting(
            prompt=input.prompt,
            image=self.parse_image_url(input.image_url),
            mask=self.parse_image_url(input.mask_url),
            generator=torch.Generator("cuda").manual_seed(input.seed),
            strength=input.strength,
        )
        return OutputModel(
            images=[Image.from_pil(image) for image in result.images],
            seed=input.seed,
        )


if __name__ == "__main__":
    # SDK usage, TBD
    app = StableDiffusion() # returns a shallow proxy object which
                            # when called, will create the stateful
                            # app if it is not already in the process
                            # and treat the calls as if they were coming
                            # in a local python process.
    result = app.text_to_image(...)

$ fal run t.py StableDiffusion

mederka · 2024-01-09T21:42:40Z

Is this how you deploy?

fal run t.py StableDiffusion

isidentical · 2024-01-09T21:43:41Z

no, that's how you start the local test server (compared to just calling it we have now), the deployment flow is the same:

❯ fal fn run t.py StableDiffusion
2024-01-09 21:39:08.976 [info     ] !!! HEADS UP !!! Scheduling a job with <isolate_controller.scheduler.nomad.scheduler.NomadJobScheduler object at 0x7f7fa1516ad0>
2024-01-09 21:39:11.370 [info     ] Access your exposed service at https://5d0299e5-479d-48a3-b26a-8fb38207ba17.gateway.alpha.fal.ai
2024-01-09 21:39:17.996 [stderr   ] 
2024-01-09 21:39:17.996 [stderr   ] Loading pipeline components...:   0%|          | 0/7 [00:00<?, ?it/s]
2024-01-09 21:39:18.714 [stderr   ] 
...

❯ fal fn serve t.py StableDiffusion --alias lolz-diffusion       
Registered a new revision for function 'lolz-diffusion' (revision='4eaeab3c-1619-4ef5-97fa-5dcc8edd7488').
URL: https://47358913-lolz-diffusion.gateway.alpha.fal.ai

mederka · 2024-01-09T21:45:39Z

How come _scheduler is a parameter?

isidentical · 2024-01-09T21:48:33Z

How come _scheduler is a parameter?

There is a set of generic parameters (the same set of arguments we have at @fal.function) and a set of host-specific parameters (fal vs local). For the latter we need to either use a dict (host_options = {"a": "b"}) or pass them as parameters in the class, although the former can also be supported in the class class T(fal.App, requirements=[...], a=b).

Open for comments on this, we can have a single way to configure everything (but it might become super cluttered due to requirements) OR have two different segments which configure two different things (the current way).

mederka · 2024-01-09T21:57:44Z

So, when testing locally, I might need to test the fal.endpoints locally as well. How would we do this?

isidentical · 2024-01-09T22:03:35Z

Given the application above, when developing, you just edit your code and run fal fn run t.py StableDiffusion (which gives you an endpoint you can perform tests against). Once you are ready to deploy, you use the fal fn serve t.py StableDiffusion as if it was just a single @fal.function (w/same options etc.).

squat · 2024-01-10T20:08:58Z

projects/fal/src/fal/app.py

+    def marker_fn(callable: EndpointT) -> EndpointT:
+        if hasattr(callable, "route_signature"):
+            raise ValueError(
+                f"Can't set multiple routes for the same function: {callable.__name__}"


Why is that?

we don't handle it yet, but in the future I want to be able to support multiple endpoints for a single function by stacking @fal.endpoint. this is done to reserve that use case (instead of treating it as an override which we would need to break)

aha ok 👍

squat

great idea and work!

chamini2 · 2024-01-12T02:53:46Z

projects/fal/src/fal/app.py

+    def _build_app(self) -> FastAPI:
+        from fastapi import FastAPI
+        from fastapi.middleware.cors import CORSMiddleware
+
+        _app = FastAPI()
+
+        _app.add_middleware(
+            CORSMiddleware,
+            allow_credentials=True,
+            allow_headers=("*"),
+            allow_methods=("*"),
+            allow_origins=("*"),
+        )
+
+        routes: dict[RouteSignature, Callable[..., Any]] = {
+            signature: endpoint
+            for _, endpoint in inspect.getmembers(self, inspect.ismethod)
+            if (signature := getattr(endpoint, "route_signature", None))
+        }
+        if not routes:
+            raise ValueError("An application must have at least one route!")
+
+        for signature, endpoint in routes.items():
+            _app.add_api_route(
+                signature.path,
+                endpoint,
+                name=endpoint.__name__,
+                methods=["POST"],
+            )
+
+        return _app
+
+    def openapi(self) -> dict[str, Any]:
+        """
+        Build the OpenAPI specification for the served function.
+        Attach needed metadata for a better integration to fal.
+        """
+        app = self._build_app()
+        spec = app.openapi()
+        self._mark_order_openapi(spec)
+        return spec
+
+    def _mark_order_openapi(self, spec: dict[str, Any]):
+        """
+        Add x-fal-order-* keys to the OpenAPI specification to help the rendering of UI.
+
+        NOTE: We rely on the fact that fastapi and Python dicts keep the order of properties.
+        """
+
+        def mark_order(obj: dict[str, Any], key: str):
+            obj[f"x-fal-order-{key}"] = list(obj[key].keys())
+
+        mark_order(spec, "paths")
+
+        def order_schema_object(schema: dict[str, Any]):
+            """
+            Mark the order of properties in the schema object.
+            They can have 'allOf', 'properties' or '$ref' key.
+            """
+            if "allOf" in schema:
+                for sub_schema in schema["allOf"]:
+                    order_schema_object(sub_schema)
+            if "properties" in schema:
+                mark_order(schema, "properties")
+
+        for key in spec["components"].get("schemas") or {}:
+            order_schema_object(spec["components"]["schemas"][key])
+
+        return spec


this whole thing already existed for serve functions, can we merge them?

isidentical force-pushed the multiple-endpoints branch from 6724ce2 to d8b540d Compare January 9, 2024 21:36

isidentical force-pushed the multiple-endpoints branch from 43ea98a to 63347db Compare January 10, 2024 15:34

isidentical added 2 commits January 10, 2024 18:34

wip: feat: fal.App for multiple endpoints

d55dd15

fix: build the metadata

1ed49e6

isidentical force-pushed the multiple-endpoints branch 2 times, most recently from 1adae46 to 31f9c63 Compare January 10, 2024 15:50

isidentical marked this pull request as ready for review January 10, 2024 15:56

add tests

6340ecb

isidentical force-pushed the multiple-endpoints branch from 31f9c63 to 6340ecb Compare January 10, 2024 16:58

isidentical changed the title ~~wip: feat: fal.App for multiple endpoints~~ feat: fal.App for multiple endpoints Jan 10, 2024

turbo1912 approved these changes Jan 10, 2024

View reviewed changes

squat reviewed Jan 10, 2024

View reviewed changes

squat approved these changes Jan 10, 2024

View reviewed changes

isidentical merged commit 74afc08 into main Jan 10, 2024
4 checks passed

isidentical deleted the multiple-endpoints branch January 10, 2024 20:19

chamini2 reviewed Jan 12, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: `fal.App` for multiple endpoints #27

feat: `fal.App` for multiple endpoints #27

isidentical commented Jan 9, 2024 •

edited

Loading

mederka commented Jan 9, 2024

isidentical commented Jan 9, 2024

mederka commented Jan 9, 2024

isidentical commented Jan 9, 2024

mederka commented Jan 9, 2024

isidentical commented Jan 9, 2024

squat Jan 10, 2024

isidentical Jan 10, 2024

squat Jan 10, 2024

squat left a comment

chamini2 Jan 12, 2024

feat: fal.App for multiple endpoints #27

feat: fal.App for multiple endpoints #27

Conversation

isidentical commented Jan 9, 2024 • edited Loading

mederka commented Jan 9, 2024

isidentical commented Jan 9, 2024

mederka commented Jan 9, 2024

isidentical commented Jan 9, 2024

mederka commented Jan 9, 2024

isidentical commented Jan 9, 2024

squat Jan 10, 2024

Choose a reason for hiding this comment

isidentical Jan 10, 2024

Choose a reason for hiding this comment

squat Jan 10, 2024

Choose a reason for hiding this comment

squat left a comment

Choose a reason for hiding this comment

chamini2 Jan 12, 2024

Choose a reason for hiding this comment

feat: `fal.App` for multiple endpoints #27

feat: `fal.App` for multiple endpoints #27

isidentical commented Jan 9, 2024 •

edited

Loading