Update main.py

main.py

@@ -1,7 +1,8 @@
 run_api = False
 is_ssd = False
 is_sdxl = False
-is_sdxl_turbo=True
+is_sdxl_turbo=False
+use_request=True
 import os
 # Use GPU
 gpu_info = os.popen("nvidia-smi").read()
@@ -60,6 +61,9 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1024"))
 SECRET_TOKEN = os.getenv("SECRET_TOKEN", "default_secret")
 
+API_TOKEN = os.environ.get("HF_READ_TOKEN")
+headers = {"Authorization": f"Bearer {API_TOKEN}"}
+
 # Uncomment the following line if you are using PyTorch 1.10 or later
 # os.environ["TORCH_USE_CUDA_DSA"] = "1"
 
@@ -129,6 +133,7 @@ if is_sdxl_turbo:
     use_cuda=is_gpu
     pipe = load_pipeline(use_cuda)
 
+
 def generate(
     prompt: str,
     negative_prompt: str = "",
@@ -146,16 +151,31 @@ def generate(
 
     generator = torch.Generator().manual_seed(seed)
 
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        width=width,
-        height=height,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        generator=generator,
-        output_type="pil",
-    ).images[0]
+
+    if not use_request:
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+            output_type="pil",
+        ).images[0]
+
+    else:
+        API_URL = "https://api-inference.huggingface.co/models/segmind/SSD-1B"
+        payload = {
+            "inputs": prompt ,
+            "is_negative": negative_prompt,
+            "steps": num_inference_steps,
+            "cfg_scale": guidance_scale,
+            "seed": seed if seed is not None else random.randint(-1, 2147483647)
+        }
+
+        image_bytes = requests.post(API_URL, headers=headers, json=payload).content
+        image = Image.open(io.BytesIO(image_bytes))    
     return image