Update Working_TTI_GPU.py

Working_TTI_GPU.py

@@ -0,0 +1,24 @@
+from diffusers import DiffusionPipeline
+import torch
+from PIL import Image
+
+# Load the model
+pipe = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    variant="fp16"
+)
+pipe.to("cuda")
+
+# Define the prompt
+prompt = "An astronaut riding a red horse"
+
+# Generate the image
+result = pipe(prompt=prompt)
+image = result.images[0]
+
+# Save the image
+image.save("generated_image.png")
+
+print("Image saved as generated_image.png")