cakemus commited on
Commit
71bbdfb
·
1 Parent(s): ef2375d

disk space fix test

Browse files
Files changed (1) hide show
  1. app.py +67 -30
app.py CHANGED
@@ -3,6 +3,7 @@ import spaces
3
  #import gradio.helpers
4
  import torch
5
  import os
 
6
  from glob import glob
7
  from pathlib import Path
8
  from typing import Optional
@@ -17,19 +18,38 @@ from huggingface_hub import hf_hub_download
17
 
18
  #gradio.helpers.CACHED_FOLDER = '/data/cache'
19
 
 
 
 
 
 
 
 
 
20
  # Load the pipeline with authentication token
21
  pipe = StableVideoDiffusionPipeline.from_pretrained(
22
  "stabilityai/stable-video-diffusion-img2vid-xt",
23
  torch_dtype=torch.float16,
24
  variant="fp16",
25
- use_auth_token=os.getenv("HUGGINGFACE_TOKEN") # Fetch the token from the environment variable
26
  )
27
  pipe.to("cuda")
28
- #pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
29
- #pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
 
30
 
31
  max_64_bit_int = 2**63 - 1
32
 
 
 
 
 
 
 
 
 
 
 
33
  @spaces.GPU(duration=250)
34
  def sample(
35
  image: Image,
@@ -55,62 +75,79 @@ def sample(
55
  base_count = len(glob(os.path.join(output_folder, "*.mp4")))
56
  video_path = os.path.join(output_folder, f"{base_count:06d}.mp4")
57
 
58
- frames = pipe(image, decode_chunk_size=decoding_t, generator=generator, motion_bucket_id=motion_bucket_id, noise_aug_strength=0.1, num_frames=25).frames[0]
 
 
 
 
 
 
 
 
 
59
  export_to_video(frames, video_path, fps=fps_id)
60
  torch.manual_seed(seed)
 
 
 
61
 
62
  return video_path, seed
63
 
64
  def resize_image(image, output_size=(1024, 576)):
65
  # Calculate aspect ratios
66
- target_aspect = output_size[0] / output_size[1] # Aspect ratio of the desired size
67
- image_aspect = image.width / image.height # Aspect ratio of the original image
68
 
69
  # Resize then crop if the original image is larger
70
  if image_aspect > target_aspect:
71
- # Resize the image to match the target height, maintaining aspect ratio
72
  new_height = output_size[1]
73
  new_width = int(new_height * image_aspect)
74
  resized_image = image.resize((new_width, new_height), Image.LANCZOS)
75
- # Calculate coordinates for cropping
76
  left = (new_width - output_size[0]) / 2
77
  top = 0
78
  right = (new_width + output_size[0]) / 2
79
  bottom = output_size[1]
80
  else:
81
- # Resize the image to match the target width, maintaining aspect ratio
82
  new_width = output_size[0]
83
  new_height = int(new_width / image_aspect)
84
  resized_image = image.resize((new_width, new_height), Image.LANCZOS)
85
- # Calculate coordinates for cropping
86
  left = 0
87
  top = (new_height - output_size[1]) / 2
88
  right = output_size[0]
89
  bottom = (new_height + output_size[1]) / 2
90
 
91
- # Crop the image
92
  cropped_image = resized_image.crop((left, top, right, bottom))
93
  return cropped_image
94
 
95
  with gr.Blocks() as demo:
96
- gr.Markdown('''# Community demo for Stable Video Diffusion - Img2Vid - XT ([model](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt), [paper](https://stability.ai/research/stable-video-diffusion-scaling-latent-video-diffusion-models-to-large-datasets), [stability's ui waitlist](https://stability.ai/contact))
97
- #### Research release ([_non-commercial_](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt/blob/main/LICENSE)): generate `4s` vid from a single image at (`25 frames` at `6 fps`). this demo uses [🧨 diffusers for low VRAM and fast generation](https://huggingface.co/docs/diffusers/main/en/using-diffusers/svd).
98
- ''')
99
- with gr.Row():
100
- with gr.Column():
101
- image = gr.Image(label="Upload your image", type="pil")
102
- generate_btn = gr.Button("Generate")
103
- video = gr.Video()
104
- with gr.Accordion("Advanced options", open=False):
105
- seed = gr.Slider(label="Seed", value=42, randomize=True, minimum=0, maximum=max_64_bit_int, step=1)
106
- randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
107
- motion_bucket_id = gr.Slider(label="Motion bucket id", info="Controls how much motion to add/remove from the image", value=127, minimum=1, maximum=255)
108
- fps_id = gr.Slider(label="Frames per second", info="The length of your video in seconds will be 25/fps", value=6, minimum=5, maximum=30)
109
-
110
- image.upload(fn=resize_image, inputs=image, outputs=image, queue=False)
111
- generate_btn.click(fn=sample, inputs=[image, seed, randomize_seed, motion_bucket_id, fps_id], outputs=[video, seed], api_name="video")
112
-
 
 
 
 
 
 
 
 
 
113
 
114
  if __name__ == "__main__":
115
- #demo.queue(max_size=20, api_open=False)
116
- demo.launch(share=True, show_api=False)
 
 
3
  #import gradio.helpers
4
  import torch
5
  import os
6
+ import shutil
7
  from glob import glob
8
  from pathlib import Path
9
  from typing import Optional
 
18
 
19
  #gradio.helpers.CACHED_FOLDER = '/data/cache'
20
 
21
+ # OPTIONAL: Clear caches at startup to free space (comment out if you prefer not to)
22
+ hf_cache = os.path.expanduser("~/.cache/huggingface")
23
+ torch_cache = os.path.expanduser("~/.cache/torch")
24
+ if os.path.exists(hf_cache):
25
+ shutil.rmtree(hf_cache)
26
+ if os.path.exists(torch_cache):
27
+ shutil.rmtree(torch_cache)
28
+
29
  # Load the pipeline with authentication token
30
  pipe = StableVideoDiffusionPipeline.from_pretrained(
31
  "stabilityai/stable-video-diffusion-img2vid-xt",
32
  torch_dtype=torch.float16,
33
  variant="fp16",
34
+ use_auth_token=os.getenv("HUGGINGFACE_TOKEN") # Fetch the token from environment if set
35
  )
36
  pipe.to("cuda")
37
+ # Uncomment these lines only if you want experimental compilation (and if your environment supports it)
38
+ # pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
39
+ # pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
40
 
41
  max_64_bit_int = 2**63 - 1
42
 
43
+ def clean_outputs(output_folder: str, keep: int = 1):
44
+ """
45
+ Remove old video files to prevent using all disk space.
46
+ Keeps the most recent <keep> files.
47
+ """
48
+ files = sorted(glob(os.path.join(output_folder, "*.mp4")), key=os.path.getmtime)
49
+ if len(files) > keep:
50
+ for old_file in files[:-keep]:
51
+ os.remove(old_file)
52
+
53
  @spaces.GPU(duration=250)
54
  def sample(
55
  image: Image,
 
75
  base_count = len(glob(os.path.join(output_folder, "*.mp4")))
76
  video_path = os.path.join(output_folder, f"{base_count:06d}.mp4")
77
 
78
+ # Reduce num_frames from 25 to 10 to consume less space
79
+ frames = pipe(
80
+ image,
81
+ decode_chunk_size=decoding_t,
82
+ generator=generator,
83
+ motion_bucket_id=motion_bucket_id,
84
+ noise_aug_strength=0.1,
85
+ num_frames=10 # reduced from 25
86
+ ).frames[0]
87
+
88
  export_to_video(frames, video_path, fps=fps_id)
89
  torch.manual_seed(seed)
90
+
91
+ # Clean up old videos to prevent filling disk
92
+ clean_outputs(output_folder, keep=2)
93
 
94
  return video_path, seed
95
 
96
  def resize_image(image, output_size=(1024, 576)):
97
  # Calculate aspect ratios
98
+ target_aspect = output_size[0] / output_size[1]
99
+ image_aspect = image.width / image.height
100
 
101
  # Resize then crop if the original image is larger
102
  if image_aspect > target_aspect:
 
103
  new_height = output_size[1]
104
  new_width = int(new_height * image_aspect)
105
  resized_image = image.resize((new_width, new_height), Image.LANCZOS)
 
106
  left = (new_width - output_size[0]) / 2
107
  top = 0
108
  right = (new_width + output_size[0]) / 2
109
  bottom = output_size[1]
110
  else:
 
111
  new_width = output_size[0]
112
  new_height = int(new_width / image_aspect)
113
  resized_image = image.resize((new_width, new_height), Image.LANCZOS)
 
114
  left = 0
115
  top = (new_height - output_size[1]) / 2
116
  right = output_size[0]
117
  bottom = (new_height + output_size[1]) / 2
118
 
 
119
  cropped_image = resized_image.crop((left, top, right, bottom))
120
  return cropped_image
121
 
122
  with gr.Blocks() as demo:
123
+ gr.Markdown('''# Community demo for Stable Video Diffusion - Img2Vid - XT ([model](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt), [paper](https://stability.ai/research/stable-video-diffusion-scaling-latent-video-diffusion-models-to-large-datasets), [stability's ui waitlist](https://stability.ai/contact))
124
+ #### Research release ([_non-commercial_](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt/blob/main/LICENSE)): generate `~4s` vid from a single image at (`10 frames` at `6 fps`). This demo uses [🧨 diffusers](https://huggingface.co/docs/diffusers/main/en/using-diffusers/svd) for low VRAM usage.
125
+ ''')
126
+ with gr.Row():
127
+ with gr.Column():
128
+ image = gr.Image(label="Upload your image", type="pil")
129
+ generate_btn = gr.Button("Generate")
130
+ video = gr.Video()
131
+ with gr.Accordion("Advanced options", open=False):
132
+ seed = gr.Slider(label="Seed", value=42, randomize=True, minimum=0, maximum=max_64_bit_int, step=1)
133
+ randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
134
+ motion_bucket_id = gr.Slider(label="Motion bucket id", info="Controls how much motion to add/remove from the image", value=127, minimum=1, maximum=255)
135
+ fps_id = gr.Slider(label="Frames per second", info="The length of your video in seconds will be num_frames/fps", value=6, minimum=5, maximum=30)
136
+
137
+ # Resize on upload
138
+ image.upload(fn=resize_image, inputs=image, outputs=image, queue=False)
139
+
140
+ # Generate with sample() function
141
+ generate_btn.click(
142
+ fn=sample,
143
+ inputs=[image, seed, randomize_seed, motion_bucket_id, fps_id],
144
+ outputs=[video, seed],
145
+ api_name="video"
146
+ )
147
+
148
+ # REMOVED: gr.Examples(...) that referenced local images like "blink_meme.png" which don't exist
149
 
150
  if __name__ == "__main__":
151
+ # remove `share=True` because it is not supported on Hugging Face Spaces
152
+ demo.launch(show_api=False)
153
+