Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import torch
 import numpy as np
 import cv2
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionPipeline, UniPCMultistepScheduler
 from model import UNet2DConditionModelEx
 from pipeline import StableDiffusionControlLoraV3Pipeline 
 from PIL import Image
@@ -14,7 +14,7 @@ login(token=os.environ.get("HF_TOKEN"))
 
 # Initialize the models
 base_model = "runwayml/stable-diffusion-v1-5"
-dtype = torch.float32
+dtype = torch.float16  # A100 works better with float16
 
 # Load the custom UNet
 unet = UNet2DConditionModelEx.from_pretrained(
@@ -33,16 +33,17 @@ pipe = StableDiffusionControlLoraV3Pipeline.from_pretrained(
     torch_dtype=dtype
 )
 
+# Use a faster scheduler
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+
 # Load the ControlLoRA weights
 pipe.load_lora_weights(
     "models",
     weight_name="40kHalf.safetensors"
 )
 
-# Enable CPU offload
+# Enable optimizations
 pipe.enable_model_cpu_offload()
-
-# Enable memory efficient attention if available
 if hasattr(pipe, "enable_xformers_memory_efficient_attention"):
     pipe.enable_xformers_memory_efficient_attention()