Update app.py

app.py

@@ -6,12 +6,11 @@ import gradio as gr
 import numpy as np
 import torch
 from PIL import Image
-from diffusers import StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, AutoencoderKL
+from diffusers import StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, AutoencoderKL, EulerAncestralDiscreteScheduler
 from custom_pipeline import CosStableDiffusionXLInstructPix2PixPipeline
 from huggingface_hub import hf_hub_download
 from huggingface_hub import InferenceClient
 
-
 help_text = """
 To optimize image results:
 - Adjust the **Image CFG weight** if the image isn't changing enough or is changing too much. Lower it to allow bigger changes, or raise it to preserve original details.
@@ -41,6 +40,7 @@ normal_file = hf_hub_download(repo_id="stabilityai/cosxl", filename="cosxl.safet
 
 EDMEulerScheduler.set_timesteps = set_timesteps_patched
 
+# Use a pre-converted fp16 VAE for faster loading and inference
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 
 pipe_edit = StableDiffusionXLInstructPix2PixPipeline.from_single_file(
@@ -49,21 +49,15 @@ pipe_edit = StableDiffusionXLInstructPix2PixPipeline.from_single_file(
 pipe_edit.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction")
 pipe_edit.to("cuda")
 
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
-
-if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-
 # Image Generator
+# Keep the models loaded globally for reuse
 if torch.cuda.is_available():
     pipe = StableDiffusionXLPipeline.from_pretrained(
         "fluently/Fluently-XL-v4",
         torch_dtype=torch.float16,
         use_safetensors=True,
-    )
+    ).to("cuda")
+    pipe.enable_xformers_memory_efficient_attention() # Enable memory efficient attention for less VRAM usage
     pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
     pipe.load_lora_weights("ehristoforu/dalle-3-xl-v2", weight_name="dalle-3-xl-lora-v2.safetensors", adapter_name="dalle")
     pipe.set_adapters("dalle")
@@ -103,10 +97,9 @@ def king(type = "Image Generation",
             num_inference_steps=steps, generator=generator).images[0]
         return seed, output_image
     else :
-        pipe.to(device)
         seed = int(randomize_seed_fn(seed, randomize_seed))
-        generator = torch.Generator().manual_seed(seed)   
-        
+        generator = torch.Generator(device="cuda").manual_seed(seed) # Move generator to cuda for speed
+
         options = {
             "prompt":instruction,
             "width":width,
@@ -118,14 +111,14 @@ def king(type = "Image Generation",
             "output_type":"pil",
         }
         
-        output_image = pipe(**options).images[0]
+        with torch.autocast("cuda"): # Use autocast for faster inference
+            output_image = pipe(**options).images[0]
         return seed, output_image
 
 # Prompt classifier
 def response(instruction, input_image=None):
     if input_image is None:
         output="Image Generation"
-        yield output
     else:
         client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")