Update space

app.py

@@ -9,6 +9,9 @@ from model.waifu import generate_waifu
 from Text2Long_text import generate_korean_story
 
 # 모델 함수 매핑
+
+model_id = "hakurei/waifu-diffusion"
+
 MODEL_FUNCTIONS = {
     "Animagine XL": generate_animagine_xl,
     "Animesai": generate_animesai,
@@ -16,7 +19,7 @@ MODEL_FUNCTIONS = {
     "Ghibli": generate_ghibli,
     "Realistic": generate_realistic,
     "SD Turbo": generate_sd_turbo,
-    "Waifu Diffusion": generate_waifu,
+    "Waifu Diffusion": generate_waifu(model_id),
     "한국어 긴 이야기 생성": lambda prompt: generate_korean_story(prompt, max_length=100),
 }
 

model/animagine_xl.py

@@ -10,14 +10,36 @@ def generate_animagine_xl(prompt: str):
         variant="fp16"
     )
     pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-    device = "cuda" if torch.cuda.is_available() 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
+
+# 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,
+#         use_safetensors=True,
+#         variant="fp16"
+#     )
+#     pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+#     device = "cuda" if use_fp16 else "cpu"
+#     pipe = pipe.to(device)
+#     image = pipe(prompt=prompt, width=1024, height=1024).images[0]
+#     image.save("output_animagine_xl.png")
+#     return image
+
+# if __name__ == "__main__":
+#     prompt = "그녀를 바라보는 한 남자의 야망"
+#     generate_animagine_xl(prompt)

model/waifu.py

@@ -1,21 +1,46 @@
+# from diffusers import StableDiffusionPipeline
+# import torch
+
+# def generate_waifu(prompt: str):
+#     model_id = "hakurei/waifu-diffusion"
+#     pipe = StableDiffusionPipeline.from_pretrained(
+#         model_id,
+#         torch_dtype=torch.float16,
+#         use_safetensors=False,
+#         revision="fp16"
+#     ).to("cuda")
+
+#     image = pipe(prompt=prompt).images[0]
+#     return image
+
+# if __name__ == "__main__":
+#     prompt = '그녀를 바라보는 한 남자의 야망'
+#     image = generate_waifu(prompt)
+#     image.save("output_waifu.png")
+#     print("✅ 저장 완료: output_waifu.png")
+#     image.show()
+
 from diffusers import StableDiffusionPipeline
 import torch
 
-def generate_waifu(prompt: str):
+
+def generate_waifu(prompt: str, model_id):
     model_id = "hakurei/waifu-diffusion"
     pipe = StableDiffusionPipeline.from_pretrained(
         model_id,
-        torch_dtype=torch.float16,
-        use_safetensors=False,
-        revision="fp16"
-    ).to("cuda")
+        torch_dtype=torch.float32,         # CPU는 float32만 지원
+        use_safetensors=False
+        # revision=None                       # revision="fp16" 삭제
+    ).to("cpu")                            # 반드시 CPU로 지정
+
 
     image = pipe(prompt=prompt).images[0]
     return image
 
+
 if __name__ == "__main__":
     prompt = '그녀를 바라보는 한 남자의 야망'
     image = generate_waifu(prompt)
-    image.save("output_waifu.png")
-    print("✅ 저장 완료: output_waifu.png")
-    image.show()
+    image.save("output_waifu_cpu.png")
+    print("✅ 저장 완료: output_waifu_cpu.png")
+    image.show()