Update app.py

app.py

@@ -17,9 +17,11 @@ client = Groq(api_key=api_key)
 
 # Set device: CUDA if available, else CPU
 # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
 
 model_id1 = os.getenv("API_KEY")
-pipe = StableDiffusionPipeline.from_pretrained(model_id1, torch_dtype=torch.float16, use_safetensors=True)
+pipe = StableDiffusionPipeline.from_pretrained(model_id1, torch_dtype=torch_dtype, use_safetensors=True)
 pipe = pipe.to('cpu')
 
 # Updated function for text generation using the new API structure