Arslan17121 commited on
Commit
8c06470
Β·
verified Β·
1 Parent(s): 70dd853

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -3
app.py CHANGED
@@ -2,11 +2,16 @@ from diffusers import StableDiffusionPipeline
2
  from PIL import Image
3
  import imageio
4
  import gradio as gr
 
 
 
 
 
5
 
6
  # Load Stable Diffusion model
7
- model_id = "runwayml/stable-diffusion-v1-5" # Or your custom fine-tuned model
8
- pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype="float16", safety_checker=None)
9
- pipe.to("cuda") # Use GPU if available
10
 
11
  def generate_frames(prompt, num_frames=5, style="realistic"):
12
  """
 
2
  from PIL import Image
3
  import imageio
4
  import gradio as gr
5
+ import torch
6
+
7
+ # Dynamic device and data type allocation
8
+ device = "cuda" if torch.cuda.is_available() else "cpu"
9
+ dtype = torch.float16 if torch.cuda.is_available() else torch.float32
10
 
11
  # Load Stable Diffusion model
12
+ model_id = "runwayml/stable-diffusion-v1-5" # Replace with your desired model ID
13
+ pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype, safety_checker=None)
14
+ pipe.to(device)
15
 
16
  def generate_frames(prompt, num_frames=5, style="realistic"):
17
  """