Spaces:
Running
Running
| import os | |
| import io | |
| import requests | |
| from PIL import Image | |
| import gradio as gr | |
| # ------------------------------------------------------------------------------ | |
| # Make sure you export your API token before running this script: | |
| # export CHUTES_API_TOKEN="your_real_token_here" | |
| # ------------------------------------------------------------------------------ | |
| API_TOKEN = os.getenv("CHUTES_API_TOKEN") | |
| if API_TOKEN is None: | |
| raise ValueError("Please set the environment variable CHUTES_API_TOKEN before running.") | |
| def generate_image(prompt: str, | |
| resolution: str, | |
| guidance_scale: float, | |
| num_inference_steps: int, | |
| shift: int) -> Image.Image: | |
| """ | |
| Calls the HiDream (Chutes) API to generate an image from the given prompt | |
| and parameters. Assumes the API responds with raw image bytes (e.g. PNG or JPEG) | |
| rather than JSON. Returns a PIL.Image on success. Raises a RuntimeError on failure. | |
| """ | |
| headers = { | |
| "Authorization": f"Bearer {API_TOKEN}", | |
| "Content-Type": "application/json" | |
| } | |
| payload = { | |
| "seed": None, | |
| "shift": shift, | |
| "prompt": prompt, | |
| "resolution": resolution, | |
| "guidance_scale": guidance_scale, | |
| "num_inference_steps": num_inference_steps | |
| } | |
| response = requests.post( | |
| "https://chutes-hidream.chutes.ai/generate", | |
| headers=headers, | |
| json=payload | |
| ) | |
| # If the status code isn’t 200, raise immediately | |
| if response.status_code != 200: | |
| # Try to show any text body if present | |
| raw_text = response.text.strip() | |
| raise RuntimeError(f"API returned HTTP {response.status_code}:\n{raw_text}") | |
| # At this point we assume response.content is raw image bytes. | |
| try: | |
| img = Image.open(io.BytesIO(response.content)).convert("RGB") | |
| return img | |
| except Exception as e: | |
| # If decoding as an image fails, include a hex preview of the first few bytes | |
| prefix = response.content[:16].hex() | |
| raise RuntimeError( | |
| f"Failed to decode image from response content (first 16 bytes: {prefix}...):\n{str(e)}" | |
| ) | |
| # Build the Gradio interface | |
| with gr.Blocks(title="HiDream Unlimited") as demo: | |
| gr.Markdown( | |
| "## HiDream Unlimited\n\n" | |
| "Generate unlimited AI-driven images powered by HiDream.\n\n" | |
| "Enter your prompt and tweak the parameters, then click **Generate**." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| prompt_in = gr.Textbox( | |
| label="Prompt", | |
| placeholder="e.g. a serene sunset over a coral reef, digital painting", | |
| lines=2 | |
| ) | |
| resolution_in = gr.Dropdown( | |
| choices=["512x512", "768x768", "1024x1024"], | |
| value="1024x1024", | |
| label="Resolution" | |
| ) | |
| guidance_in = gr.Slider( | |
| minimum=0.0, maximum=20.0, step=0.5, | |
| value=5.0, label="Guidance Scale" | |
| ) | |
| steps_in = gr.Slider( | |
| minimum=1, maximum=100, step=1, | |
| value=50, label="Num Inference Steps" | |
| ) | |
| shift_in = gr.Slider( | |
| minimum=0, maximum=10, step=1, | |
| value=3, label="Shift" | |
| ) | |
| generate_btn = gr.Button("Generate Image", variant="primary") | |
| with gr.Column(scale=1): | |
| output_img = gr.Image(label="Generated Image", interactive=False) | |
| generate_btn.click( | |
| fn=generate_image, | |
| inputs=[prompt_in, resolution_in, guidance_in, steps_in, shift_in], | |
| outputs=output_img | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |