Update app.py

app.py

@@ -16,9 +16,9 @@ from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipelin
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-#DESCRIPTION = "Animagine XL 3.1"
-#if not torch.cuda.is_available():
-#    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU. </p>"
+DESCRIPTION = "Animagine XL 3.1"
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo is slower on CPU. </p>"
 IS_COLAB = utils.is_google_colab() or os.getenv("IS_COLAB") == "1"
 HF_TOKEN = os.getenv("HF_TOKEN")
 CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
@@ -42,7 +42,7 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 def load_pipeline(model_name):
     vae = AutoencoderKL.from_pretrained(
         "madebyollin/sdxl-vae-fp16-fix",
-        torch_dtype=torch.float16,
+        torch_dtype=torch.float32,
     )
     pipeline = (
         StableDiffusionXLPipeline.from_single_file
@@ -53,7 +53,7 @@ def load_pipeline(model_name):
     pipe = pipeline(
         model_name,
         vae=vae,
-        torch_dtype=torch.float16,
+        torch_dtype=torch.float32,
         custom_pipeline="lpw_stable_diffusion_xl",
         use_safetensors=True,
         add_watermarker=False,
@@ -64,7 +64,6 @@ def load_pipeline(model_name):
     return pipe
 
 
-@spaces.GPU
 def generate(
     prompt: str,
     negative_prompt: str = "",
@@ -139,39 +138,40 @@ def generate(
     logger.info(json.dumps(metadata, indent=4))
 
     try:
-        if use_upscaler:
-            latents = pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                width=width,
-                height=height,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                generator=generator,
-                output_type="latent",
-            ).images
-            upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
-            images = upscaler_pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                image=upscaled_latents,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                strength=upscaler_strength,
-                generator=generator,
-                output_type="pil",
-            ).images
-        else:
-            images = pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                width=width,
-                height=height,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                generator=generator,
-                output_type="pil",
-            ).images
+        with torch.no_grad():
+            if use_upscaler:
+                latents = pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    width=width,
+                    height=height,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    generator=generator,
+                    output_type="latent",
+                ).images
+                upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
+                images = upscaler_pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    image=upscaled_latents,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    strength=upscaler_strength,
+                    generator=generator,
+                    output_type="pil",
+                ).images
+            else:
+                images = pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    width=width,
+                    height=height,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    generator=generator,
+                    output_type="pil",
+                ).images
 
         if images:
             image_paths = [
@@ -193,11 +193,13 @@ def generate(
         utils.free_memory()
 
 
+
 if torch.cuda.is_available():
     pipe = load_pipeline(MODEL)
     logger.info("Loaded on Device!")
 else:
-    pipe = None
+    pipe = load_pipeline(MODEL)
+    logger.info("Loaded on CPU!")
 
 styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in config.style_list}
 quality_prompt = {
@@ -206,11 +208,11 @@ quality_prompt = {
 
 wildcard_files = utils.load_wildcard_files("wildcard")
 
-with gr.Blocks(css="style.css") as demo:
-#    title = gr.HTML(
-#        f"""<h1><span>{DESCRIPTION}</span></h1>""",
-#        elem_id="title",
-#    )
+with gr.Blocks(css="style.css", theme="NoCrypt/[email protected]") as demo:
+    title = gr.HTML(
+        f"""<h1><span>{DESCRIPTION}</span></h1>""",
+        elem_id="title",
+    )
     gr.Markdown(
         f"""Gradio demo for [cagliostrolab/animagine-xl-3.1](https://huggingface.co/cagliostrolab/animagine-xl-3.1)""",
         elem_id="subtitle",
@@ -392,4 +394,4 @@ with gr.Blocks(css="style.css") as demo:
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)
+    demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)