from comfykiosk import ComfyKiosk, ComfyGenerator, SimpleWorkflowLoader
from comfykiosk.image_sources.filesystem import FileSystemImageSource
from comfykiosk.image_sources.pregenerate import PreparedGenPool
from comfykiosk.fastapi import create_app
from fastapi.middleware.cors import CORSMiddleware

# You can compose ComfyKiosk as an alternate way to override defaults
filesystem_source = FileSystemImageSource(directory="./output")
comfyui_backend = ComfyGenerator()

# You can also use the default workflow loader
loader = SimpleWorkflowLoader(directory="./sample-workflows")

# In this example, each workflow can have up to 10 prepared images.
# The generator will make 5 images before switching workflows.
# This is to amortize the time spent loading models into VRAM.
pregen_pool = PreparedGenPool(bucket_max=10, batch_size=5, generator=comfyui_backend, saver=filesystem_source)

# When a request comes in, each image source will be tried in this order.
comfy = ComfyKiosk(loaders=loader,
                   image_sources=[filesystem_source, pregen_pool])

# Attach FastAPI routes to it
app = create_app(comfy)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info", log_config=None)