make model a parameter to constructor

livekit · Dec 4, 2023 · 1bb8db8 · 1bb8db8
1 parent 1d1645a
commit 1bb8db8
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/livekit-plugins/livekit-plugins-fal/livekit/plugins/fal/sdxl.py b/livekit-plugins/livekit-plugins-fal/livekit/plugins/fal/sdxl.py
@@ -28,11 +28,13 @@ class SDXLPlugin:
     Requires FAL_KEY_ID and FAL_KEY_SECRET environment variables to be set.
     """
 
-    def __init__(self):
+    def __init__(self, model='110602490-fast-sdxl'):
         if not os.getenv("FAL_KEY_ID") and os.getenv("FAL_KEY_SECRET"):
             raise ValueError(
                 "The Fal plugin requires FAL_KEY_ID and FAL_KEY_SECRET environment variables to be set.")
 
+        self.model = model
+
     async def generate_image_from_prompt(self, prompt: str) -> rtc.VideoFrame:
         """Generate an image from a prompt
 
@@ -44,7 +46,7 @@ async def generate_image_from_prompt(self, prompt: str) -> rtc.VideoFrame:
         """
 
         handler = fal.apps.submit(
-            '110602490-fast-sdxl',
+            self.model,
             arguments={
                 'prompt': prompt,
                 'sync_mode': False