Update app.py

app.py

@@ -430,36 +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)
+
+
+# text generation
 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)
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+st.title("Text Generation with Bloom")
+
+tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom")
+model = AutoModelForCausalLM.from_pretrained("bigscience/bloom")
+
+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)