Update app.py

app.py

@@ -48,7 +48,28 @@ import cv2
 import numpy as np
 import sys
 import io
+# FluxPipeline import 부분을 수정
+from diffusers import StableDiffusionPipeline, DiffusionPipeline
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+
+# 기존의 pipe 초기화 부분을 다음과 같이 수정
+pipe = StableDiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-2-1",
+    torch_dtype=torch.float16,
+    safety_checker=None
+)
+pipe.to("cuda")
+
+# LoRA 가중치 로드
+pipe.load_lora_weights(
+    hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors")
+)
+pipe.fuse_lora(lora_scale=0.125)
 
+# 안전 검사기 설정
+pipe.safety_checker = safety_checker.StableDiffusionSafetyChecker.from_pretrained(
+    "CompVis/stable-diffusion-safety-checker"
+)
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -138,15 +159,6 @@ class timer:
         end = time.time()
         print(f"{self.method} took {str(round(end - self.start, 2))}s")
 
-# Model initialization
-if not path.exists(cache_path):
-    os.makedirs(cache_path, exist_ok=True)
-
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
-pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
-pipe.fuse_lora(lora_scale=0.125)
-pipe.to(device="cuda", dtype=torch.bfloat16)
-pipe.safety_checker = safety_checker.StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
 
 def upload_to_catbox(image_path):
     """catbox.moe API를 사용하여 이미지 업로드"""