Update app.py

app.py

@@ -20,16 +20,20 @@ class StableBuddyApp:
         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
+        # Use float16 and fp16 so that stable diffusion can work on 4GB VRAM   float 32 for cpu
         self.pipe = StableDiffusionPipeline.from_pretrained(
-            model_id, revision='fp16', torch_dtype=torch.float16, use_auth_token=auth_token
+            model_id, revision='fp16', torch_dtype=torch.float32, use_auth_token=auth_token
         )
         self.pipe.to(device)
 
     def generate_image(self, prompt):
         """Generate an image based on the prompt."""
         try:
-            with autocast(device):
+            # Use autocast only for GPU
+            if self.device == "cuda":
+                with autocast(self.device):
+                    image = self.pipe(prompt, guidance_scale=8.5).images[0]
+            else:
                 image = self.pipe(prompt, guidance_scale=8.5).images[0]
             
             # Save the generated image temporarily
@@ -40,8 +44,7 @@ class StableBuddyApp:
             
         except Exception as e:
             print(f"An error occurred: {e}")
-            return None  # In case of an error, return None
-
+            return None  
 # Create an instance of the StableBuddyApp
 stable_buddy_app = StableBuddyApp()