Update app.py

app.py

@@ -1,31 +1,22 @@
-import requests
-import os
-from dotenv import load_dotenv  # Install the 'python-dotenv' package
-
-# Load environment variables from a .env file
-load_dotenv()
-
-API_URL = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
-API_TOKEN = os.getenv("HF_API_TOKEN")  # Use os.getenv to retrieve the API token from environment variables
-
-if API_TOKEN is None:
-    raise ValueError("HF_API_TOKEN environment variable not set")
-
-headers = {"Authorization": f"Bearer {API_TOKEN}"}
+import streamlit as st
+from transformers import CLIPProcessor, CLIPModel, DiffusionModel
+import torch
+from PIL import Image
 
-def query(payload):
-    response = requests.post(API_URL, headers=headers, json=payload)
-    return response.content
+st.title("Text to Image Generation")
 
-image_bytes = query({
-    "inputs": "cat is drinking milk ",
-})
+# Load pretrained models
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
+diffusion_model = DiffusionModel.from_pretrained("openai/guided-diffusion-clipped-coco")
 
-import io
-from PIL import Image
+text = st.text_area("Enter a description:")
+if st.button("Generate Image") and text:
+    # Process text and get CLIP features
+    text_features = clip_processor(text, return_tensors="pt", padding=True)
 
-image = Image.open(io.BytesIO(image_bytes))
-image.show()
+    # Generate image from text using Guided Diffusion
+    image = diffusion_model.generate_text_to_image(text_features["pixel_values"])
 
-# If you want to save the generated image to a file, you can use the save method
-image.save("generated_image.png")
+    # Display the generated image
+    st.image(image, caption="Generated Image", use_column_width=True)