Update app.py

app.py

@@ -4,7 +4,6 @@ from diffusers import StableDiffusionPipeline
 from safetensors.torch import load_file
 import os
 
-
 # safetensors 모델 로드 함수
 def load_safetensors_model(model_id, model_file, use_auth_token):
     # safetensors 파일 로드
@@ -15,19 +14,21 @@ def load_safetensors_model(model_id, model_file, use_auth_token):
     model.model.load_state_dict(state_dict)
     return model
 
-model_id = "gagong/Traditional-Korean-Painting-Model-v2.0"
+# 올바른 모델 ID로 설정
+model_id = "stabilityai/stable-diffusion-2"
 model_file = "./Traditional_Korean_Painting_Model_2.safetensors"  # 루트 디렉토리에 있는 파일 경로 설정
 
 try:
     pipe = load_safetensors_model(model_id, model_file, use_auth_token=HUGGINGFACE_TOKEN)
     pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+    print("Model loaded successfully")
 except Exception as e:
     print(f"Error loading model: {e}")
     pipe = None
 
 def generate_image(prompt):
     if not pipe:
-        return None
+        return "Model not loaded properly"
     image = pipe(prompt).images[0]
     return image