Update app.py

app.py

@@ -1,103 +1,47 @@
+# app.py
+
 import streamlit as st
-from streamlit_aggrid import AgGrid, GridOptionsBuilder, GridUpdateMode, DataReturnMode
-from diffusers import DiffusionPipeline, StableDiffusionPipeline, DDIMScheduler, LMSDiscreteScheduler, EulerDiscreteScheduler
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from datasets import load_dataset
 import torch
-
-# Set page title and favicon
-st.set_page_config(page_title="Diffusers UI", page_icon=":art:")
-
-# Sidebar for selecting models and tasks
-st.sidebar.title("Diffusers UI")
-model_type = st.sidebar.selectbox("Select Model Type", ["Text-to-Image", "Image-to-Image", "Text-to-Text"])
-task = st.sidebar.selectbox("Select Task", ["Image Generation", "Image Editing", "Text Generation"])
-
-# Load Hugging Face Hub models based on selected model type and task
-if model_type == "Text-to-Image":
-    if task == "Image Generation":
-        model_id = "CompVis/stable-diffusion-v1-4"
-    elif task == "Image Editing":
-        model_id = "runwayml/stable-diffusion-inpainting"
-elif model_type == "Image-to-Image":
-    if task == "Image Generation":
-        model_id = "CompVis/stable-diffusion-v1-4"
-    elif task == "Image Editing":
-        model_id = "CompVis/stable-diffusion-v1-4"
-elif model_type == "Text-to-Text":
-    if task == "Text Generation":
-        model_id = "gpt2"
-
-# Load model and tokenizer
-if model_type in ["Text-to-Image", "Image-to-Image"]:
-    try:
-        pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
-    except Exception as e:
-        st.error(f"Error loading model: {e}")
-        st.stop()
-elif model_type == "Text-to-Text":
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        model = AutoModelForCausalLM.from_pretrained(model_id)
-    except Exception as e:
-        st.error(f"Error loading model: {e}")
-        st.stop()
-
-# Main app layout
-st.title("Diffusers UI")
-
-# Input section
-with st.form("input_form"):
-    if model_type in ["Text-to-Image", "Image-to-Image"]:
-        prompt = st.text_input("Enter a prompt", "a photo of an astronaut riding a horse")
-        num_inference_steps = st.slider("Number of inference steps", 1, 50, 25)
-        guidance_scale = st.slider("Guidance scale", 1.0, 10.0, 7.5)
-        image = st.file_uploader("Upload an image (optional)", type=["png", "jpg", "jpeg"])
-    elif model_type == "Text-to-Text":
-        input_text = st.text_input("Enter input text", "Hello, my name is")
-        max_length = st.slider("Maximum length of generated text", 1, 100, 20)
-
-    submitted = st.form_submit_button("Generate")
-
-# Output section
-if submitted:
-    if model_type in ["Text-to-Image", "Image-to-Image"]:
-        if image is not None:
-            image = Image.open(image)
-            image = image.resize((768, 768))
-            image = np.array(image).astype(np.float32) / 255.0
-            image = image[None].transpose(0, 3, 1, 2)
-            image = torch.from_numpy(image)
-            image = 2.0 * image - 1.0
-            image = image.to(pipe.device)
-
-            with st.spinner("Generating image..."):
-                try:
-                    image = pipe(prompt, image=image, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
-                except Exception as e:
-                    st.error(f"Error generating image: {e}")
-                    st.stop()
-
-            st.image(image, caption="Generated Image", use_column_width=True)
-        else:
-            with st.spinner("Generating image..."):
-                try:
-                    image = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
-                except Exception as e:
-                    st.error(f"Error generating image: {e}")
-                    st.stop()
-
-            st.image(image, caption="Generated Image", use_column_width=True)
-    elif model_type == "Text-to-Text":
-        with st.spinner("Generating text..."):
-            try:
-                input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(model.device)
-                generated_ids = model.generate(input_ids, max_length=max_length)
-                generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-            except Exception as e:
-                st.error(f"Error generating text: {e}")
-                st.stop()
-
-        st.write("Generated Text:")
-        st.write(generated_text)
+from diffusers import DiffusionPipeline
+from PIL import Image
+
+# Set up the Streamlit app
+st.set_page_config(page_title="Huggingface Diffusers Showcase", page_icon=":guardsman:", layout="wide")
+
+# Add a sidebar
+st.sidebar.title("Navigation")
+st.sidebar.markdown("## Choose a Model")
+model_options = ["stable-diffusion", "ddim", "ddpm"]
+selected_model = st.sidebar.selectbox("Select a model", model_options)
+
+# Load the selected model
+if selected_model == "stable-diffusion":
+    pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-512-v2")
+elif selected_model == "ddim":
+    pipe = DiffusionPipeline.from_pretrained("CompVis/ddim-512-fp16")
+else:
+    pipe = DiffusionPipeline.from_pretrained("CompVis/ddpm-512-fp16")
+
+# Add a text input for the user to enter a prompt
+prompt = st.text_input("Enter a text prompt:", "A sunset over a mountain")
+
+# Add a button to generate the image
+if st.button("Generate Image"):
+    # Generate the image
+    image = pipe(prompt, num_inference_steps=50).images[0]
+
+    # Display the generated image
+    col1, col2 = st.columns(2)
+    with col1:
+        st.subheader("Prompt")
+        st.markdown(prompt)
+    with col2:
+        st.subheader("Generated Image")
+        st.image(image, caption="Generated using the selected diffusion model.", use_column_width=True)
+
+# Add a section for the app info
+st.sidebar.markdown("---")
+st.sidebar.markdown("## App Info")
+st.sidebar.markdown("This app showcases the Huggingface Diffusers library and its features.")
+st.sidebar.markdown("You can select a model from the sidebar and enter a text prompt to generate an image.")
+st.sidebar.markdown("Visit the [Huggingface Diffusers documentation](https://huggingface.co/docs/diffusers) for more information.")