improve

run_local.py

@@ -4,21 +4,24 @@ import time
 from huggingface_hub import HfApi
 import torch
 import sys
+from pathlib import Path
 
 path = sys.argv[1]
 
 api = HfApi()
 start_time = time.time()
-pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
+pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16, device_map="auto")
 pipe.scheduler = HeunDiscreteScheduler.from_config(pipe.scheduler.config)
-pipe = pipe.to("cuda")
+#pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
+#pipe.scheduler = HeunDiscreteScheduler.from_config(pipe.scheduler.config)
+#pipe = pipe.to("cuda")
 
 prompt = "a highly realistic photo of green turtle"
 generator = torch.Generator(device="cuda").manual_seed(0)
 image = pipe(prompt, generator=generator, num_inference_steps=25).images[0]
 print("Time", time.time() - start_time)
 
-path = "/home/patrick_huggingface_co/images/aa.png"
+path = os.path.join(Path.home(), "images", "aa.png")
 image.save(path)
 
 api.upload_file(