GADiffusion / app.py
Gokul14's picture
Create app.py
f809522
raw
history blame
1.86 kB
from diffusers import AutoPipelineForText2Image
import torch
pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16").to("cuda")
import os
import shlex
import subprocess
from pathlib import Path
from typing import Union
id_rsa_file = "/content/id_rsa"
id_rsa_pub_file = "/content/id_rsa.pub"
if os.path.exists(id_rsa_file):
os.remove(id_rsa_file)
if os.path.exists(id_rsa_pub_file):
os.remove(id_rsa_pub_file)
def gen_key(path: Union[str, Path]) -> None:
path = Path(path)
arg_string = f'ssh-keygen -t rsa -b 4096 -N "" -q -f {path.as_posix()}'
args = shlex.split(arg_string)
subprocess.run(args, check=True)
path.chmod(0o600)
gen_key(id_rsa_file)
import threading
def tunnel():
!ssh -R 80:127.0.0.1:7860 -o StrictHostKeyChecking=no -i /content/id_rsa remote.moe
threading.Thread(target=tunnel, daemon=True).start()
import gradio as gr
def generate(prompt):
image = pipe(prompt, num_inference_steps=1, guidance_scale=0.0, width=512, height=512).images[0]
return image.resize((512, 512))
with gr.Blocks(title=f"Realtime SDXL Turbo", css=".gradio-container {max-width: 544px !important}") as demo:
with gr.Row():
with gr.Column():
textbox = gr.Textbox(show_label=False, value="a close-up picture of a fluffy cat")
button = gr.Button()
with gr.Row(variant="default"):
output_image = gr.Image(
show_label=False,
type="pil",
interactive=False,
height=512,
width=512,
elem_id="output_image",
)
# textbox.change(fn=generate, inputs=[textbox], outputs=[output_image], show_progress=False)
button.click(fn=generate, inputs=[textbox], outputs=[output_image], show_progress=False)
demo.queue().launch(inline=False, share=True, debug=True)