Update app.py

app.py

@@ -13,6 +13,8 @@ import random
 from share_btn import community_icon_html, loading_icon_html, share_js
 
 MODELS = ["stabilityai/stable-diffusion-2-1", "stabilityai/stable-diffusion-2-1-base", "stabilityai/stable-diffusion-2"]
+OUTPUT_PATH = ".output"
+os.makedirs(OUTPUT_PATH, exist_ok=True)
 
 #word_list_dataset = load_dataset("stabilityai/word-list", data_files="list.txt")
 #word_list = word_list_dataset["train"]['text']
@@ -34,7 +36,9 @@ def infer(prompt, negative, scale, model_id=MODELS[0], hf_token="", images=[]):
     client = InferenceClient(model=model_id, token=hf_token)
     for i in range(4):
         image = client.text_to_image(prompt=prompt, negative_prompt=negative, guidance_scale=scale, num_inference_steps=45, seed=gen_random_seed())
-        if image: images.append(image)
+        file_path = f"{OUTPUT_PATH}/{uuid.uuid4()}.jpg"
+        image.save(file_path)
+        if image: images.append(file_path)
         yield images
     #payload = {'prompt': prompt, 'negative_prompt': negative, 'guidance_scale': scale}
     #images_request = requests.post(url, json = payload)
@@ -56,7 +60,9 @@ def infer_advance(prompt, negative, scale, width, height, steps, seed, samples,
     for i in range(samples):
         image = client.text_to_image(prompt=prompt, negative_prompt=negative, width=width, height=height,
                                      guidance_scale=scale, num_inference_steps=steps, seed=seed if i==0 else gen_random_seed())
-        if image: images.append(image)
+        file_path = f"{OUTPUT_PATH}/{uuid.uuid4()}.jpg"
+        image.save(file_path)
+        if image: images.append(file_path)
         yield images 
 
 css = """