KarthikAI commited on
Commit
72c3233
·
verified ·
1 Parent(s): cdee956

Update utils.py

Browse files
Files changed (1) hide show
  1. utils.py +22 -7
utils.py CHANGED
@@ -1,25 +1,40 @@
 
 
 
 
 
 
1
  import torch
2
  from diffusers import StableDiffusionImg2ImgPipeline
3
  from PIL import Image
4
 
 
5
  # --- Place any download or path setup here ---
6
  MODEL_ID = "runwayml/stable-diffusion-v1-5" # Can swap for custom path if using IP-Adapter
7
  DEVICE = "cpu"
8
- MODEL_CACHE = "./data"
9
 
10
  # (Optional) Download IP-Adapter weights and patch pipeline if desired
11
 
 
 
 
 
 
 
 
 
12
  def generate_sticker(input_image, prompt):
13
  """
14
  Given a user image and a prompt, generates a sticker/emoji-style portrait.
15
  """
16
  # Load the model (download if not present)
17
- pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
18
- MODEL_ID,
19
- torch_dtype=torch.float32,
20
- cache_dir=MODEL_CACHE,
21
- safety_checker=None, # Disable for demo/testing
22
- ).to(DEVICE)
23
 
24
  # Preprocess the image (resize, etc)
25
  init_image = input_image.convert("RGB").resize((512, 512))
 
1
+
2
+ import os
3
+ # Set Hugging Face cache dir to a safe writable location (works in Spaces & Docker)
4
+ os.environ["TRANSFORMERS_CACHE"] = "/workspace/.cache/huggingface"
5
+ os.makedirs("/workspace/.cache/huggingface", exist_ok=True)
6
+
7
  import torch
8
  from diffusers import StableDiffusionImg2ImgPipeline
9
  from PIL import Image
10
 
11
+
12
  # --- Place any download or path setup here ---
13
  MODEL_ID = "runwayml/stable-diffusion-v1-5" # Can swap for custom path if using IP-Adapter
14
  DEVICE = "cpu"
15
+ MODEL_CACHE = "/workspace/.cache/huggingface"
16
 
17
  # (Optional) Download IP-Adapter weights and patch pipeline if desired
18
 
19
+ # Load the model ONCE at startup, not per request!
20
+ pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
21
+ MODEL_ID,
22
+ torch_dtype=torch.float32,
23
+ cache_dir=MODEL_CACHE,
24
+ safety_checker=None, # Disable for demo/testing; enable in prod
25
+ ).to(DEVICE)
26
+
27
  def generate_sticker(input_image, prompt):
28
  """
29
  Given a user image and a prompt, generates a sticker/emoji-style portrait.
30
  """
31
  # Load the model (download if not present)
32
+ # pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
33
+ # MODEL_ID,
34
+ # torch_dtype=torch.float32,
35
+ # cache_dir=MODEL_CACHE,
36
+ # safety_checker=None, # Disable for demo/testing
37
+ # ).to(DEVICE)
38
 
39
  # Preprocess the image (resize, etc)
40
  init_image = input_image.convert("RGB").resize((512, 512))