import gradio as gr import os import zipfile import requests from PIL import Image from io import BytesIO from datetime import datetime # ---------- CONFIG ---------- BYTEZ_API_KEY = os.getenv("BYTEZ_API_KEY") # put your real key here BYTEZ_ENDPOINT = "https://api.bytez.com/models/v2/google/imagen-4.0-fast-generate-001" BASE_DIR = os.path.abspath(os.getcwd()) HISTORY_ROOT = os.path.join(BASE_DIR, "user_image_history") os.makedirs(HISTORY_ROOT, exist_ok=True) # ---------- HELPERS ---------- def get_user_dir(username: str) -> str: """Return (and create) per-user history directory.""" if not username: username = "_anonymous" safe_username = "".join(c for c in username if c.isalnum() or c in ("_", "-")) or "_anonymous" user_dir = os.path.join(HISTORY_ROOT, safe_username) os.makedirs(user_dir, exist_ok=True) return user_dir def get_history(username: str): """Return sorted list of image paths for this user.""" user_dir = get_user_dir(username) history = sorted([ os.path.join(user_dir, f) for f in os.listdir(user_dir) if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp")) ]) # First output for gallery, second output unused but kept to match callbacks return history, history def generate_image_txt2img(username: str, prompt: str): """Call Bytez Imagen API with text prompt, download image, store it.""" user_dir = get_user_dir(username) if not prompt.strip(): # No prompt, just return history history_txt, history_img = get_history(username) return None, history_txt, None headers = { "Authorization": BYTEZ_API_KEY, "Content-Type": "application/json", } payload = { "text": prompt } try: resp = requests.post(BYTEZ_ENDPOINT, json=payload, headers=headers) resp.raise_for_status() data = resp.json() except Exception as e: history_txt, history_img = get_history(username) # Show error in label, no image return gr.update(label=f"Error calling API: {e}", value=None), history_txt, None # Bytez sample response structure (from your Postman test): # { # "error": null, # "output": "https://cdn.bytez.com/model/output/google/imagen-4.0-fast-generate-001/....png", # "provider": { "generatedImages": [ { "image": { "imageBytes": "..." } } ] } # } img = None # Prefer using the public URL in "output" try: img_url = data.get("output") if img_url: img_resp = requests.get(img_url) img_resp.raise_for_status() img = Image.open(BytesIO(img_resp.content)).convert("RGB") else: raise ValueError("No 'output' URL in response.") except Exception as e: history_txt, history_img = get_history(username) return gr.update(label=f"Error downloading image: {e}", value=None), history_txt, None # Save image to user's history if img: timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f") filename = os.path.join(user_dir, f"{timestamp}.png") img.save(filename) history_txt, history_img = get_history(username) return img, history_txt, None def select_image(image_path): """Return selected image path for popup display.""" return image_path def delete_selected_image(username: str, selected_image_path: str): """Delete the selected image from the user's directory.""" if selected_image_path: user_dir = get_user_dir(username) # Only delete if path is inside this user's folder try: if os.path.commonpath( [os.path.abspath(selected_image_path), user_dir] ) == os.path.abspath(user_dir): if os.path.exists(selected_image_path): os.remove(selected_image_path) except Exception: pass history_txt, history_img = get_history(username) return None, history_txt, None def download_history(username: str): """Create a ZIP of the user's history folder and return its path.""" user_dir = get_user_dir(username) zip_path = os.path.join(user_dir, "image_history.zip") with zipfile.ZipFile(zip_path, "w") as zipf: for root, dirs, files in os.walk(user_dir): for file in files: if file == "image_history.zip": continue full_path = os.path.join(root, file) arcname = os.path.relpath(full_path, user_dir) zipf.write(full_path, arcname=arcname) return zip_path # ---------- GRADIO UI ---------- with gr.Blocks() as demo: gr.Markdown("# Bytez Imagen 4.0 Fast – Txt2Img (per‑user history)") with gr.Row(): with gr.Column(scale=2): username = gr.Textbox( label="Username (for personal image history)", placeholder="Enter a username, e.g. alice", value="" ) prompt_txt2img = gr.Textbox( label="Prompt", placeholder="Describe the image...", lines=5 ) generate_btn_txt2img = gr.Button("Generate Image") delete_btn_txt2img = gr.Button("Delete Selected Image") download_btn_txt2img = gr.Button("Download My History ZIP") with gr.Column(scale=3): output_image_txt2img = gr.Image( label="Generated Image", interactive=False, type="pil" ) with gr.Accordion("My History", open=False): gallery_txt2img = gr.Gallery( label="Image History", show_label=False ) popup_image_txt2img = gr.Image( label="Full Image View", visible=False, type="filepath" ) # Reload history when username changes username.change( fn=get_history, inputs=[username], outputs=[gallery_txt2img, gallery_txt2img] ) # Load initial history for default username demo.load( fn=get_history, inputs=[username], outputs=[gallery_txt2img, gallery_txt2img] ) # Generate image generate_btn_txt2img.click( fn=generate_image_txt2img, inputs=[username, prompt_txt2img], outputs=[output_image_txt2img, gallery_txt2img, popup_image_txt2img] ) # Select image from gallery gallery_txt2img.select( fn=select_image, inputs=gallery_txt2img, outputs=popup_image_txt2img ) # Delete selected image delete_btn_txt2img.click( fn=delete_selected_image, inputs=[username, popup_image_txt2img], outputs=[output_image_txt2img, gallery_txt2img, popup_image_txt2img] ) # Download ZIP download_btn_txt2img.click( fn=download_history, inputs=[username], outputs=[gr.File(label="Download ZIP")] ) demo.launch()