Update app.py

app.py

@@ -7,7 +7,7 @@ import numpy as np
 from PIL import Image
 import spaces
 import torch
-from diffusers import StableDiffusion3Pipeline,StableDiffusion3Img2ImgPipeline, DPMSolverMultistepScheduler, AutoencoderKL
+from diffusers import StableDiffusion3Pipeline, DPMSolverMultistepScheduler, AutoencoderKL, StableDiffusion3Img2ImgPipeline
 from huggingface_hub import snapshot_download
 
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
@@ -33,8 +33,11 @@ ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-pipe = StableDiffusion3Pipeline.from_pretrained(model_path, torch_dtype=torch.float16)
-img2img_pipe = StableDiffusion3Img2ImgPipeline.from_pretrained(model_path, torch_dtype=torch.float16)
+def load_pipeline(pipeline_type):
+    if pipeline_type == "text2img":
+        return StableDiffusion3Pipeline.from_pretrained(model_path, torch_dtype=torch.float16)
+    elif pipeline_type == "img2img":
+        return StableDiffusion3Img2ImgPipeline.from_pretrained(model_path, torch_dtype=torch.float16)
 
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
@@ -63,6 +66,7 @@ def generate(
     use_resolution_binning: bool = True,
     progress=gr.Progress(track_tqdm=True),
 ):
+    pipe = load_pipeline("text2img")
     pipe.to(device)
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator().manual_seed(seed)
@@ -100,7 +104,8 @@ def img2img_generate(
     use_resolution_binning: bool = True,
     progress=gr.Progress(track_tqdm=True),
 ):
-    img2img_pipe.to(device)
+    pipe = load_pipeline("img2img")
+    pipe.to(device)
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator().manual_seed(seed)
     
@@ -109,7 +114,7 @@ def img2img_generate(
     
     init_image = init_image.resize((768, 768))
     
-    output = img2img_pipe(
+    output = pipe(
         prompt=prompt,
         image=init_image,
         negative_prompt=negative_prompt,
@@ -177,7 +182,7 @@ with gr.Blocks(css=css) as demo:
             )
         seed = gr.Slider(
             label="Seed",
-            min=0,
+            minimum=0,
             maximum=MAX_SEED,
             step=1,
             value=0,
@@ -185,14 +190,14 @@ with gr.Blocks(css=css) as demo:
 
         steps = gr.Slider(
             label="Steps",
-            min=0,
+            minimum=0,
             maximum=60,
             step=1,
             value=25,
         )
         number_image = gr.Slider(
             label="Number of Images",
-            min=1,
+            minimum=1,
             maximum=4,
             step=1,
             value=1,
@@ -201,14 +206,14 @@ with gr.Blocks(css=css) as demo:
         with gr.Row(visible=True):
             width = gr.Slider(
                 label="Width",
-                min=256,
+                minimum=256,
                 maximum=MAX_IMAGE_SIZE,
                 step=32,
                 value=1024,
             )
             height = gr.Slider(
                 label="Height",
-                min=256,
+                minimum=256,
                 maximum=MAX_IMAGE_SIZE,
                 step=32,
                 value=1024,
@@ -216,7 +221,7 @@ with gr.Blocks(css=css) as demo:
         with gr.Row():
             guidance_scale = gr.Slider(
                 label="Guidance Scale",
-                min=0.1,
+                minimum=0.1,
                 maximum=10,
                 step=0.1,
                 value=7.0,