From 174635d90359e1900623eabe4123b41b4899fb70 Mon Sep 17 00:00:00 2001
From: hodanov <1031hoda@gmail.com>
Date: Tue, 18 Jul 2023 23:12:18 +0900
Subject: [PATCH] Modify a little.

---
 setup_files/__main__.py | 6 +-----
 setup_files/txt2img.py  | 4 ++++
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/setup_files/__main__.py b/setup_files/__main__.py
index 0075491..7500a76 100644
--- a/setup_files/__main__.py
+++ b/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")
diff --git a/setup_files/txt2img.py b/setup_files/txt2img.py
index bc12a6d..748d4c8 100644
--- a/setup_files/txt2img.py
+++ b/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.