Update app.py

app.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-#patch 0.01
+#patch 0.01yle(collage_style, prompt, negative_prompt)
 import os
 import random
 import uuid
@@ -10,8 +10,8 @@ import numpy as np
 from PIL import Image
 import spaces
 import torch
-from diffusers import StableDiffusion3Pipeline, FlowMatchEulerDiscreteScheduler
-from typing import Tuple
+from diffusers import StableDiffusion3Pipeline, DPMSolverMultistepScheduler, AutoencoderKL
+from huggingface_hub import snapshot_download
 
 # BaseConditions
 bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))
@@ -141,13 +141,20 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
 
+# Download the model to a local directory
+model_path = snapshot_download(
+    repo_id="stabilityai/stable-diffusion-3-medium", 
+    revision="refs/pr/26",
+    repo_type="model", 
+    ignore_patterns=["*.md", "*.gitattributes"],
+    local_dir="stable-diffusion-3-medium",
+    token=huggingface_token,
+)
+
 if torch.cuda.is_available():
-    repo = "stabilityai/stable-diffusion-3-medium"  # You might want to replace with your specific repo if different
     pipe = StableDiffusion3Pipeline.from_pretrained(
-        repo,
+        model_path,
         torch_dtype=torch.float16,
-        revision="refs/pr/26",
-        use_auth_token=huggingface_token
     ).to(device)
 
     if ENABLE_CPU_OFFLOAD:
@@ -214,13 +221,12 @@ def generate(
     options = {
         "prompt": prompt,
         "negative_prompt": negative_prompt,
-        "width": width,
-        "height": height,
         "guidance_scale": guidance_scale,
-        "num_inference_steps": 20,
         "generator": generator,
+        "num_inference_steps": 25,
+        "height": height,
+        "width": width,
         "num_images_per_prompt": num_images,
-        "use_resolution_binning": use_resolution_binning,
         "output_type": "pil",
     }