Update app.py

app.py

@@ -1,20 +1,36 @@
 import gradio as gr
 from transformers import pipeline
+from diffusers import StableDiffusionPipeline
+import torch
 
-# Load models
-classifier_env = pipeline("text-classification", model="bhadresh-savani/distilbert-base-uncased-emotion")  # Replace with your own if needed
-image_generator = pipeline("text-to-image", model="CompVis/stable-diffusion-v1-4")
+# ==============================
+# Load Pipelines
+# ==============================
+
+# Sentence classification (use a multi-class model if available)
+classifier_env = pipeline("text-classification", model="bhadresh-savani/distilbert-base-uncased-emotion")  # You can fine-tune your own later
+
+# NER pipeline
 ner_model = pipeline("ner", grouped_entities=True)
+
+# Masked word fill
 fill_mask = pipeline("fill-mask", model="roberta-large")
 
-# Functions
+# Stable Diffusion for image generation
+pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32)
+pipe = pipe.to("cpu")
+
+# ==============================
+# Define Functions
+# ==============================
+
 def classify_environment(text):
     result = classifier_env(text)
     return result[0]['label'], result[0]['score']
 
 def generate_image(prompt):
-    result = image_generator(prompt)
-    return result[0]
+    image = pipe(prompt).images[0]
+    return image
 
 def ner_mapping(text):
     entities = ner_model(text)
@@ -24,9 +40,12 @@ def fill_masked_text(text):
     results = fill_mask(text)
     return "\n".join([r['sequence'] for r in results])
 
-# Gradio Interfaces
+# ==============================
+# Gradio Interface
+# ==============================
+
 with gr.Blocks() as demo:
-    gr.Markdown("# 🌍 Multi-Model NLP + Image App (Hugging Face)")
+    gr.Markdown("# 🤖 Multi-Model NLP + Image Generator (Hugging Face Space)")
 
     with gr.Tab("Sentence Classification (Environment)"):
         input_text = gr.Textbox(placeholder="Type a sentence about the environment...", lines=2)
@@ -53,5 +72,4 @@ with gr.Blocks() as demo:
         mask_btn = gr.Button("Fill Mask")
         mask_btn.click(fn=fill_masked_text, inputs=mask_input, outputs=mask_output)
 
-# Launch
 demo.launch()