qqwjq1981 commited on
Commit
ddb961b
·
verified ·
1 Parent(s): 4b5729a

Update utils/keyframe_utils.py

Browse files
Files changed (1) hide show
  1. utils/keyframe_utils.py +2 -2
utils/keyframe_utils.py CHANGED
@@ -16,7 +16,7 @@ CACHE_DIR.mkdir(exist_ok=True)
16
  LOG_PATH = Path("prompt_log.jsonl")
17
 
18
  pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
19
- pipe = pipe.to("cuda")
20
 
21
  def generate_keyframe_prompt(segment):
22
  segment_id = segment.get("segment_id")
@@ -91,4 +91,4 @@ def generate_all_keyframe_images(script_data, output_dir="keyframes"):
91
  with open("all_prompts_output.json", "w", encoding="utf-8") as f:
92
  json.dump(keyframe_outputs, f, ensure_ascii=False, indent=2)
93
 
94
- return keyframe_outputs
 
16
  LOG_PATH = Path("prompt_log.jsonl")
17
 
18
  pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
19
+ pipe = pipe.to("cpu")
20
 
21
  def generate_keyframe_prompt(segment):
22
  segment_id = segment.get("segment_id")
 
91
  with open("all_prompts_output.json", "w", encoding="utf-8") as f:
92
  json.dump(keyframe_outputs, f, ensure_ascii=False, indent=2)
93
 
94
+ return keyframe_outputs