Update app.py

app.py

@@ -17,29 +17,27 @@ model_id = "stabilityai/stable-diffusion-3.5-large"
 pipe = StableDiffusion3Pipeline.from_pretrained(model_id)
 
 # Check if GPU is available, then move the model to the appropriate device
-pipe.to('cuda' if torch.cuda.is_available() else 'cpu')
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+pipe.to(device)
 
 # Define the path to the LoRA model
 lora_model_path = "./lora_model.pth"  # Assuming the file is saved locally
 
 # Custom method to load and apply LoRA weights to the Stable Diffusion pipeline
 def load_lora_model(pipe, lora_model_path):
-    # Set device to 'cuda' if available, otherwise 'cpu'
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
     # When loading the LoRA weights
     lora_weights = torch.load(lora_model_path, map_location=device, weights_only=True)
+    
+    # Check if the transformer folder has the necessary attributes
+    print(dir(pipe.transformer))  # List available attributes of the transformer (formerly 'unet')
 
-    # 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
+    # Apply weights to the transformer submodule
     try:
-        for name, param in pipe.unet.named_parameters():  # Accessing unet parameters
+        for name, param in pipe.transformer.named_parameters():  # Accessing transformer 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.")
+        print("The model doesn't have 'transformer' attributes. Please check the model structure.")
         # Add alternative handling or exit
 
     return pipe
@@ -47,7 +45,7 @@ def load_lora_model(pipe, lora_model_path):
 # Load and apply the LoRA model weights
 pipe = load_lora_model(pipe, lora_model_path)
 
-# Use the @space.gpu decorator to ensure compatibility with GPU or CPU as needed
+# Use the @spaces.gpu decorator to ensure compatibility with GPU or CPU as needed
 @spaces.gpu
 def generate(prompt, seed=None):
     generator = torch.manual_seed(seed) if seed is not None else None