segmindRT

app.py

@@ -1,10 +1,7 @@
 import numpy as np
 import PIL.Image
 import torch
-from typing import List
-from diffusers.utils import numpy_to_pil
-from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
-from diffusers.pipelines.wuerstchen import DEFAULT_STAGE_C_TIMESTEPS
+from diffusers import LCMScheduler, AutoPipelineForText2Image
 from fastapi import FastAPI
 import uvicorn
 from fastapi.middleware.cors import CORSMiddleware
@@ -20,11 +17,8 @@ from fastapi.middleware.cors import CORSMiddleware
 
 logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))
 
-MAX_SEED = np.iinfo(np.int32).max
-USE_TORCH_COMPILE = os.environ.get("USE_TORCH_COMPILE", "0") == "1"
 SPACE_ID = os.environ.get("SPACE_ID", "")
 DEV = os.environ.get("DEV", "0") == "1"
-os.environ['TOKENIZERS_PARALLELISM'] = 'false'
 
 DB_PATH = Path("/data/cache") if SPACE_ID else Path("./cache")
 IMGS_PATH = DB_PATH / "imgs"
@@ -33,64 +27,39 @@ IMGS_PATH.mkdir(exist_ok=True, parents=True)
 
 database = Database(DB_PATH)
 
+
+model_id = "segmind/Segmind-Vega"
+adapter_id = "segmind/Segmind-VegaRT"
+
 dtype = torch.bfloat16
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 if torch.cuda.is_available():
-    prior_pipeline = StableCascadePriorPipeline.from_pretrained(
-        "stabilityai/stable-cascade-prior", torch_dtype=dtype
-    )  # .to(device)
-    decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained(
-        "stabilityai/stable-cascade", torch_dtype=dtype
-    )  # .to(device)
-    prior_pipeline.to(device)
-    decoder_pipeline.to(device)
-
-    if USE_TORCH_COMPILE:
-        prior_pipeline.prior = torch.compile(
-            prior_pipeline.prior, mode="reduce-overhead", fullgraph=True
-        )
-        decoder_pipeline.decoder = torch.compile(
-            decoder_pipeline.decoder, mode="max-autotune", fullgraph=True
-        )
+    pipe = AutoPipelineForText2Image.from_pretrained(
+        model_id, torch_dtype=torch.float16, variant="fp16"
+    )
+    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+    pipe.to("cuda")
+    pipe.load_lora_weights(adapter_id)
+    pipe.fuse_lora()
 
 
 def generate(
     prompt: str,
     negative_prompt: str = "",
     seed: int = 0,
-    width: int = 1024,
-    height: int = 1024,
-    prior_num_inference_steps: int = 20,
-    prior_guidance_scale: float = 4.0,
-    decoder_num_inference_steps: int = 10,
-    decoder_guidance_scale: float = 0.0,
-    num_images_per_prompt: int = 2,
 ) -> PIL.Image.Image:
 
     generator = torch.Generator().manual_seed(seed)
-    prior_output = prior_pipeline(
-        prompt=prompt,
-        height=height,
-        width=width,
-        num_inference_steps=prior_num_inference_steps,
-        timesteps=DEFAULT_STAGE_C_TIMESTEPS,
-        negative_prompt=negative_prompt,
-        guidance_scale=prior_guidance_scale,
-        num_images_per_prompt=num_images_per_prompt,
-        generator=generator,
-    )
-    decoder_output = decoder_pipeline(
-        image_embeddings=prior_output.image_embeddings,
+
+    image = pipe(
         prompt=prompt,
-        num_inference_steps=decoder_num_inference_steps,
-        # timesteps=decoder_timesteps,
-        guidance_scale=decoder_guidance_scale,
         negative_prompt=negative_prompt,
         generator=generator,
-        output_type="pil",
-    ).images
+        num_inference_steps=4,
+        guidance_scale=0,
+    ).images[0]
 
-    return decoder_output[0]
+    return image
 
 
 app = FastAPI()
@@ -120,7 +89,9 @@ app.add_middleware(
 
 
 @app.get("/image")
-async def generate_image(prompt: str, negative_prompt: str = "", seed: int = 2134213213):
+async def generate_image(
+    prompt: str, negative_prompt: str = "", seed: int = 2134213213
+):
     cached_img = database.check(prompt, negative_prompt, seed)
     if cached_img:
         logging.info(f"Image found in cache: {cached_img[0]}")