Martim-Ramos-Neural commited on
Commit
4372b17
·
verified ·
1 Parent(s): 08c3e82

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +149 -0
app.py ADDED
@@ -0,0 +1,149 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ import os
4
+ import spaces
5
+ import uuid
6
+
7
+ from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
8
+ from diffusers.utils import export_to_video
9
+ from huggingface_hub import hf_hub_download
10
+ from safetensors.torch import load_file
11
+ from PIL import Image
12
+
13
+ # Constants
14
+ bases = {
15
+ "Cartoon": "frankjoshua/toonyou_beta6",
16
+ "Realistic": "emilianJR/epiCRealism",
17
+ "3d": "Lykon/DreamShaper",
18
+ "Anime": "Yntec/mistoonAnime2"
19
+ }
20
+ step_loaded = None
21
+ base_loaded = "Realistic"
22
+ motion_loaded = None
23
+
24
+ # Ensure model and scheduler are initialized in GPU-enabled function
25
+ if not torch.cuda.is_available():
26
+ raise NotImplementedError("No GPU detected!")
27
+
28
+ device = "cuda"
29
+ dtype = torch.float16
30
+ pipe = AnimateDiffPipeline.from_pretrained(bases[base_loaded], torch_dtype=dtype).to(device)
31
+ pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")
32
+
33
+ # Safety checkers
34
+ from transformers import CLIPFeatureExtractor
35
+
36
+ feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")
37
+
38
+ # Function
39
+ @spaces.GPU(duration=30,queue=False)
40
+ def generate_image(prompt, base="Realistic", motion="", step=8, progress=gr.Progress()):
41
+ global step_loaded
42
+ global base_loaded
43
+ global motion_loaded
44
+ print(prompt, base, step)
45
+
46
+ if step_loaded != step:
47
+ repo = "ByteDance/AnimateDiff-Lightning"
48
+ ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
49
+ pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device), strict=False)
50
+ step_loaded = step
51
+
52
+ if base_loaded != base:
53
+ pipe.unet.load_state_dict(torch.load(hf_hub_download(bases[base], "unet/diffusion_pytorch_model.bin"), map_location=device), strict=False)
54
+ base_loaded = base
55
+
56
+ if motion_loaded != motion:
57
+ pipe.unload_lora_weights()
58
+ if motion != "":
59
+ pipe.load_lora_weights(motion, adapter_name="motion")
60
+ pipe.set_adapters(["motion"], [0.7])
61
+ motion_loaded = motion
62
+
63
+ progress((0, step))
64
+ def progress_callback(i, t, z):
65
+ progress((i+1, step))
66
+
67
+ output = pipe(prompt=prompt, guidance_scale=1.2, num_inference_steps=step, callback=progress_callback, callback_steps=1)
68
+
69
+ name = str(uuid.uuid4()).replace("-", "")
70
+ path = f"/tmp/{name}.mp4"
71
+ export_to_video(output.frames[0], path, fps=10)
72
+ return path
73
+
74
+
75
+ # Gradio Interface
76
+ with gr.Blocks(css="style.css") as demo:
77
+ gr.HTML(
78
+ "<h1><center>Textual Imagination : A Text To Video Synthesis</center></h1>"
79
+ )
80
+ with gr.Group():
81
+ with gr.Row():
82
+ prompt = gr.Textbox(
83
+ label='Prompt'
84
+ )
85
+ with gr.Row():
86
+ select_base = gr.Dropdown(
87
+ label='Base model',
88
+ choices=[
89
+ "Cartoon",
90
+ "Realistic",
91
+ "3d",
92
+ "Anime",
93
+ ],
94
+ value=base_loaded,
95
+ interactive=True
96
+ )
97
+ select_motion = gr.Dropdown(
98
+ label='Motion',
99
+ choices=[
100
+ ("Default", ""),
101
+ ("Zoom in", "guoyww/animatediff-motion-lora-zoom-in"),
102
+ ("Zoom out", "guoyww/animatediff-motion-lora-zoom-out"),
103
+ ("Tilt up", "guoyww/animatediff-motion-lora-tilt-up"),
104
+ ("Tilt down", "guoyww/animatediff-motion-lora-tilt-down"),
105
+ ("Pan left", "guoyww/animatediff-motion-lora-pan-left"),
106
+ ("Pan right", "guoyww/animatediff-motion-lora-pan-right"),
107
+ ("Roll left", "guoyww/animatediff-motion-lora-rolling-anticlockwise"),
108
+ ("Roll right", "guoyww/animatediff-motion-lora-rolling-clockwise"),
109
+ ],
110
+ value="guoyww/animatediff-motion-lora-zoom-in",
111
+ interactive=True
112
+ )
113
+ select_step = gr.Dropdown(
114
+ label='Inference steps',
115
+ choices=[
116
+ ('1-Step', 1),
117
+ ('2-Step', 2),
118
+ ('4-Step', 4),
119
+ ('8-Step', 8),
120
+ ],
121
+ value=4,
122
+ interactive=True
123
+ )
124
+ submit = gr.Button(
125
+ scale=1,
126
+ variant='primary'
127
+ )
128
+ video = gr.Video(
129
+ label='AnimateDiff-Lightning',
130
+ autoplay=True,
131
+ height=512,
132
+ width=512,
133
+ elem_id="video_output"
134
+ )
135
+
136
+ gr.on(triggers=[
137
+ submit.click,
138
+ prompt.submit
139
+ ],
140
+ fn = generate_image,
141
+ inputs = [prompt, select_base, select_motion, select_step],
142
+ outputs = [video],
143
+ api_name = "instant_video",
144
+ queue = False
145
+ )
146
+
147
+ demo.queue().launch()
148
+
149
+ Translate