Update app.py

app.py

@@ -430,55 +430,55 @@
 
  # text to Image
 
-# import streamlit as st
-# import torch
-# from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
-# from huggingface_hub import hf_hub_download
-# from safetensors.torch import load_file
-
-# # Model Path/Repo Information
-# base = "stabilityai/stable-diffusion-xl-base-1.0"
-# repo = "ByteDance/SDXL-Lightning"
-# ckpt = "sdxl_lightning_4step_unet.safetensors" 
-
-# # Load model (Executed only once for efficiency)
-# @st.cache_resource
-# def load_sdxl_pipeline():
-#     unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
-#     unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
-#     pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
-#     pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
-#     return pipe
-
-# # Streamlit UI
-# st.title("Image Generation")
-# prompt = st.text_input("Enter your image prompt:")
-
-# if st.button("Generate Image"):
-#     if not prompt:
-#         st.warning("Please enter a prompt.")
-#     else:
-#         pipe = load_sdxl_pipeline()  # Load the pipeline from cache
-#         with torch.no_grad():
-#             image = pipe(prompt).images[0] 
-
-#         st.image(image)
+import streamlit as st
+import torch
+from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+# Model Path/Repo Information
+base = "stabilityai/stable-diffusion-xl-base-1.0"
+repo = "ByteDance/SDXL-Lightning"
+ckpt = "sdxl_lightning_4step_unet.safetensors" 
+
+# Load model (Executed only once for efficiency)
+@st.cache_resource
+def load_sdxl_pipeline():
+    unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
+    unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
+    pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
+    pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
+    return pipe
+
+# Streamlit UI
+st.title("Image Generation")
+prompt = st.text_input("Enter your image prompt:")
+
+if st.button("Generate Image"):
+    if not prompt:
+        st.warning("Please enter a prompt.")
+    else:
+        pipe = load_sdxl_pipeline()  # Load the pipeline from cache
+        with torch.no_grad():
+            image = pipe(prompt).images[0] 
+
+        st.image(image)
 
 
 # text generation
-import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
+# import streamlit as st
+# from transformers import AutoTokenizer, AutoModelForCausalLM
 
-st.title("Text Generation with Bloom")
+# st.title("Text Generation with Bloom")
 
-tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom")
-model = AutoModelForCausalLM.from_pretrained("bigscience/bloom")
+# tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom")
+# model = AutoModelForCausalLM.from_pretrained("bigscience/bloom")
 
-user_input = st.text_area("Enter your prompt:", height=100)
+# user_input = st.text_area("Enter your prompt:", height=100)
 
-if st.button('Generate Text'):
-    inputs = tokenizer(user_input, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=100)  # Adjust max_length as needed
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True) 
-    st.write("Generated Text:")
-    st.write(generated_text)
+# if st.button('Generate Text'):
+#     inputs = tokenizer(user_input, return_tensors="pt")
+#     outputs = model.generate(**inputs, max_length=100)  # Adjust max_length as needed
+#     generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True) 
+#     st.write("Generated Text:")
+#     st.write(generated_text)