diff --git a/setup_files/__main__.py b/setup_files/__main__.py
index e5fb5f5..833fecb 100644
--- a/setup_files/__main__.py
+++ b/setup_files/__main__.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from setup import stub
-from txt2img import StableDiffusion, new_stable_diffusion
+from txt2img import new_stable_diffusion
 
 
 @stub.function(gpu="A10G")