Update app.py

app.py

@@ -19,9 +19,13 @@ if token:
 else:
     raise ValueError("Hugging Face token not found. Please set it as a repository secret in the Space settings.")
 
-# Load the Stable Diffusion 3.5 model with lower precision (float16)
+# Load the Stable Diffusion 3.5 model with lower precision (float16) if GPU is available
 model_id = "stabilityai/stable-diffusion-3.5-large"
-pipe = StableDiffusion3Pipeline.from_pretrained(model_id, torch_dtype=torch.float16)  # Use float16 precision
+if device == "cuda":
+    pipe = StableDiffusion3Pipeline.from_pretrained(model_id, torch_dtype=torch.float16)  # Use float16 precision
+else:
+    pipe = StableDiffusion3Pipeline.from_pretrained(model_id)  # Default precision for CPU
+
 pipe.to(device)  # Ensuring the model is on the correct device (GPU or CPU)
 
 # Define the path to the LoRA model
@@ -32,10 +36,17 @@ def load_lora_model(pipe, lora_model_path):
     # Load the LoRA weights
     lora_weights = torch.load(lora_model_path, map_location=device)  # Load LoRA model to the correct device
     
+    # Print available attributes of the model to check access to `unet` (optional)
+    print(dir(pipe))  # This will list all attributes and methods of the `pipe` object
+
     # Apply weights to the UNet submodule
-    for name, param in pipe.unet.named_parameters():  # Accessing unet parameters
-        if name in lora_weights:
-            param.data += lora_weights[name]
+    try:
+        for name, param in pipe.unet.named_parameters():  # Accessing unet parameters
+            if name in lora_weights:
+                param.data += lora_weights[name]
+    except AttributeError:
+        print("The model doesn't have 'unet' attributes. Please check the model structure.")
+        # Add alternative handling or exit
 
     return pipe