Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,778 Bytes
c2c42ca 91dd651 d3a77ee 91dd651 c2c42ca 91dd651 de56cd9 9d41bd5 c2c42ca 9d41bd5 91dd651 c2c42ca 91dd651 c2c42ca 91dd651 143f063 91dd651 143f063 91dd651 143f063 c2c42ca 34cb1b5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import gradio as gr
from gradio_client import Client
from diffusers import AutoencoderKL, StableDiffusionXLPipeline
import concurrent.futures
client_lightning = Client("AP123/SDXL-Lightning")
client_hyper = Client("ByteDance/Hyper-SDXL-1Step-T2I")
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
### SDXL Turbo ####
pipe_turbo = StableDiffusionXLPipeline.from_pretrained("stabilityai/sdxl-turbo",
vae=vae,
torch_dtype=torch.float16,
variant="fp16"
)
pipe_turbo.to("cuda")
def get_lighting_result(prompt):
result_lighting = client_lightning.predict(
prompt, # Your prompt
"1-Step", # Number of inference steps
api_name="/generate_image"
)
return result_lighting
def get_hyper_result(prompt):
result_hyper = client_hyper.predict(
num_images=1,
height=1024,
width=1024,
prompt=prompt,
seed=3413,
api_name="/process_image"
)
return result_hyper
@spaces.GPU
def get_turbo_result(prompt):
image_turbo = pipe_turbo(prompt=prompt, num_inference_steps=1, guidance_scale=0).images[0]
return image_turbo
def run_in_parallel(prompt):
with concurrent.futures.ThreadPoolExecutor() as executor:
# Submit tasks to the executor
future_lighting = executor.submit(get_lighting_result, prompt)
future_hyper = executor.submit(get_hyper_result, prompt)
future_turbo = executor.submit(get_turbo_result, prompt)
# Wait for all futures to complete
results = concurrent.futures.wait(
[future_lighting, future_hyper, future_turbo],
return_when=concurrent.futures.ALL_COMPLETED
)
# Extract results from futures
result_lighting = future_lighting.result()
result_hyper = future_hyper.result()
image_turbo = future_turbo.result()
print(result_lighting)
print(result_hyper)
return image_turbo, result_lighting, result_hyper
# Example usage
prompt = "Enter your prompt here"
image_turbo, result_lighting, result_hyper = run_in_parallel(prompt)
css = '''
.gradio-container{max-width: 768px !important}
'''
with gr.Blocks(css=css) as demo:
prompt = gr.Textbox(label="Prompt")
run = gr.Button("Run")
with gr.Row():
image_turbo = gr.Image(label="SDXL Turbo")
image_lightning = gr.Image(label="SDXL Lightning")
image_hyper = gr.Image("Hyper SDXL")
run.click(fn=run_comparison, inputs=prompt, outputs=[image_turbo, image_lightning, image_hyper])
|