Update app.py

app.py

@@ -1,36 +1,15 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline
-from safetensors.torch import load_file
-import os
+from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
 
-# Hugging Face 토큰 환경 변수에서 가져오기
-HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+model_id = "gagong/Traditional-Korean-Painting-Model-v2.0"
+scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
+pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, torch_dtype=torch.float16)
+pipe = pipe.to("cuda")
 
-# safetensors 모델 로드 함수
-def load_safetensors_model(model_id, model_file, use_auth_token):
-    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"  # 루트 디렉토리에 있는 파일 경로 설정
+# prompt = "a photo of an astronaut riding a horse on mars"
+# image = pipe(prompt).images[0]
 
-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")
-else:
-    print("Failed to load the model")
 
 def generate_image(prompt):
     if not pipe: