Try a different model provider

app.py

@@ -4,7 +4,7 @@ from PIL import Image
 import base64 
 
 import torch
-from diffusers import StableDiffusionPipeline
+#from diffusers import StableDiffusionPipeline
 
 #from transformers import pipeline
 import gradio as gr
@@ -13,48 +13,48 @@ import gradio as gr
 hf_api_key = os.environ.get('HF_API_KEY')
 
 # Load the Stable Diffusion pipeline
-model_id = "runwayml/stable-diffusion-v1-5"
-pipe = StableDiffusionPipeline.from_pretrained(
-    model_id,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,  # Use float16 on GPU, float32 on CPU
-    use_auth_token=hf_api_key  # Required for gated model
-)
+#model_id = "runwayml/stable-diffusion-v1-5"
+#pipe = StableDiffusionPipeline.from_pretrained(
+#    model_id,
+#    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,  # Use float16 on GPU, float32 on CPU
+#    use_auth_token=hf_api_key  # Required for gated model
+#)
 
 # Move pipeline to GPU if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = pipe.to(device)
+#device = "cuda" if torch.cuda.is_available() else "cpu"
+#pipe = pipe.to(device)
 
 # Text-to-image endpoint
-#get_completion = pipeline("text-to-image", model="stable-diffusion-v1-5/stable-diffusion-v1-5")
+get_completion = pipeline("text-to-image", model="stabilityai/stable-diffusion-xl-base-1.0")
 
 # A helper function to convert the PIL image to base64,
 # so you can send it to the API
-#def base64_to_pil(img_base64):
-#    base64_decoded = base64.b64decode(img_base64)
-#    byte_stream = io.BytesIO(base64_decoded)
-#    pil_image = Image.open(byte_stream)
-#    return pil_image
+def base64_to_pil(img_base64):
+    base64_decoded = base64.b64decode(img_base64)
+    byte_stream = io.BytesIO(base64_decoded)
+    pil_image = Image.open(byte_stream)
+    return pil_image
 
-#def generate(prompt):
-#    output = get_completion(prompt)
-#    result_image = base64_to_pil(output)
-#    return result_image
+def generate(prompt):
+    output = get_completion(prompt)
+    result_image = base64_to_pil(output)
+    return result_image
 
-def generate(prompt, negative_prompt, steps, guidance, width, height):
-    # Ensure width and height are multiples of 8 (required by Stable Diffusion)
-    width = int(width) - (int(width) % 8)
-    height = int(height) - (int(height) % 8)
+#def generate(prompt, negative_prompt, steps, guidance, width, height):
+#    # Ensure width and height are multiples of 8 (required by Stable Diffusion)
+#    width = int(width) - (int(width) % 8)
+#    height = int(height) - (int(height) % 8)
     
-    # Generate image with Stable Diffusion
-    output = pipe(
-        prompt,
-        negative_prompt=negative_prompt or None,  # Handle empty negative prompt
-        num_inference_steps=int(steps),
-        guidance_scale=float(guidance),
-        width=width,
-        height=height
-    )
-    return output.images[0]  # Return the first generated image (PIL format)
+#    # Generate image with Stable Diffusion
+#    output = pipe(
+#        prompt,
+#        negative_prompt=negative_prompt or None,  # Handle empty negative prompt
+#        num_inference_steps=int(steps),
+#        guidance_scale=float(guidance),
+#        width=width,
+#        height=height
+#    )
+#    return output.images[0]  # Return the first generated image (PIL format)
 
 # Create Gradio interface
 with gr.Blocks() as demo: