Update app.py

app.py

@@ -1,19 +1,24 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import StableDiffusionPipeline
+import torch
 
-# Hugging Face의 트랜스포머 라이브러리를 사용해 모델을 로드합니다.
-model = pipeline("text-classification", model="InstantX/InstantID")
+# 모델 초기화
+pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", 
+                                                use_auth_token=True)
+pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
-# 모델을 호출하는 함수를 정의합니다.
-def predict(text):
-    return model(text)
+# 텍스트 프롬프트로부터 이미지를 생성하는 함수
+def generate_image(prompt):
+    with torch.autocast("cuda"):
+        image = pipe(prompt).images[0]
+    return image
 
-# Gradio 인터페이스를 만듭니다.
+# Gradio 인터페이스 정의
 iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Textbox(lines=2, placeholder="Type something here..."),
-    outputs=[gr.Label(num_top_classes=3), gr.Dataframe()],
+    fn=generate_image,
+    inputs=gr.Textbox(label="Enter a prompt for the AI to generate an image"),
+    outputs=gr.Image(type="pil"),
 )
 
-# 인터페이스를 실행하여 테스트합니다.
+# 인터페이스 실행
 iface.launch()