Harumiiii commited on
Commit
ebb1079
·
verified ·
1 Parent(s): 0f71ed2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +43 -42
app.py CHANGED
@@ -1,66 +1,75 @@
1
  import gradio as gr
2
  import numpy as np
3
  import random
4
- #import spaces #[uncomment to use ZeroGPU]
5
  from diffusers import DiffusionPipeline
6
  import torch
7
 
 
8
  device = "cuda" if torch.cuda.is_available() else "cpu"
9
- model_repo_id = "stabilityai/sdxl-turbo" #Replace to the model you would like to use
10
 
 
 
 
 
11
  if torch.cuda.is_available():
12
  torch_dtype = torch.float16
13
  else:
14
  torch_dtype = torch.float32
15
 
 
16
  pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
17
  pipe = pipe.to(device)
18
 
 
19
  MAX_SEED = np.iinfo(np.int32).max
20
  MAX_IMAGE_SIZE = 1024
21
 
22
- #@spaces.GPU #[uncomment to use ZeroGPU]
23
  def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
24
 
 
25
  if randomize_seed:
26
  seed = random.randint(0, MAX_SEED)
27
 
28
- generator = torch.Generator().manual_seed(seed)
29
-
30
- image = pipe(
31
- prompt = prompt,
32
- negative_prompt = negative_prompt,
33
- guidance_scale = guidance_scale,
34
- num_inference_steps = num_inference_steps,
35
- width = width,
36
- height = height,
37
- generator = generator
38
- ).images[0]
 
39
 
40
- return image, seed
41
 
 
42
  examples = [
43
- "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
44
- "An astronaut riding a green horse",
45
- "A delicious ceviche cheesecake slice",
46
  ]
47
 
48
- css="""
 
49
  #col-container {
50
  margin: 0 auto;
51
  max-width: 640px;
52
  }
53
  """
54
 
 
55
  with gr.Blocks(css=css) as demo:
56
 
57
  with gr.Column(elem_id="col-container"):
58
  gr.Markdown(f"""
59
- # Text-to-Image Gradio Template
60
  """)
61
 
62
  with gr.Row():
63
-
64
  prompt = gr.Text(
65
  label="Prompt",
66
  show_label=False,
@@ -68,10 +77,9 @@ with gr.Blocks(css=css) as demo:
68
  placeholder="Enter your prompt",
69
  container=False,
70
  )
71
-
72
  run_button = gr.Button("Run", scale=0)
73
 
74
- result = gr.Image(label="Result", show_label=False)
75
 
76
  with gr.Accordion("Advanced Settings", open=False):
77
 
@@ -79,7 +87,7 @@ with gr.Blocks(css=css) as demo:
79
  label="Negative prompt",
80
  max_lines=1,
81
  placeholder="Enter a negative prompt",
82
- visible=False,
83
  )
84
 
