update

app.py

@@ -6,10 +6,11 @@ import uuid
 
 import gradio as gr
 import numpy as np
-import spaces
+#import spaces
 import torch
 from PIL import Image
 from evosdxl_jp_v1 import load_evosdxl_jp
+import devicetorch
 
 DESCRIPTION = """# 🐟 EvoSDXL-JP
 🤗 [モデル一覧](https://huggingface.co/SakanaAI) | 📚 [技術レポート](https://arxiv.org/abs/2403.13187) | 📝 [ブログ](https://sakana.ai/evosdxl-jp/) | 🐦 [Twitter](https://twitter.com/SakanaAILabs)
@@ -23,12 +24,14 @@ if not torch.cuda.is_available():
 MAX_SEED = np.iinfo(np.int32).max
 CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1"
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+#device = "cuda" if torch.cuda.is_available() else "cpu"
+device = devicetorch.get(torch)
 
 NUM_IMAGES_PER_PROMPT = 1
 ENABLE_CPU_OFFLOAD = False
 USE_TORCH_COMPILE = False
-SAFETY_CHECKER = True
+#SAFETY_CHECKER = True
+SAFETY_CHECKER = False
 DEVELOP_MODE = True
 if SAFETY_CHECKER:
     from safety_checker import StableDiffusionSafetyChecker
@@ -53,7 +56,8 @@ if SAFETY_CHECKER:
         return images, has_nsfw_concepts
 
 
-pipe = load_evosdxl_jp("cpu").to("cuda")
+#pipe = load_evosdxl_jp("cpu").to("cuda")
+pipe = load_evosdxl_jp("cpu").to(device)
 
 def show_warning(warning_text: str) -> gr.Blocks:
     with gr.Blocks() as demo:
@@ -154,4 +158,4 @@ with gr.Blocks(css=css) as demo:
                 Sakana AIは、本モデルの使用によって生じた直接的または間接的な損失に対して、結果に関わらず、一切の責任を負いません。
                 利用者は、本モデルの使用に伴うリスクを十分に理解し、自身の判断で使用することが必要です。""")
 
-demo.queue().launch()
+demo.queue().launch()

evosdxl_jp_v1.py

@@ -11,6 +11,7 @@ from diffusers import (
     EulerDiscreteScheduler,
 )
 from diffusers.loaders import LoraLoaderMixin
+import devicetorch
 
 SDXL_REPO = "stabilityai/stable-diffusion-xl-base-1.0"
 JSDXL_REPO = "stabilityai/japanese-stable-diffusion-xl"
@@ -137,7 +138,8 @@ def load_evosdxl_jp(device="cuda") -> StableDiffusionXLPipeline:
         ],
     )
     del sdxl_weights, dpo_weights, jn_weights, jsdxl_weights
-    torch.cuda.empty_cache()
+    devicetorch.empty_cache(torch)
+    #torch.cuda.empty_cache()
     unet_config = UNet2DConditionModel.load_config(SDXL_REPO, subfolder="unet")
     unet = UNet2DConditionModel.from_config(unet_config).to(device=device)
     unet.load_state_dict({**new_conv, **new_attn})

requirements.txt

@@ -1,6 +1,6 @@
-torch
+#torch
 diffusers==0.26.0
 transformers
 safetensors
 accelerate
-sentencepiece
+sentencepiece