Yaron Koresh commited on
Commit
94eb3b9
·
verified ·
1 Parent(s): e95a30b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -6
app.py CHANGED
@@ -28,14 +28,19 @@ def port_inc():
28
  else:
29
  os.environ["CUSTOM_PORT"]=str(int(env)+1)
30
 
31
- def init_pool():
 
 
 
 
32
  port_inc()
33
 
34
  #pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, token=os.getenv("hf_token")).to(device)
35
  #pipe2 = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to(device)
36
  #pipe2.unet = torch.compile(pipe2.unet, mode="reduce-overhead", fullgraph=True)
37
 
38
- PIPE = None
 
39
 
40
  def pipe_t2i():
41
  global PIPE
@@ -96,7 +101,7 @@ def generate_random_string(length):
96
  return ''.join(random.choice(characters) for _ in range(length))
97
 
98
  @spaces.GPU(duration=55)
99
- def Piper(pipe,_do):
100
  pipe = pipe_t2i()
101
  try:
102
  retu = pipe(
@@ -139,7 +144,7 @@ def main(p1,p2,*result):
139
  rng = list(range(ln))
140
  arr = [p for _ in rng]
141
  out = None
142
- with Pool(ln, initializer=init_pool) as pool:
143
  out = pool.imap(infer,arr)
144
  return list(out)
145
 
@@ -155,7 +160,7 @@ def infer(p):
155
  neg = _dont
156
  else:
157
  neg = None
158
- output = Piper(pipe_t2i(),'A '+" ".join(_do))
159
  if output == None:
160
  return None
161
  else:
@@ -163,7 +168,7 @@ def infer(p):
163
  if neg == None:
164
  return name
165
  img = load_image(name).convert("RGB")
166
- output2 = Piper2(pipe_i2i(),img,p1,neg)
167
  if output2 == None:
168
  return None
169
  else:
 
28
  else:
29
  os.environ["CUSTOM_PORT"]=str(int(env)+1)
30
 
31
+ def init_pool(_1,_2):
32
+ global infer1
33
+ global infer2
34
+ infer1 = _1
35
+ infer2 = _2
36
  port_inc()
37
 
38
  #pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, token=os.getenv("hf_token")).to(device)
39
  #pipe2 = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to(device)
40
  #pipe2.unet = torch.compile(pipe2.unet, mode="reduce-overhead", fullgraph=True)
41
 
42
+ pp1=pipe_t2i()
43
+ pp2=pipe_i2i()
44
 
45
  def pipe_t2i():
46
  global PIPE
 
101
  return ''.join(random.choice(characters) for _ in range(length))
102
 
103
  @spaces.GPU(duration=55)
104
+ def Piper1(pipe,_do):
105
  pipe = pipe_t2i()
106
  try:
107
  retu = pipe(
 
144
  rng = list(range(ln))
145
  arr = [p for _ in rng]
146
  out = None
147
+ with Pool(ln, initializer=init_pool, initargs=(pp1,pp2)) as pool:
148
  out = pool.imap(infer,arr)
149
  return list(out)
150
 
 
160
  neg = _dont
161
  else:
162
  neg = None
163
+ output = Piper1(infer1,'A '+" ".join(_do))
164
  if output == None:
165
  return None
166
  else:
 
168
  if neg == None:
169
  return name
170
  img = load_image(name).convert("RGB")
171
+ output2 = Piper2(infer2,img,p1,neg)
172
  if output2 == None:
173
  return None
174
  else: