KingNish commited on
Commit
87a083e
1 Parent(s): a70c106

Update app.py

Browse files

Addes Dalle lora for better image (test)

Files changed (1) hide show
  1. app.py +7 -11
app.py CHANGED
@@ -8,7 +8,7 @@ import numpy as np
8
  from PIL import Image
9
  import spaces
10
  import torch
11
- from diffusers import DiffusionPipeline
12
 
13
  if not torch.cuda.is_available():
14
  DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"
@@ -22,21 +22,17 @@ ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
22
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
23
 
24
  if torch.cuda.is_available():
25
- pipe = DiffusionPipeline.from_pretrained(
26
  "sd-community/sdxl-flash",
27
  torch_dtype=torch.float16,
28
  use_safetensors=True,
29
  add_watermarker=False
30
  )
31
- if ENABLE_CPU_OFFLOAD:
32
- pipe.enable_model_cpu_offload()
33
- else:
34
- pipe.to(device)
35
- print("Loaded on Device!")
36
-
37
- if USE_TORCH_COMPILE:
38
- pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
39
- print("Model Compiled!")
40
 
41
 
42
  def save_image(img):
 
8
  from PIL import Image
9
  import spaces
10
  import torch
11
+ from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
12
 
13
  if not torch.cuda.is_available():
14
  DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"
 
22
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
23
 
24
  if torch.cuda.is_available():
25
+ pipe = StableDiffusionXLPipeline.from_pretrained(
26
  "sd-community/sdxl-flash",
27
  torch_dtype=torch.float16,
28
  use_safetensors=True,
29
  add_watermarker=False
30
  )
31
+ pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
32
+ pipe.load_lora_weights("ehristoforu/dalle-3-xl-v2", weight_name="dalle-3-xl-lora-v2.safetensors", adapter_name="dalle")
33
+ pipe.set_adapters("dalle")
34
+
35
+ pipe.to("cuda")
 
 
 
 
36
 
37
 
38
  def save_image(img):