Update app.py

app.py

@@ -1,52 +1,57 @@
-import streamlit as st
-import replicate
-import requests
-from PIL import Image
-from io import BytesIO
-import threading
-
-# Set the replicate API token
-import os
-os.environ["REPLICATE_API_TOKEN"] = "r8_JSR8xlRoCk6cmq3qEOOThVTn3dAgdPq1bWXdj"
-
-model_name = "fofr/sdxl-turbo:6244ebc4d96ffcc48fa1270d22a1f014addf79c41732fe205fb1ff638c409267"
-
-class PromptGenerator(threading.Thread):
-    def __init__(self, prompt, update_prompt, lock):
-        self.prompt = prompt
-        self.update_prompt = update_prompt
-        self.lock = lock
-        super(PromptGenerator, self).__init__()
-
-    def run(self):
-        while True:
-            new_prompt = input('Enter a new prompt: ')
-            with self.lock:
-                self.prompt = new_prompt
-                self.update_prompt(self.prompt)
-
-def generate_output(prompt):
-    return replicate.run(model_name, input={"prompt": prompt})
+from diffusers import AutoPipelineForText2Image
+import torch
 
-st.title("Hugging Face Model Real-time Interface")
+pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16").to("cuda")
 
-lock = threading.Lock()
-
-prompt = st.text_input("Enter your prompt here")
-update_prompt = st.empty()
-
-def update_promt_thread(prompt):
-    with lock:
-        update_prompt.text(prompt)
-
-t = PromptGenerator(prompt, update_prompt, lock)
-t.start()
-
-output_url = generate_output(prompt)
+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)
 
-try:
-    response = requests.get(output_url)
-    img = Image.open(BytesIO(response.content))
-    st.image(img, caption="Model Output", use_column_width=True)
-except Exception as e:
-    st.write("Unable to display the image.")
+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)