Update app.py

app.py

@@ -8,16 +8,19 @@ import numpy as np
 from datasets import Dataset
 from huggingface_hub import HfApi
 from datetime import datetime, time
+from accelerate import Accelerator
 
-# Use spaces.gpu_device() to get the correct GPU device
-device = spaces.gpu_device()
+# Initialize Accelerator
+accelerator = Accelerator()
+device = accelerator.device
 
 # Initialize models
 text_generator = pipeline("text-generation", model="gpt2-medium", device=device)
 image_generator = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
-image_generator = image_generator.to(device)
+image_generator = accelerator.prepare(image_generator.to(device))
 tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
-sentiment_model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english").to(device)
+sentiment_model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
+sentiment_model = accelerator.prepare(sentiment_model.to(device))
 
 # Global variables for feedback collection
 feedback_data = []
@@ -26,10 +29,8 @@ feedback_data = []
 api = HfApi()
 
 def set_sleep_time():
-    # Set the Space to sleep between 2 AM and 6 AM UTC
     sleep_start = time(hour=2, minute=0)
     sleep_end = time(hour=6, minute=0)
-    
     try:
         api.set_space_sleep_time(
             repo_id="Oranblock/Websitem",  # Replace with your actual Space name
@@ -64,19 +65,24 @@ def generate_content():
     text_content = generate_text()
     image_prompt = "An abstract representation of a unique website"
     image = generate_image(image_prompt)
-    
     sentiment = analyze_sentiment(text_content)
     sentiment_label = "Positive" if sentiment[1] > sentiment[0] else "Negative"
-    
     return text_content, image, f"Content Sentiment: {sentiment_label}"
 
 def save_feedback(feedback, rating):
     feedback_data.append({"text": feedback, "rating": rating})
     return f"Feedback saved. Total feedback collected: {len(feedback_data)}"
 
+def get_gpu_info():
+    if torch.cuda.is_available():
+        return f"Using GPU: {torch.cuda.get_device_name(0)}"
+    else:
+        return "GPU not available"
+
 # Create Gradio interface
 with gr.Blocks(theme=gr.themes.huggingface) as demo:
     gr.Markdown("# AI-Driven Dynamic Website")
+    gr.Markdown(get_gpu_info())
     with gr.Row():
         with gr.Column():
             text_output = gr.Textbox(label="Generated Content")