Update app.py

app.py

@@ -1,5 +1,4 @@
 import streamlit as st
-import streamlit_ace
 from streamlit_aggrid import AgGrid, GridOptionsBuilder, GridUpdateMode, DataReturnMode
 from diffusers import DiffusionPipeline, StableDiffusionPipeline, DDIMScheduler, LMSDiscreteScheduler, EulerDiscreteScheduler
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
@@ -31,11 +30,19 @@ elif model_type == "Text-to-Text":
 
 # Load model and tokenizer
 if model_type in ["Text-to-Image", "Image-to-Image"]:
-    pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+    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":
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(model_id)
+    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")
@@ -66,121 +73,31 @@ if submitted:
             image = image.to(pipe.device)
 
             with st.spinner("Generating image..."):
-                image = pipe(prompt, image=image, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
+                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..."):
-                image = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
+                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..."):
-            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]
+            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)
-
-# Code editor section
-st.header("Code Editor")
-with st.expander("View Code"):
-    code = 
-import streamlit as st
-import streamlit_ace
-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"]:
-    pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
-elif model_type == "Text-to-Text":
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(model_id)
-
-# 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..."):
-                image = pipe(prompt, image=image, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
-
-            st.image(image, caption="Generated Image", use_column_width=True)
-        else:
-            with st.spinner("Generating image..."):
-                image = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
-
-            st.image(image, caption="Generated Image", use_column_width=True)
-    elif model_type == "Text-to-Text":
-        with st.spinner("Generating text..."):
-            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]
-
-        st.write("Generated Text:")
-        st.write(generated_text)
-"""
-# Code editor section
-st.header("Code Editor")
-with st.expander("View Code"):
-    code = 
-    streamlit_ace.ace(code, type="python", theme="monokai", height=500)
-    streamlit_ace.ace(code, type="python", theme="monokai", height=500)
-"""
-
-# Run the app
-if __name__ == "__main__":
-    main()
-    streamlit_ace.ace(code, type="python", theme="monokai", height=500)
+        st.write(generated_text)