structure

config.py

@@ -0,0 +1,15 @@
+import os
+
+from pydantic import BaseModel
+
+
+class Settings(BaseModel):
+    hf_token: str = os.environ.get("hf_token")
+    base_sd_model: str = os.environ.get("base_sd_model", "stabilityai/stable-diffusion-xl-base-1.0")
+    refiner_sd_model: str = os.environ.get("refiner_sd_model", "stabilityai/stable-diffusion-xl-refiner-1.0")
+    version: str = "0.1.0"
+    url_version: str = "v1"
+    prefix: str = "v1/unik-ml"
+
+
+settings = Settings()

main.py

@@ -1,33 +1,19 @@
-from io import BytesIO
-
-import torch
-from diffusers import DiffusionPipeline
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import StreamingResponse
+from huggingface_hub import login
 
-# load both base & refiner
-base = DiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
-)
+from config import settings
+from routers.intference import stable_diffusion
 
-base.to("cuda")
-# base.enable_model_cpu_offload()
-base.enable_attention_slicing()
-refiner = DiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-refiner-1.0",
-    text_encoder_2=base.text_encoder_2,
-    vae=base.vae,
-    torch_dtype=torch.float16,
-    use_safetensors=True,
-    variant="fp16",
-)
-refiner.to("cuda")
-# refiner.enable_model_cpu_offload()
-refiner.enable_attention_slicing()
+login(settings.hf_token)
 
-# Create a new FastAPI app instance
-app = FastAPI()
+app = FastAPI(
+    title="UNIK ML",
+    version=settings.version,
+    openapi_url=f"{settings.prefix}/openapi.json",
+    docs_url=f"{settings.prefix}/docs",
+    redoc_url=f"{settings.prefix}/redoc",
+    swagger_ui_oauth2_redirect_url=f"{settings.prefix}/docs/oauth2-redirect")
 
 app.add_middleware(
     CORSMiddleware,
@@ -43,41 +29,4 @@ async def root():
     return {"message": "UNIK ML API"}
 
 
-@app.get("/generate")
-async def generate(text: str):
-    """
-    generate image
-    """
-    # Define how many steps and what % of steps to be run on each experts (80/20) here
-    n_steps = 40
-    high_noise_frac = 0.8
-    negative = "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly. bad anatomy, disfigured, poorly drawn face, mutation, mutated, extra limb, ugly, disgusting, poorly drawn hands, missing limb, floating limbs, disconnected limbs, malformed hands, blurry, mutated hands and fingers, watermark, watermarked, oversaturated, censored, distorted hands, amputation, missing hands, obese, doubled face, double hands, two women, anime style, cartoon, toon."
-    prompt = "Designs should play with different textures and layering but stick to a monochrome palette. Think leather jackets over mesh tops, or satin draped over matte cotton. in a studio. zoomed-in. single model."
-
-    # run both experts
-    image = base(
-        prompt=prompt,
-        negative_prompt=negative,
-        num_inference_steps=n_steps,
-        denoising_end=high_noise_frac,
-        output_type="latent",
-    ).images
-    final_image = refiner(
-        prompt=prompt,
-        negative_prompt=negative,
-        num_inference_steps=n_steps,
-        denoising_start=high_noise_frac,
-        image=image,
-    ).images[0]
-
-    # buffer = BytesIO()
-    # final_image.save(buffer, format="PNG")
-    # image_bytes = buffer.getvalue()
-    #
-    # return StreamingResponse(BytesIO(image_bytes), media_type="image/png")
-    #
-    memory_stream = BytesIO()
-    final_image.save(memory_stream, format="PNG")
-    memory_stream.seek(0)
-    return StreamingResponse(memory_stream, media_type="image/png")
-
+app.include_router(stable_diffusion.router, prefix=settings.prefix, tags=["Inference", "sd"])

requirements.txt

@@ -9,5 +9,4 @@ accelerate==0.21.0
 diffusers==0.19.3
 torchvision==0.15.2
 safetensors==0.3.1
-# invisible-watermark==0.2.0
-# opencv-python-headless==4.8.0.74
+huggingface-hub==0.16.4

routers/intference/stable_diffusion.py

@@ -0,0 +1,69 @@
+# load both base & refiner
+from io import BytesIO
+
+import torch
+from diffusers import DiffusionPipeline
+from fastapi import APIRouter
+from fastapi.responses import StreamingResponse
+
+from config import settings
+
+router = APIRouter()
+
+base = DiffusionPipeline.from_pretrained(
+    settings.base_sd_model, torch_dtype=torch.float16, variant="fp16", use_safetensors=True
+)
+
+base.to("cuda")
+# base.enable_model_cpu_offload()
+base.enable_attention_slicing()
+refiner = DiffusionPipeline.from_pretrained(
+    settings.refiner_sd_model,
+    text_encoder_2=base.text_encoder_2,
+    vae=base.vae,
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    variant="fp16",
+)
+refiner.to("cuda")
+# refiner.enable_model_cpu_offload()
+refiner.enable_attention_slicing()
+
+
+@router.get("/generate")
+async def generate(prompt: str):
+    """
+    generate image
+    """
+    # Define how many steps and what % of steps to be run on each experts (80/20) here
+    n_steps = 40
+    high_noise_frac = 0.8
+    negative = "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly. bad anatomy, disfigured, poorly drawn face, mutation, mutated, extra limb, ugly, disgusting, poorly drawn hands, missing limb, floating limbs, disconnected limbs, malformed hands, blurry, mutated hands and fingers, watermark, watermarked, oversaturated, censored, distorted hands, amputation, missing hands, obese, doubled face, double hands, two women, anime style, cartoon, toon."
+    # prompt = "Designs should play with different textures and layering but stick to a monochrome palette. Think leather jackets over mesh tops, or satin draped over matte cotton. in a studio. zoomed-in. single model."
+
+    # run both experts
+    image = base(
+        prompt=prompt,
+        negative_prompt=negative,
+        num_inference_steps=n_steps,
+        denoising_end=high_noise_frac,
+        output_type="latent",
+    ).images
+    final_image = refiner(
+        prompt=prompt,
+        negative_prompt=negative,
+        num_inference_steps=n_steps,
+        denoising_start=high_noise_frac,
+        image=image,
+    ).images[0]
+
+    # buffer = BytesIO()
+    # final_image.save(buffer, format="PNG")
+    # image_bytes = buffer.getvalue()
+    #
+    # return StreamingResponse(BytesIO(image_bytes), media_type="image/png")
+    #
+    memory_stream = BytesIO()
+    final_image.save(memory_stream, format="PNG")
+    memory_stream.seek(0)
+    return StreamingResponse(memory_stream, media_type="image/png")