85
  seed = gr.Slider(
@@ -93,50 +101,43 @@ with gr.Blocks(css=css) as demo:
93
  randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
94
 
95
  with gr.Row():
96
-
97
  width = gr.Slider(
98
  label="Width",
99
  minimum=256,
100
  maximum=MAX_IMAGE_SIZE,
101
  step=32,
102
- value=1024, #Replace with defaults that work for your model
103
  )
104
-
105
  height = gr.Slider(
106
  label="Height",
107
  minimum=256,
108
  maximum=MAX_IMAGE_SIZE,
109
  step=32,
110
- value=1024, #Replace with defaults that work for your model
111
  )
112
 
113
  with gr.Row():
114
-
115
  guidance_scale = gr.Slider(
116
  label="Guidance scale",
117
  minimum=0.0,
118
  maximum=10.0,
119
  step=0.1,
120
- value=0.0, #Replace with defaults that work for your model
121
  )
122
-
123
  num_inference_steps = gr.Slider(
124
  label="Number of inference steps",
125
  minimum=1,
126
  maximum=50,
127
  step=1,
128
- value=2, #Replace with defaults that work for your model
129
  )
130
-
 
131
  gr.Examples(
132
- examples = examples,
133
- inputs = [prompt]
134
  )
135
- gr.on(
136
- triggers=[run_button.click, prompt.submit],
137
- fn = infer,
138
- inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
139
- outputs = [result, seed]
140
- )
141
 
142
- demo.queue().launch()
 
 
 
1
  import gradio as gr
2
  import numpy as np
3
  import random
 
4
  from diffusers import DiffusionPipeline
5
  import torch
6
 
7
+ # Set the device based on availability
8
  device = "cuda" if torch.cuda.is_available() else "cpu"
 
9
 
10
+ # Use the ByteDance/AnimateDiff-Lightning model
11
+ model_repo_id = "ByteDance/AnimateDiff-Lightning"
12
+
13
+ # Set the torch dtype based on available hardware
14
  if torch.cuda.is_available():
15
  torch_dtype = torch.float16
16
  else:
17
  torch_dtype = torch.float32
18
 
19
+ # Load the pipeline from the pretrained model repository
20
  pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
21
  pipe = pipe.to(device)
22
 
23
+ # Maximum values for seed and image size
24
  MAX_SEED = np.iinfo(np.int32).max
25
  MAX_IMAGE_SIZE = 1024
26
 
27
+ # Define the inference function
28
  def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
29
 
30
+ # Randomize seed if the checkbox is selected
31
  if randomize_seed:
32
  seed = random.randint(0, MAX_SEED)
33
 
34
+ generator = torch.Generator(device=device).manual_seed(seed)
35
+
36
+ # Generate the animation using the pipeline
37
+ animation = pipe(
38
+ prompt=prompt,
39
+ negative_prompt=negative_prompt,
40
+ guidance_scale=guidance_scale,
41
+ num_inference_steps=num_inference_steps,
42
+ width=width,
43
+ height=height,
44
+ generator=generator
45
+ ).images[0] # Assuming the model generates images in the `.images` property
46
 
47
+ return animation, seed
48
 
49
+ # Sample prompts for the UI
50
  examples = [
51
+ "A cat playing with a ball in a garden",
52
+ "A dancing astronaut in space",
53
+ "A flying dragon in the sky at sunset",
54
  ]
55
 
56
+ # Define CSS for styling
57
+ css = """
58
  #col-container {
59
  margin: 0 auto;
60
  max-width: 640px;
61
  }
62
  """
63
 
64
+ # Build the Gradio UI
65
  with gr.Blocks(css=css) as demo:
66
 
67
  with gr.Column(elem_id="col-container"):
68
  gr.Markdown(f"""
69
+ # AnimateDiff Lightning Model Text-to-Animation
70
  """)
71
 
72
  with gr.Row():
 
73
  prompt = gr.Text(
74
  label="Prompt",
75
  show_label=False,
 
77
  placeholder="Enter your prompt",
78
  container=False,
79
  )
 
80
  run_button = gr.Button("Run", scale=0)
81
 
82
+ result = gr.Image(label="Generated Animation", show_label=False)
83
 
84
  with gr.Accordion("Advanced Settings", open=False):
85
 
 
87
  label="Negative prompt",
88
  max_lines=1,
89
  placeholder="Enter a negative prompt",
90
+ visible=True,
91
  )
92
 
93
  seed = gr.Slider(
 
101
  randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
102
 
103
  with gr.Row():
 
104
  width = gr.Slider(
105
  label="Width",
106
  minimum=256,
107
  maximum=MAX_IMAGE_SIZE,
108
  step=32,
109
+ value=1024,
110
  )
 
111
  height = gr.Slider(
112
  label="Height",
113
  minimum=256,
114
  maximum=MAX_IMAGE_SIZE,
115
  step=32,
116
+ value=1024,
117
  )
118
 
119
  with gr.Row():
 
120
  guidance_scale = gr.Slider(
121
  label="Guidance scale",
122
  minimum=0.0,
123
  maximum=10.0,
124
  step=0.1,
125
+ value=7.5,
126
  )
 
127
  num_inference_steps = gr.Slider(
128
  label="Number of inference steps",
129
  minimum=1,
130
  maximum=50,
131
  step=1,
132
+ value=30,
133
  )
134
+
135
+ # Example prompts for user selection
136
  gr.Examples(
137
+ examples=examples,
138
+ inputs=[prompt]
139
  )
 
 
 
 
 
 
140
 
141
+ # Create an API endpoint for the model
142
+ demo.api(fn=infer, inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps], outputs=[result, seed])
143
+ demo.launch()