Update app.py

app.py

@@ -13,18 +13,19 @@ class StableBuddyApp:
     def __init__(self):
         # Set up the Stable Diffusion pipeline
         model_id = "CompVis/stable-diffusion-v1-4"
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"  # Store device as a class attribute
         
         # Get the auth_token from the environment variable
         auth_token = os.getenv("AUTH_TOKEN")
         if not auth_token:
             raise ValueError("AUTH_TOKEN environment variable is not set.")
         
-        # Use float16 and fp16 so that stable diffusion can work on 4GB VRAM   float 32 for cpu
+        # Use float16 for GPU and float32 for CPU to manage VRAM
+        torch_dtype = torch.float16 if self.device == "cuda" else torch.float32
         self.pipe = StableDiffusionPipeline.from_pretrained(
-            model_id, revision='fp16', torch_dtype=torch.float32, use_auth_token=auth_token
+            model_id, revision='fp16', torch_dtype=torch_dtype, use_auth_token=auth_token
         )
-        self.pipe.to(device)
+        self.pipe.to(self.device)
 
     def generate_image(self, prompt):
         """Generate an image based on the prompt."""
@@ -44,7 +45,8 @@ class StableBuddyApp:
             
         except Exception as e:
             print(f"An error occurred: {e}")
-            return None  
+            return None
+
 # Create an instance of the StableBuddyApp
 stable_buddy_app = StableBuddyApp()