Adding test

main.py

@@ -0,0 +1,39 @@
+import torch
+from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
+import gradio as gr
+
+# Check if CUDA is available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+if device == "cuda":
+    torch.cuda.empty_cache()
+
+model_id = "stabilityai/stable-diffusion-2-1"
+
+# Use appropriate dtype based on device
+if device == "cuda":
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+else:
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
+
+pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+pipe = pipe.to(device)
+
+def generate_image(prompt, width, height):
+    image = pipe(prompt, width=int(width), height=int(height)).images[0]
+    return image
+
+iface = gr.Interface(
+    fn=generate_image,
+    inputs=[
+        gr.Textbox(label="Prompt", value="a house in front of the ocean and a dog is running in the field"),
+        gr.Number(label="Width", value=1000),
+        gr.Number(label="Height", value=1000)
+    ],
+    outputs=gr.Image(type="pil"),
+    title="Stable Diffusion Image Generator"
+)
+
+if __name__ == "__main__":
+    iface.launch()

requirements.txt

@@ -0,0 +1,18 @@
+# PyTorch with CUDA support (for GPU acceleration)
+# Install from PyTorch official channel for CUDA 11.8 or 12.1
+torch>=2.0.0
+torchvision>=0.15.0
+torchaudio>=2.0.0
+
+# Diffusion models and related packages
+diffusers 
+transformers 
+accelerate 
+scipy 
+safetensors 
+
+# Optional for better performance
+xformers
+
+# Gradio for web interface
+gradio