Modify a little.

setup_files/__main__.py

@@ -1,11 +1,7 @@
 from __future__ import annotations
 
 from setup import stub
-from txt2img import StableDiffusion
-
-
-def new_stable_diffusion() -> StableDiffusion:
-    return StableDiffusion()
+from txt2img import StableDiffusion, new_stable_diffusion
 
 
 @stub.function(gpu="A10G")

setup_files/txt2img.py

@@ -14,6 +14,10 @@ from setup import (BASE_CACHE_PATH, BASE_CACHE_PATH_CONTROLNET,
                    stub)
 
 
+def new_stable_diffusion() -> StableDiffusionInterface:
+    return StableDiffusion()
+
+
 class StableDiffusionInterface(metaclass=abc.ABCMeta):
     """
     A StableDiffusionInterface is an interface that will be used for StableDiffusion class creation.