Update space

model/animagine_xl.py

@@ -1,45 +1,47 @@
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
-import torch
-
-def generate_animagine_xl(prompt: str):
-    model_id = "Linaqruf/animagine-xl"
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-        variant="fp16"
-    )
-    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-    pipe = pipe.to("cuda")
-    image = pipe(prompt=prompt, width=1024, height=1024).images[0]
-    image.save("output_animagine_xl.png")
-    print("✅ 저장 완료: output_animagine_xl.png")
-    return image
-
-if __name__ == "__main__":
-    prompt = "그녀를 바라보는 한 남자의 야망"
-    generate_animagine_xl(prompt)
-
-
 # from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 # import torch
 
 # def generate_animagine_xl(prompt: str):
 #     model_id = "Linaqruf/animagine-xl"
-#     use_fp16 = torch.cuda.is_available()
 #     pipe = StableDiffusionXLPipeline.from_pretrained(
 #         model_id,
-#         torch_dtype=torch.float16 if use_fp16 else torch.float32,
+#         torch_dtype=torch.float16,
 #         use_safetensors=True,
 #         variant="fp16"
 #     )
 #     pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-#     device = "cuda" if use_fp16 else "cpu"
-#     pipe = pipe.to(device)
+#     pipe = pipe.to("cuda")
 #     image = pipe(prompt=prompt, width=1024, height=1024).images[0]
 #     image.save("output_animagine_xl.png")
+#     print("✅ 저장 완료: output_animagine_xl.png")
 #     return image
 
 # if __name__ == "__main__":
 #     prompt = "그녀를 바라보는 한 남자의 야망"
 #     generate_animagine_xl(prompt)
+
+
+from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+import torch
+
+# (1) 모델과 스케줄러를 전역에서 한 번만 초기화
+model_id = "Linaqruf/animagine-xl"
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    variant="fp16"
+)
+pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+pipe = pipe.to("cuda")  # 또는 "cpu"로 변경 가능
+
+# (2) 이미지 생성 함수
+def generate_animagine_xl(prompt: str):
+    image = pipe(prompt=prompt, width=1024, height=1024).images[0]
+    image.save("output_animagine_xl.png")
+    print("✅ 저장 완료: output_animagine_xl.png")
+    return image
+
+if __name__ == "__main__":
+    prompt = "그녀를 바라보는 한 남자의 야망"
+    generate_animagine_xl(prompt)