Update app.py

app.py

@@ -4,33 +4,43 @@ from diffusers import StableDiffusionPipeline
 from safetensors.torch import load_file
 import os
 
+# Hugging Face 토큰 환경 변수에서 가져오기
+HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+
 # safetensors 모델 로드 함수
 def load_safetensors_model(model_id, model_file, use_auth_token):
-    # safetensors 파일 로드
-    state_dict = load_file(model_file)
-    # 모델 초기화
-    model = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=use_auth_token)
-    # 모델에 state_dict 로드
-    model.model.load_state_dict(state_dict)
-    return model
+    try:
+        # safetensors 파일 로드
+        state_dict = load_file(model_file)
+        # 모델 초기화
+        model = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=use_auth_token)
+        # 모델에 state_dict 로드
+        model.model.load_state_dict(state_dict)
+        return model
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        return None
 
 # 올바른 모델 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 = load_safetensors_model(model_id, model_file, use_auth_token=HUGGINGFACE_TOKEN)
+if pipe:
     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
+else:
+    print("Failed to load the model")
 
 def generate_image(prompt):
     if not pipe:
         return "Model not loaded properly"
-    image = pipe(prompt).images[0]
-    return image
+    try:
+        image = pipe(prompt).images[0]
+        return image
+    except Exception as e:
+        print(f"Error generating image: {e}")
+        return "Error generating image"
 
 # Gradio 인터페이스 설정
 with gr.Blocks() as demo:
@@ -42,10 +52,9 @@ with gr.Blocks() as demo:
             prompt = gr.Textbox(label="Prompt", placeholder="Describe the scene...")
             generate_btn = gr.Button("Generate")
         with gr.Column():
-            output_image = gr.Image(label="Generated Image")
-    
+            output_image = gr.Image(label="Generated Image", type="pil")
+
     generate_btn.click(fn=generate_image, inputs=prompt, outputs=output_image)
 
 if __name__ == "__main__":
     demo.launch()
-