base

app.py

@@ -0,0 +1,29 @@
+import gradio as gr
+from PIL import Image
+from io import BytesIO
+import torch
+from torch.cuda.amp import autocast
+from diffusers import DiffusionPipeline
+import os
+
+model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32, use_safetensors=True, variant="fp16")
+pipe = pipe.to(device)
+
+def generate_image(prompt: str, secret: str):
+    if secret != os.getenv("SECRET_KEY"):
+        return "Not Authorized"
+
+    with autocast(device):
+        image = pipe(prompt, guidance_scale=8.5).images[0]
+
+    buffer = BytesIO()
+    image.save(buffer, format="PNG") 
+    buffer.seek(0)
+
+    return buffer
+
+iface = gr.Interface(fn=generate_image, inputs=[gr.inputs.Textbox(label="prompt"), gr.inputs.Textbox(label="secret")], outputs="bytes")
+iface.launch()

requirements.txt

@@ -0,0 +1,16 @@
+# fastapi
+# uvicorn
+# torch
+# transformers
+# diffusers
+
+gradio
+torch
+diffusers
+transformers
+accelerate
+uvicorn
+datasets
+torch
+soundfile
+sentencepiece