diff --git a/.github/workflows/lint_python.yml b/.github/workflows/lint_python.yml
index 0c94b00..6c9533d 100644
--- a/.github/workflows/lint_python.yml
+++ b/.github/workflows/lint_python.yml
@@ -20,5 +20,5 @@ jobs:
           pip install pycodestyle pyflakes
       - name: Analysing the code with pycodestyle
         run: |
-          pycodestyle --first --ignore='E501' $(git ls-files '*.py')
+          pycodestyle --first --ignore='E501,E401' $(git ls-files '*.py')
           pyflakes $(git ls-files '*.py')
diff --git a/sdcli/util.py b/sdcli/util.py
index 32708c3..643be2e 100644
--- a/sdcli/util.py
+++ b/sdcli/util.py
@@ -43,13 +43,13 @@ def save_prompts(inputs: dict):
         print(f"Save prompts: {prompts_filename}.txt")
 
 
-def save_images(directory: Path, images: list[bytes], seed: int, i: int):
+def save_images(directory: Path, images: list[bytes], seed: int, i: int, output_format: str = "png"):
     """
     Save images to a file.
     """
     for j, image_bytes in enumerate(images):
         formatted_time = time.strftime("%Y%m%d%H%M%S", time.localtime(time.time()))
-        output_path = directory / f"{formatted_time}_{seed}_{i}_{j}.png"
+        output_path = directory / f"{formatted_time}_{seed}_{i}_{j}.{output_format}"
         print(f"Saving it to {output_path}")
         with open(output_path, "wb") as file:
             file.write(image_bytes)
diff --git a/setup_files/__main__.py b/setup_files/__main__.py
index 7fb70fe..d1ff313 100644
--- a/setup_files/__main__.py
+++ b/setup_files/__main__.py
@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+from setup import stub
 from txt2img import StableDiffusion
 
 
+@stub.function(gpu="A10G")
 def main():
     StableDiffusion
 
diff --git a/setup_files/requirements.txt b/setup_files/requirements.txt
index c3cc197..b04530e 100644
--- a/setup_files/requirements.txt
+++ b/setup_files/requirements.txt
@@ -14,6 +14,7 @@ gfpgan>=1.3.8
 scipy==1.11.4
 opencv-python
 Pillow
+pillow-avif-plugin
 torchvision
 tqdm
 
diff --git a/setup_files/txt2img.py b/setup_files/txt2img.py
index 0ed9d65..0a4907f 100644
--- a/setup_files/txt2img.py
+++ b/setup_files/txt2img.py
@@ -3,9 +3,7 @@ from __future__ import annotations
 import io
 import os
 
-import diffusers
 import PIL.Image
-import torch
 from modal import Secret, method
 from setup import (
     BASE_CACHE_PATH,
@@ -26,6 +24,8 @@ class StableDiffusion:
     """
 
     def __enter__(self):
+        import diffusers
+        import torch
         import yaml
 
         config = {}
@@ -133,13 +133,17 @@ class StableDiffusion:
         upscaler: str = "",
         use_face_enhancer: bool = False,
         fix_by_controlnet_tile: bool = False,
+        output_format: str = "png",
     ) -> list[bytes]:
         """
         Runs the Stable Diffusion pipeline on the given prompt and outputs images.
         """
+        import pillow_avif  # noqa: F401
+        import torch
+
         max_embeddings_multiples = self._count_token(p=prompt, n=n_prompt)
         generator = torch.Generator("cuda").manual_seed(seed)
-        self.pipe = self.pipe.to("cuda")
+        self.pipe.to("cuda")
         self.pipe.enable_vae_tiling()
         self.pipe.enable_xformers_memory_efficient_attention()
         with torch.autocast("cuda"):
@@ -161,7 +165,7 @@ class StableDiffusion:
         https://huggingface.co/lllyasviel/control_v11f1e_sd15_tile
         """
         if fix_by_controlnet_tile:
-            self.controlnet_pipe = self.controlnet_pipe.to("cuda")
+            self.controlnet_pipe.to("cuda")
             self.controlnet_pipe.enable_vae_tiling()
             self.controlnet_pipe.enable_xformers_memory_efficient_attention()
             for image in base_images:
@@ -193,7 +197,7 @@ class StableDiffusion:
         image_output = []
         for image in generated_images:
             with io.BytesIO() as buf:
-                image.save(buf, format="PNG")
+                image.save(buf, format=output_format)
                 image_output.append(buf.getvalue())
 
         return image_output