Upload 2 files

demo.py

@@ -0,0 +1,53 @@
+from diffusers import (
+    StableDiffusionPipeline,
+    PNDMScheduler,
+)
+from diffusers.models import AutoencoderKL
+import os
+import torch
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+vae = AutoencoderKL.from_pretrained(
+    "stabilityai/sd-vae-ft-ema", torch_dtype=torch.float32
+)
+common_config = {"beta_start": 0.00085, "beta_end": 0.012, "beta_schedule": "scaled_linear"}
+SCHEDULER = PNDMScheduler(**common_config, skip_prk_steps=True, steps_offset=1,)
+HF_API_TOKEN = os.getenv("HF_API_TOKEN")
+shared_pipe_kwargs = dict(
+    vae=vae,
+    torch_dtype=torch.float32,
+    revision="fp16",
+    use_auth_token=HF_API_TOKEN,
+    scheduler=SCHEDULER,
+)
+
+
+base_sd_pipe = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5", **shared_pipe_kwargs
+).to(device)
+
+pai_model_dir = '.'
+
+playground_v1_model_dir = os.path.join(
+    pai_model_dir, "snapshots/36c9e19103f6b897886fb019ebc4d8e86b566032"
+)
+
+pgv1_shared_pipe_kwargs = dict(
+    vae=vae,
+    torch_dtype=torch.float32,
+    tokenizer=base_sd_pipe.tokenizer,
+    feature_extractor=base_sd_pipe.feature_extractor,
+    text_encoder=base_sd_pipe.text_encoder,
+    scheduler=SCHEDULER,
+    local_files_only=True, 
+)
+
+pgv1_pipe = StableDiffusionPipeline.from_pretrained(
+    playground_v1_model_dir,
+    **pgv1_shared_pipe_kwargs,
+).to(device)
+
+img = pgv1_pipe("Frog", num_inference_steps=30)
+
+img[0][0].save('frog.png')

requirements.txt

@@ -0,0 +1,2 @@
+diffusers==0.22.3
+torch==2.1.0