Mr-Vicky-01 commited on
Commit
44ad93d
·
verified ·
1 Parent(s): 63b3ced

Update image_engine.py

Browse files
Files changed (1) hide show
  1. image_engine.py +6 -2
image_engine.py CHANGED
@@ -1,13 +1,17 @@
1
  from PIL import Image
2
  import requests
 
3
  import io
4
  import os
5
 
6
  # Define the ImageGenerator class for Hugging Face API
7
  class ImageGenerator:
8
  def __init__(self) -> None:
9
- # self.api_url = "https://api-inference.huggingface.co/models/stablediffusionapi/mklan-xxx-nsfw-pony"
10
- self.api_url = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large"
 
 
 
11
  self.headers = {"Authorization": f"Bearer {os.getenv('HUGGING_FACE')}"}
12
 
13
  def generate_image(self, prompt):
 
1
  from PIL import Image
2
  import requests
3
+ import random
4
  import io
5
  import os
6
 
7
  # Define the ImageGenerator class for Hugging Face API
8
  class ImageGenerator:
9
  def __init__(self) -> None:
10
+ self.all_urls = [
11
+ "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell",
12
+ "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large"
13
+ ]
14
+ self.api_url = random.choice(self.all_urls)
15
  self.headers = {"Authorization": f"Bearer {os.getenv('HUGGING_FACE')}"}
16
 
17
  def generate_image(self, prompt):