Use library manually

app.py

@@ -3,29 +3,57 @@ import io
 from PIL import Image
 import base64 
 
+import torch
+from diffusers import StableDiffusionPipeline
+
 from transformers import pipeline
 import gradio as gr
 
-hf_api_key = os.environ['HF_API_KEY']
+# Set Hugging Face API (needed for gated models)
+hf_api_key = os.environ.get('HF_API_KEY')
 
-# Text-to-image endpoint
-get_completion = pipeline("text-to-image", model="stable-diffusion-v1-5/stable-diffusion-v1-5")
+# Load the Stable Diffusion pipeline
+model_id = "runwayml/stable-diffusion-v1-5"
+pipe = StableDiffusionPipeline.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16,  # Use float16 for better performance on GPU
+    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)
+
+# Text-to-image endpoint
+#get_completion = pipeline("text-to-image", model="stable-diffusion-v1-5/stable-diffusion-v1-5")
 
 # 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):
+    # Generate image with Stable Diffusion
+    output = pipe(
+        prompt,
+        negative_prompt=negative_prompt,
+        num_inference_steps=int(steps),
+        guidance_scale=float(guidance),
+        width=int(width),
+        height=int(height)
+    )
+    return output.images[0]  # Return the first generated image (PIL format)
 
 
+# Create Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Image Generation with Stable Diffusion")
     prompt = gr.Textbox(label="Your prompt")
@@ -45,6 +73,7 @@ with gr.Blocks() as demo:
 
     btn.click(fn=generate, inputs=[prompt,negative_prompt,steps,guidance,width,height], outputs=[output])
 
+# Launch the app
 demo.launch(
     share=True,
     #server_port=int(os.environ['PORT3'])