Update app.py

app.py

@@ -1,34 +1,26 @@
 import gradio as gr
 import torch
-import os
 from diffusers import StableDiffusion3Pipeline
-from huggingface_hub import InferenceApi
+import os
 
+# Retrieve Hugging Face API key from environment variables
 hf_api_key = os.getenv("PRODIGY_GA_02")
 if hf_api_key is None:
     raise ValueError("Hugging Face API key 'PRODIGY_GA_02' not found. Ensure it is set as a secret.")
 
-# Initialize the Hugging Face API with the restricted model and token
-inference = InferenceApi(repo_id="stabilityai/stable-diffusion-3.5-medium", token=hf_api_key)
-
-# Example inference request
-response = inference(inputs="Your input text here")
-print(response)
-
-# Load model
-model_id = "stabilityai/stable-diffusion-3.5-medium"
-pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium")
+# Load model with authentication token
+model_id = "stabilityai/stable-diffusion-3.5-large"
+pipe = StableDiffusion3Pipeline.from_pretrained(model_id, use_auth_token=hf_api_key)
 
-pipe.to("cuda")# If you have GPU access; otherwise, use "cpu"
+# Move model to the correct device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe.to(device)
 
-
-
-# Define Gradio interface
+# Define Gradio interface function
 def generate_image(prompt):
     images = pipe(prompt).images
     return images[0]
 
-
 # Create Gradio UI
 iface = gr.Interface(
     fn=generate_image,