Update app.py

app.py

@@ -202,6 +202,13 @@ with gr.Blocks() as demo:
         polish_button.click(get_image, inputs=[height, width, seed, prompt, cfg, timesteps, stage_a_tiled], outputs=output_img)           
         polish_button.click(clear_image, inputs=[], outputs=output_img)
    
+
+
+def download_and_save_model(model_name_or_path, save_directory):
+    from transformers import AutoModel
+    model = AutoModel.from_pretrained(model_name_or_path)
+    model.save_pretrained(save_directory)
+    print(f"Model saved to {save_directory}", model_name_or_path)
 def download_model():
     urls = [
         'https://huggingface.co/stabilityai/StableWurst/resolve/main/stage_a.safetensors',
@@ -212,13 +219,7 @@ def download_model():
         'https://huggingface.co/roubaofeipi/UltraPixel/blob/main/ultrapixel_t2i.safetensors'  
     ]
     for file_url in urls:
-        
-        response = requests.get(file_url)
-        if response.status_code == 200:
-            file_path = os.path.join('models', file_url.split('/')[-1])
-            with open(file_path, "wb") as f:
-                f.write(response.content)
-            print('download successful! ', file_path)
+        download_and_save_model(file_url, 'models')        
 if __name__ == "__main__":
    
     args = parse_args()