Update app.py

app.py

@@ -32,8 +32,11 @@ vlm_processor = PaliGemmaProcessor.from_pretrained("gokaygokay/sd3-long-captione
 enhancer_medium = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device=device)
 enhancer_long = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device=device)
 
-# SD3
-sd3_pipe = StableDiffusion3Pipeline.from_pretrained(model_path, torch_dtype=dtype).to(device)
+def load_pipeline(pipeline_type):
+    if pipeline_type == "text2img":
+        return StableDiffusion3Pipeline.from_pretrained(model_path, torch_dtype=dtype).to(device)
+    elif pipeline_type == "img2img":
+        return StableDiffusion3Img2ImgPipeline.from_pretrained(model_path, torch_dtype=dtype).to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1344
@@ -90,8 +93,10 @@ def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height,
         seed = random.randint(0, MAX_SEED)
     
     generator = torch.Generator().manual_seed(seed)
+
+    pipe = load_pipeline("text2img")
     
-    image = sd3_pipe(
+    image = pipe(
         prompt=prompt, 
         negative_prompt=negative_prompt,
         guidance_scale=guidance_scale, 
@@ -143,7 +148,7 @@ def img2img_generate(
     
     generator = torch.Generator().manual_seed(seed)
     
-    img2img_pipe = StableDiffusion3Img2ImgPipeline.from_pretrained(model_path, torch_dtype=dtype).to(device)
+    img2img_pipe = load_pipeline("img2img")
     
     init_image = init_image.resize((768, 768))