Spaces:
Running
Running
| import gradio as gr | |
| from huggingface_hub import InferenceClient | |
| import os | |
| # --- Configuration --- | |
| # Use environment variable for token, fallback to default if not set | |
| # For Spaces, set the HF_TOKEN secret | |
| API_TOKEN = os.getenv("HF_TOKEN", None) | |
| MODEL = "HuggingFaceH4/zephyr-7b-beta" # Or choose another suitable model | |
| # --- Initialize Inference Client --- | |
| try: | |
| print(f"Attempting to initialize Inference Client for model: {MODEL}") | |
| if API_TOKEN: | |
| print("Using HF Token found in environment.") | |
| client = InferenceClient(model=MODEL, token=API_TOKEN) | |
| else: | |
| print("HF Token not found. Running without token (may lead to rate limits).") | |
| client = InferenceClient(model=MODEL) | |
| # Optional: Add a quick health check if needed, though client init usually suffices | |
| # client.get_model_status(MODEL) # Example check, might raise if model invalid | |
| print("Inference Client initialized successfully.") | |
| except Exception as e: | |
| print(f"Error initializing Inference Client: {e}") | |
| # Provide a more informative error in the Gradio interface | |
| raise gr.Error(f"Failed to initialize the AI model client for '{MODEL}'. Please check the model name and ensure network connectivity. If running locally without a token, you might hit rate limits. If using HF Spaces, ensure the HF_TOKEN secret is set correctly. Original Error: {e}") | |
| # --- Core Code Generation Function --- | |
| def generate_code( | |
| prompt: str, | |
| backend_choice: str, | |
| file_structure: str, | |
| max_tokens: int, | |
| temperature: float, | |
| top_p: float, | |
| ): | |
| """ | |
| Generates website code based on user prompt and choices. | |
| Yields the code token by token for live updates. | |
| """ | |
| print(f"--- Generating Code ---") | |
| print(f"Prompt: {prompt[:100]}...") # Log truncated prompt | |
| print(f"Backend Context: {backend_choice}") | |
| print(f"File Structure: {file_structure}") | |
| print(f"Settings: Max Tokens={max_tokens}, Temp={temperature}, Top-P={top_p}") | |
| # --- System Message (Internal) --- | |
| # Guides the AI's behavior. Not user-editable in the UI. | |
| # Refined system prompt for clarity and stricter output formatting | |
| system_message = ( | |
| "You are an expert frontend web developer AI. Your task is to generate HTML, CSS, and potentially JavaScript code " | |
| "for a website based ONLY on the user's description. Adhere strictly to the following constraints:\n" | |
| "1. **Output ONLY Code:** Generate only the raw code for the requested files (`index.html`, `style.css`, `script.js`). Do NOT include any introductory text, explanations, apologies, markdown formatting (like ```html), or closing remarks. Your response must start *immediately* with the code (e.g., `<!DOCTYPE html>` or `<!-- index.html -->`).\n" | |
| "2. **index.html is Mandatory:** ALWAYS generate a complete `index.html` file.\n" | |
| "3. **File Structure:**\n" | |
| f" - If '{file_structure}' is 'Multiple Files':\n" | |
| " - Use clear markers EXACTLY as follows:\n" | |
| " `<!-- index.html -->`\n" | |
| " `/* style.css */`\n" | |
| " `// script.js` (only include if JavaScript is necessary for the described functionality)\n" | |
| " - Place the corresponding code directly after each marker.\n" | |
| " - Link the CSS (`<link rel='stylesheet' href='style.css'>`) in the `<head>` of `index.html`.\n" | |
| " - Include the JS (`<script src='script.js'></script>`) just before the closing `</body>` tag in `index.html` if `script.js` is generated.\n" | |
| f" - If '{file_structure}' is 'Single File':\n" | |
| " - Embed ALL CSS within `<style>` tags inside the `<head>` of the `index.html` file.\n" | |
| " - Embed ALL necessary JavaScript within `<script>` tags just before the closing `</body>` tag of the `index.html` file.\n" | |
| "4. **Backend Context ({backend_choice}):** This choice provides context. For 'Flask' or 'Node.js', you might include standard template placeholders (like `{{ variable }}` for Flask/Jinja2 or similar patterns for Node templating engines if appropriate for the frontend structure), but primarily focus on generating the static frontend assets (HTML structure, CSS styling, client-side JS interactions). For 'Static', generate standard HTML/CSS/JS without backend-specific placeholders.\n" | |
| "5. **Focus on Frontend:** Generate the client-side code. Do not generate server-side Flask or Node.js code.\n" | |
| "6. **Completeness:** Generate functional code based on the prompt. If the prompt is vague, create a reasonable default structure." | |
| ) | |
| # --- Construct the messages for the API --- | |
| messages = [ | |
| {"role": "system", "content": system_message}, | |
| {"role": "user", "content": f"Create a website based on this description: {prompt}"} # Make user role explicit | |
| ] | |
| # --- Stream the response from the API --- | |
| response_stream = "" | |
| try: | |
| print("Sending request to Hugging Face Inference API...") | |
| # CORRECTED: Removed stop_sequences | |
| for message in client.chat_completion( | |
| messages=messages, | |
| max_tokens=max_tokens, | |
| stream=True, | |
| temperature=temperature, | |
| top_p=top_p, | |
| ): | |
| token = message.choices[0].delta.content | |
| # Basic check to ensure token is a string (it should be) | |
| if isinstance(token, str): | |
| response_stream += token | |
| # Clean potential unwanted prefixes sometimes added by models | |
| # if response_stream.strip().startswith(("```html", "```")): | |
| # response_stream = response_stream.split("\n", 1)[-1] | |
| yield response_stream # Yield the cumulative response for live update | |
| print(f"API stream finished. Total length: {len(response_stream)}") | |
| # Optional: Post-process to remove leading/trailing whitespace or markdown | |
| final_response = response_stream.strip() | |
| # More aggressive cleaning if needed: | |
| # if final_response.startswith("```html"): | |
| # final_response = final_response[7:] | |
| # if final_response.endswith("```"): | |
| # final_response = final_response[:-3] | |
| # yield final_response.strip() # Yield final cleaned response | |
| except Exception as e: | |
| error_message = f"An error occurred during the API call: {e}" | |
| print(error_message) | |
| # Display the error clearly in the output box | |
| yield f"## Error\n\nFailed to generate code.\n**Reason:** {e}\n\nPlease check the model status, your connection, and API token (if applicable)." | |
| # --- Build Gradio Interface using Blocks --- | |
| with gr.Blocks(css=".gradio-container { max-width: 90% !important; }") as demo: # Add CSS for wider layout | |
| gr.Markdown("# Website Code Generator 🚀") | |
| gr.Markdown( | |
| "Describe the website you want, choose your options, and the AI will generate the frontend code (HTML, CSS, JS). " | |
| "The code will appear live in the text editor below. **Note:** The AI generates only frontend code based on your description." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| prompt_input = gr.Textbox( | |
| label="Website Description", | |
| placeholder="e.g., A simple landing page with a navigation bar (Home, About, Contact), a hero section with a title and button, and a simple footer.", | |
| lines=5, # Increased lines for better prompt visibility | |
| ) | |
| backend_radio = gr.Radio( | |
| ["Static", "Flask", "Node.js"], | |
| label="Backend Context Hint", | |
| value="Static", | |
| info="Hint for AI: influences potential template placeholders (e.g., {{var}}) but AI generates ONLY frontend code.", | |
| ) | |
| file_structure_radio = gr.Radio( | |
| ["Multiple Files", "Single File"], # Default to multiple for clarity | |
| label="Output File Structure", | |
| value="Multiple Files", | |
| info="Generate separate files (index.html, style.css, script.js) or embed all in index.html?", | |
| ) | |
| generate_button = gr.Button("Generate Website Code", variant="primary") | |
| with gr.Column(scale=3): | |
| # Use Code component which is better suited for displaying code | |
| code_output = gr.Code( | |
| label="Generated Code", | |
| language="html", # Base language, will contain CSS/JS markers if multiple files | |
| lines=28, # Increased lines | |
| interactive=False, # Read-only display | |
| ) | |
| with gr.Accordion("Advanced Generation Settings", open=False): | |
| max_tokens_slider = gr.Slider( | |
| minimum=512, # Increased minimum for potentially complex sites | |
| maximum=4096, # Match common context lengths | |
| value=2048, # Increased default | |
| step=128, | |
| label="Max New Tokens", | |
| info="Maximum number of tokens (approx. words/code elements) the AI can generate." | |
| ) | |
| temperature_slider = gr.Slider( | |
| minimum=0.1, | |
| maximum=1.2, # Allow slightly higher for more creativity if needed | |
| value=0.6, # Slightly lower default for more predictable code | |
| step=0.1, | |
| label="Temperature", | |
| info="Controls randomness. Lower values (e.g., 0.2) make output more focused, higher values (e.g., 0.9) make it more creative/random." | |
| ) | |
| top_p_slider = gr.Slider( | |
| minimum=0.1, | |
| maximum=1.0, | |
| value=0.9, # Slightly lower default top-p | |
| step=0.05, | |
| label="Top-P (Nucleus Sampling)", | |
| info="Alternative to temperature for controlling randomness. Considers only the most probable tokens with cumulative probability p." | |
| ) | |
| # --- Connect Inputs/Outputs to the Function --- | |
| generate_button.click( | |
| fn=generate_code, | |
| inputs=[ | |
| prompt_input, | |
| backend_radio, | |
| file_structure_radio, | |
| max_tokens_slider, | |
| temperature_slider, | |
| top_p_slider, | |
| ], | |
| outputs=code_output, | |
| #api_name="generate_website_code" # Optional: for API usage | |
| ) | |
| # Add examples for guidance | |
| gr.Examples( | |
| examples=[ | |
| ["A simple counter page with a number display, an increment button, and a decrement button. Use Javascript for the logic.", "Static", "Single File"], | |
| ["A login form with fields for username and password, and a submit button. Basic styling.", "Static", "Multiple Files"], | |
| ["Product cards display grid. Each card should show an image, product name, price, and an 'Add to Cart' button. Make it responsive.", "Static", "Multiple Files"], | |
| ["A personal blog homepage with a header, a list of recent posts (just placeholders), and a sidebar with categories.", "Flask", "Multiple Files"], | |
| ], | |
| inputs=[prompt_input, backend_radio, file_structure_radio], | |
| label="Example Prompts" # Optional label for the examples section | |
| ) | |
| # --- Launch the App --- | |
| if __name__ == "__main__": | |
| print("Starting Gradio app...") | |
| # Enable queuing for handling multiple users, essential for Spaces | |
| # Increase concurrency count if needed and if your hardware/Space plan supports it | |
| demo.queue(max_size=10).launch( | |
| # debug=True, # Set debug=False for production/Spaces deployment | |
| # share=False # Set share=True to create a temporary public link (useful for local testing) | |
| ) | |
| print("Gradio app launched.") |