Update app.py

app.py

@@ -3,24 +3,27 @@ import gradio as gr
 from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
 from diffusers import StableDiffusionPipeline
 import torch
-import random
 import numpy as np
 from datasets import Dataset
 from huggingface_hub import HfApi
 from datetime import datetime, time
 from accelerate import Accelerator
+from accelerate.utils import set_seed
+
+# Set a seed for reproducibility
+set_seed(42)
 
 # Initialize Accelerator
 accelerator = Accelerator()
-device = accelerator.device
+device = 0 if torch.cuda.is_available() else -1  # Use GPU 0 if available, else CPU
 
 # 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 = accelerator.prepare(image_generator.to(device))
+image_generator = image_generator.to(f"cuda:{device}" if device >= 0 else "cpu")
 tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
 sentiment_model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
-sentiment_model = accelerator.prepare(sentiment_model.to(device))
+sentiment_model = sentiment_model.to(f"cuda:{device}" if device >= 0 else "cpu")
 
 # Global variables for feedback collection
 feedback_data = []
@@ -43,21 +46,23 @@ def set_sleep_time():
         return f"Error setting sleep time: {str(e)}"
 
 @spaces.GPU
+@torch.no_grad()
 def generate_text():
     prompt = "This AI-driven website is unique because"
     return text_generator(prompt, max_length=100, num_return_sequences=1)[0]['generated_text']
 
 @spaces.GPU
+@torch.no_grad()
 def generate_image(prompt):
-    with torch.autocast(device.type):
-        image = image_generator(prompt).images[0]
+    with torch.autocast("cuda" if device >= 0 else "cpu"):
+        image = image_generator(prompt, guidance_scale=7.5).images[0]
     return image
 
 @spaces.GPU
+@torch.no_grad()
 def analyze_sentiment(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
-    with torch.no_grad():
-        outputs = sentiment_model(**inputs)
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(f"cuda:{device}" if device >= 0 else "cpu")
+    outputs = sentiment_model(**inputs)
     probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
     return probabilities.cpu().numpy()[0]
 
@@ -77,7 +82,7 @@ def get_gpu_info():
     if torch.cuda.is_available():
         return f"Using GPU: {torch.cuda.get_device_name(0)}"
     else:
-        return "GPU not available"
+        return "GPU not available, using CPU"
 
 # Create Gradio interface
 with gr.Blocks() as demo: