Update app.py

app.py

@@ -7,9 +7,6 @@ from PIL import Image
 
 from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel, AutoencoderKL, UniPCMultistepScheduler
 
-# -----------------------------------------------------------------------------
-# Configuration & Registry
-# -----------------------------------------------------------------------------
 LORA_REGISTRY = {
     "None (Base SDXL)": {
         "repo": None,
@@ -36,19 +33,18 @@ LORA_REGISTRY = {
     }
 }
 
-# -----------------------------------------------------------------------------
-# Model Initialization
-# -----------------------------------------------------------------------------
-print("Initializing SDXL Pipeline on CPU...")
+print("Loading SDXL Pipeline...")
+
+dtype = torch.float16
 
 vae = AutoencoderKL.from_pretrained(
     "madebyollin/sdxl-vae-fp16-fix", 
-    torch_dtype=torch.float16
+    torch_dtype=dtype
 )
 
 controlnet = ControlNetModel.from_pretrained(
     "diffusers/controlnet-canny-sdxl-1.0",
-    torch_dtype=torch.float16,
+    torch_dtype=dtype,
     use_safetensors=True
 )
 
@@ -56,17 +52,15 @@ pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
     controlnet=controlnet,
     vae=vae,
-    torch_dtype=torch.float16,
+    torch_dtype=dtype,
     use_safetensors=True
 )
 
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 
-print("Pipeline loaded successfully.")
+pipe.enable_model_cpu_offload()
 
-# -----------------------------------------------------------------------------
-# Helper Functions
-# -----------------------------------------------------------------------------
+print("Pipeline loaded successfully.")
 
 def get_canny_image(image, low_threshold=100, high_threshold=200):
     image_array = np.array(image)
@@ -75,10 +69,6 @@ def get_canny_image(image, low_threshold=100, high_threshold=200):
     canny_edges = np.concatenate([canny_edges, canny_edges, canny_edges], axis=2)
     return Image.fromarray(canny_edges)
 
-# -----------------------------------------------------------------------------
-# Inference Logic
-# -----------------------------------------------------------------------------
-
 @spaces.GPU(duration=120)
 def generate_controlled_image(
     input_image,
@@ -91,10 +81,6 @@ def generate_controlled_image(
 ):
     if input_image is None:
         raise gr.Error("Please upload an image first!")
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    
-    pipe.to(device)
     
     width, height = 1024, 1024
     input_image = input_image.resize((width, height))
@@ -121,7 +107,7 @@ def generate_controlled_image(
             print(f"LoRA Load Error: {e}")
             gr.Warning(f"Failed to load LoRA. Using base model.")
 
-    generator = torch.Generator(device).manual_seed(int(seed))
+    generator = torch.Generator("cuda").manual_seed(int(seed))
 
     print(f"Generating: {final_prompt[:100]}...")
     
@@ -141,16 +127,10 @@ def generate_controlled_image(
         raise e
 
     pipe.unload_lora_weights()
-    
-    if device == "cuda":
-        torch.cuda.empty_cache()
+    torch.cuda.empty_cache()
 
     return canny_image, output_image
 
-# -----------------------------------------------------------------------------
-# Gradio UI
-# -----------------------------------------------------------------------------
-
 css = """
 #col-container {max-width: 1200px; margin-left: auto; margin-right: auto;}
 .guide-text {font-size: 1.1em; color: #4a5568;}