Update app.py

app.py

@@ -5,11 +5,16 @@ from PIL import Image
 from io import BytesIO
 from tqdm import tqdm
 import time
+import random
 
 # Defining the repository information and the trigger word
 repo = "stabilityai/stable-diffusion-xl-base-1.0"
 trigger_word = "T shirt design, TshirtDesignAF, "
 
+# Directory to save images
+output_dir = "saved_designs"
+os.makedirs(output_dir, exist_ok=True)  # Create directory if it doesn't exist
+
 def generate_images(prompt):
     print("Generating 10 unique images with prompt:", prompt)
     api_url = f"https://api-inference.huggingface.co/models/{repo}"
@@ -20,7 +25,10 @@ def generate_images(prompt):
     
     images = []
     for i in range(10):
-        full_prompt = f"{prompt} {trigger_word} unique_design_{i}"
+        # Add a unique seed to each prompt to ensure different images
+        unique_seed = random.randint(1000, 9999)
+        full_prompt = f"{prompt} {trigger_word} seed_{unique_seed}"
+
         payload = {
             "inputs": full_prompt,
             "parameters": {
@@ -40,6 +48,11 @@ def generate_images(prompt):
                 print(f"Image {i+1} generation successful!")
                 img = Image.open(BytesIO(response.content))
                 images.append(img)
+                
+                # Save the image to the output directory
+                img_filename = os.path.join(output_dir, f"{prompt.replace(' ', '_')}_design_{i+1}.png")
+                img.save(img_filename)
+                print(f"Image saved as: {img_filename}")
                 break
             elif response.status_code == 503:
                 time.sleep(1)