localize

app.py

@@ -1,13 +1,12 @@
 import gradio as gr
 import torch
+import devicetorch
 from diffusers import StableDiffusionXLPipeline, StableDiffusionPipeline, LCMScheduler
 from diffusers.schedulers import TCDScheduler
 
 import spaces
 from PIL import Image
 
-SAFETY_CHECKER = True
-
 checkpoints = {
     "2-Step": ["pcm_{}_smallcfg_2step_converted.safetensors", 2, 0.0],
     "4-Step": ["pcm_{}_smallcfg_4step_converted.safetensors", 4, 0.0],
@@ -26,37 +25,26 @@ checkpoints = {
 
 loaded = None
 
-if torch.cuda.is_available():
-    pipe_sdxl = StableDiffusionXLPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-xl-base-1.0",
-        torch_dtype=torch.float16,
-        variant="fp16",
-    ).to("cuda")
-    pipe_sd15 = StableDiffusionPipeline.from_pretrained(
-        "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16"
-    ).to("cuda")
-
-if SAFETY_CHECKER:
-    from safety_checker import StableDiffusionSafetyChecker
-    from transformers import CLIPFeatureExtractor
-
-    safety_checker = StableDiffusionSafetyChecker.from_pretrained(
-        "CompVis/stable-diffusion-safety-checker"
-    ).to("cuda")
-    feature_extractor = CLIPFeatureExtractor.from_pretrained(
-        "openai/clip-vit-base-patch32"
-    )
-
-    def check_nsfw_images(
-        images: list[Image.Image],
-    ) -> tuple[list[Image.Image], list[bool]]:
-        safety_checker_input = feature_extractor(images, return_tensors="pt").to("cuda")
-        has_nsfw_concepts = safety_checker(
-            images=[images], clip_input=safety_checker_input.pixel_values.to("cuda")
-        )
-
-        return images, has_nsfw_concepts
-
+device = devicetorch.get(torch)
+
+#if torch.cuda.is_available():
+#    pipe_sdxl = StableDiffusionXLPipeline.from_pretrained(
+#        "stabilityai/stable-diffusion-xl-base-1.0",
+#        torch_dtype=torch.float16,
+#        variant="fp16",
+#    ).to("cuda")
+#    pipe_sd15 = StableDiffusionPipeline.from_pretrained(
+#        "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16"
+#    ).to("cuda")
+
+pipe_sdxl = StableDiffusionXLPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.float16,
+    variant="fp16",
+).to(device)
+pipe_sd15 = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, variant="fp16"
+).to(device)
 
 @spaces.GPU(enable_queue=True)
 def generate_image(

requirements.txt

@@ -3,6 +3,7 @@ datasets
 transformers
 accelerate
 peft
-xformers
+#xformers
 gradio==4.32.2
-spaces
+#spaces
+devicetorch