callum-canavan commited on
Commit
b6fc2cb
·
1 Parent(s): 970b026

Update app

Browse files
Files changed (2) hide show
  1. diffuse.py +0 -34
  2. requirements.txt +1 -0
diffuse.py DELETED
@@ -1,34 +0,0 @@
1
- from diffusers import DiffusionPipeline
2
- from diffusers.utils import pt_to_pil
3
- import torch
4
-
5
- # stage 1
6
- stage_1 = DiffusionPipeline.from_pretrained(
7
- "DeepFloyd/IF-I-M-v1.0", variant="fp16", torch_dtype=torch.float16
8
- )
9
- stage_1.enable_xformers_memory_efficient_attention() # remove line if torch.__version__ >= 2.0.0
10
- stage_1.enable_model_cpu_offload()
11
-
12
- # stage 2
13
- stage_2 = DiffusionPipeline.from_pretrained(
14
- "DeepFloyd/IF-II-M-v1.0",
15
- text_encoder=None,
16
- variant="fp16",
17
- torch_dtype=torch.float16,
18
- )
19
- stage_2.enable_xformers_memory_efficient_attention() # remove line if torch.__version__ >= 2.0.0
20
- stage_2.enable_model_cpu_offload()
21
-
22
- # stage 3
23
- safety_modules = {
24
- "feature_extractor": stage_1.feature_extractor,
25
- "safety_checker": stage_1.safety_checker,
26
- "watermarker": stage_1.watermarker,
27
- }
28
- stage_3 = DiffusionPipeline.from_pretrained(
29
- "stabilityai/stable-diffusion-x4-upscaler",
30
- **safety_modules,
31
- torch_dtype=torch.float16
32
- )
33
- stage_3.enable_xformers_memory_efficient_attention() # remove line if torch.__version__ >= 2.0.0
34
- stage_3.enable_model_cpu_offload()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
requirements.txt CHANGED
@@ -1,2 +1,3 @@
 
1
  transformers
2
  torch
 
1
+ gradio
2
  transformers
3
  torch