Update app.py

app.py

@@ -33,6 +33,9 @@ unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cpu"))
 pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cpu")
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
 
+os.environ['HF_API_KEY']
+api_key = os.getenv('HF_API_KEY')
+API_URL = "https://api-inference.huggingface.co/models/CompVis/stable-diffusion-v1-4"
 
 # Function to transcribe, translate, and generate an image
 def process_audio(audio_path, generate_image):
@@ -71,11 +74,31 @@ def process_audio(audio_path, generate_image):
     return tamil_text, translation, None
 
 
+def query(payload, max_retries=5):
+    for attempt in range(max_retries):
+        response = requests.post(API_URL, headers=headers, json=payload)
+
+        if response.status_code == 503:
+            print(f"Model is still loading, retrying... Attempt {attempt + 1}/{max_retries}")
+            estimated_time = min(response.json().get("estimated_time", 60), 60)
+            time.sleep(estimated_time)
+            continue
+
+        if response.status_code != 200:
+            print(f"Error: Received status code {response.status_code}")
+            print(f"Response: {response.text}")
+            return None
+
+        return response.content
+
+    print(f"Failed to generate image after {max_retries} attempts.")
+    return None
+    
 # Function for direct prompt to image generation
 def generate_image_from_prompt(prompt):
     try:
-        img = pipe(prompt, num_inference_steps=4, guidance_scale=0).images[0]
-        return img
+        image_bytes = query({"inputs": prompt})
+        return image_bytes
     except Exception as e:
         return f"An error occurred during image generation: {str(e)}"