update device

app.py

@@ -8,6 +8,7 @@ import os
 
 # Global variable to store the text pipeline
 text_pipe = None
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 def load_model():
     """Load the T5 text encoder model"""
@@ -24,23 +25,14 @@ def load_model():
             load_in_8bit=True,
             variant="8bit",
             token=token
-        )
+        ).to(device)
         text_pipe = DiffusionPipeline.from_pretrained(
             "DeepFloyd/IF-I-L-v1.0",
             text_encoder=text_encoder,
             unet=None,
             token=token,
         )
-        
-        # Set the execution device to match the text encoder's device
-        # This ensures tokenizer outputs go to the correct device
-        if hasattr(text_encoder, 'device'):
-            text_pipe._execution_device = text_encoder.device
-        else:
-            # Find the device of the first parameter
-            text_pipe._execution_device = next(text_encoder.parameters()).device
-        
-        print(f"Model loaded successfully on device: {text_pipe._execution_device}")
+        print("Model loaded successfully!")
     return text_pipe
 
 @spaces.GPU