Update app.py

app.py

@@ -10,10 +10,10 @@ HF_TOKEN = os.getenv('HF_TOKEN')
 hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "movie-img-generator")
 auth_token = os.environ.get("auth_token")
 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=auth_token)
-pipe = pipe.to("cuda")
+pipe = pipe.to(device)
 
-device = "cuda"
 generator = torch.Generator(device=device)
 seed = generator.seed()