MegaTronX commited on
Commit
3257204
·
verified ·
1 Parent(s): 3043719

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -6
app.py CHANGED
@@ -1,16 +1,19 @@
1
  import gradio as gr
2
- #import spaces
3
  from diffusers import DiffusionPipeline
4
  import torch
5
 
 
 
 
 
6
  pipe = DiffusionPipeline.from_pretrained(
7
  "black-forest-labs/FLUX.1-dev",
8
- torch_dtype=torch.float16
9
- )
10
- pipe.load_lora_weights("MegaTronX/MetartLoRA", weight_name="MetartLoRA.safetensors")
11
- pipe = pipe.to("cuda")
12
 
13
- #@spaces.GPU(duration=75)
14
  def generate_image(prompt, num_inference_steps=25, guidance_scale=7.5, seed=None):
15
  """Generates an image using the FLUX.1-dev LoRA model."""
16
  generator = torch.Generator("cuda").manual_seed(seed) if seed else None
 
1
  import gradio as gr
2
+ import spaces
3
  from diffusers import DiffusionPipeline
4
  import torch
5
 
6
+
7
+ device = "cuda" if torch.cuda.is_available() else "cpu"
8
+
9
+
10
  pipe = DiffusionPipeline.from_pretrained(
11
  "black-forest-labs/FLUX.1-dev",
12
+ torch_dtype=torch.bfloat16
13
+ ).to(device)
14
+ pipe.load_lora_weights("MegaTronX/MetartLoRA", weight_name="MetartLoRA.safetensors").to(device)
 
15
 
16
+ @spaces.GPU(duration=75)
17
  def generate_image(prompt, num_inference_steps=25, guidance_scale=7.5, seed=None):
18
  """Generates an image using the FLUX.1-dev LoRA model."""
19
  generator = torch.Generator("cuda").manual_seed(seed) if seed else None