Update app.py

app.py

@@ -1,34 +1,18 @@
 import gradio as gr
 import torch
-#from diffusers import StableDiffusionPipeline
-
-
 from diffusers import DiffusionPipeline
 
-
-
-
+# Define the image generation function
 def image_generation(prompt):
     # Check if GPU is available
     device = "cuda" if torch.cuda.is_available() else "cpu"
     
-    # Load the Stable Diffusion 3 pipeline
+    # Load the Stable Diffusion 2.1 pipeline (as you're using DiffusionPipeline now)
     pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1",
-        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-        text_encoder_3=None,
-        tokenizer_3=None)
+                                                 torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+                                                 text_encoder_3=None,
+                                                 tokenizer_3=None)
 
-    #pipeline = StableDiffusionPipeline.from_pretrained(
-       # "stabilityai/stable-diffusion-3-medium-diffusers",
-       # torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-       # use_auth_token=HUGGINGFACE_TOKEN,  # Use the Hugging Face token for authentication
-       # text_encoder_3=None,
-       # tokenizer_3=None
-   # )
-    
-    # Enable efficient model execution
-    pipeline.enable_model_cpu_offload()
-    
     # Generate an image based on the prompt
     image = pipeline(
         prompt=prompt,