Update app.py

app.py

@@ -1,57 +1,23 @@
-from diffusers import AutoPipelineForText2Image
+import streamlit as st
 import torch
+from diffusers import StableDiffusionXLPipeline
 
-pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16").to("cuda")
+# Load the stable diffusion model
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
+)
+pipe = pipe.to("cuda")  # Move the model to GPU for faster processing
 
-import os
-import shlex
-import subprocess
-from pathlib import Path
-from typing import Union
+# Create a Streamlit app
+st.title("Image Generation App")
 
-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)
+# Add a text input for the prompt
+prompt = st.text_input("Prompt", "a photo of an astronaut riding a horse on mars")
 
-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)
+# Generate the image based on the prompt
+if st.button("Generate Image"):
+    # Use the model to generate the image
+    image = pipe(prompt).images[0]
 
-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)
+    # Display the generated image
+    st.image(image, caption="Generated Image", use_column_width=True)