import os import json import gradio as gr from huggingface_hub import InferenceClient from huggingface_hub import get_token as hf_get_token from gradio.context import LocalContext import contextvars workflow_token = contextvars.ContextVar("workflow_token", default=None) def get_hf_token() -> str | None: """ Retrieves the HF API token from either the workflow context, the user's Gradio OAuth session, or falls back to the system environment. """ w_token = workflow_token.get() if w_token: return w_token request = LocalContext.request.get(None) if request is not None: session = getattr(request, "session", {}) oauth_info = session.get("oauth_info", {}) if oauth_info: token = oauth_info.get("access_token") if token and token != "mock-oauth-token-for-local-dev": return token try: return hf_get_token() except Exception: return None def generate_prompt(concept: str) -> str: """ Expands a simple concept into a detailed image prompt using the NVIDIA Nemotron model. """ if not concept: return "a ginger cat wearing a tiny wizard hat reading a spellbook" try: token = get_hf_token() or os.environ.get("HF_TOKEN") or os.environ.get("HF_API_TOKEN") client = InferenceClient( provider="together", api_key=token, bill_to="huggingface", ) system_instruction = ( "You are an expert prompt engineer for text-to-image models. " "Your task is to take a simple concept and expand it into a detailed, " "vivid, and high-quality image prompt for FLUX.1-dev. " "Describe the scene, lighting, materials, and aesthetic in detail. " "Provide ONLY the final prompt text. Do not include any introductory or concluding text, " "do not provide multiple options, and do not wrap the prompt in quotes." ) messages = [ {"role": "system", "content": system_instruction}, {"role": "user", "content": f"Concept: {concept}"} ] response = client.chat_completion( model="nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4", messages=messages, temperature=0.7, max_tokens=256 ) result = response.choices[0].message.content clean_result = str(result).strip() if clean_result.startswith('"') and clean_result.endswith('"'): clean_result = clean_result[1:-1] elif clean_result.startswith("'") and clean_result.endswith("'"): clean_result = clean_result[1:-1] return clean_result except Exception as e: print(f"Error calling Nemotron model: {e}") return f"A detailed, high-quality, professional commercial product photograph of {concept}" def show_nemotron_prompt(prompt: str) -> str: """ Displays the detailed image prompt generated by the NVIDIA Nemotron model. """ return prompt def generate_z_image(prompt: str) -> dict: """ Generates an image from a prompt using the Tongyi-MAI/Z-Image-Turbo model. Returns a dictionary structure compatible with Gradio's image viewer. """ if not prompt: prompt = "a ginger cat wearing a tiny wizard hat reading a spellbook" try: token = get_hf_token() or os.environ.get("HF_TOKEN") or os.environ.get("HF_API_TOKEN") client = InferenceClient( provider="auto", api_key=token, bill_to="huggingface", ) image = client.text_to_image( prompt, model="Tongyi-MAI/Z-Image-Turbo", ) import tempfile import uuid temp_dir = tempfile.gettempdir() filepath = os.path.join(temp_dir, f"{uuid.uuid4()}.png") image.save(filepath) return { "path": filepath, "url": f"/gradio_api/file={filepath}", "is_file": True } except Exception as e: print(f"Error calling Z-Image-Turbo model: {e}") raise e class LockedWorkflow(gr.Workflow): def _build(self): from gradio.components.workflowcanvas import WorkflowCanvas from gradio.workflow import ( get_token, call_space, call_model, fetch_dataset, get_dataset_schema, _workflow_from_bind, ) from gradio.utils import get_space def _load_initial() -> str | None: try: with open(self._workflow_file, encoding="utf-8") as f: return f.read() except FileNotFoundError: if self._bound: return _workflow_from_bind( self._bound, self._edges, name=self._workflow_name ) return None bound = self._bound def call_fn(data, _token=None) -> str: if _token: t_str = _token.token if hasattr(_token, "token") else _token if t_str: workflow_token.set(str(t_str)) fn_name = data[0] if data else "" try: args_json = data[1] if len(data) > 1 else "[]" fn = bound.get(fn_name) if fn is None: return json.dumps( { "error": f"No function '{fn_name}' bound to this workflow", "error_type": "unknown", "suggestion": "Check the bind= argument to Workflow()", } ) args = json.loads(args_json) if not isinstance(args, list): args = [args] result = fn(*args) result = list(result) if isinstance(result, (list, tuple)) else [result] return json.dumps(result) except Exception as e: return json.dumps( {"error": str(e), "error_type": "unknown", "suggestion": ""} ) # Mocked/neutralized methods for locked read-only mode: def save_workflow(data, _token=None) -> str: return "ok" def list_bound_fns(_data=None, _token=None) -> str: return "[]" def search_spaces(data, token=None) -> str: return "[]" def search_models(data, token=None) -> str: return "[]" def search_datasets(data, token=None) -> str: return "[]" server_functions = [ get_token, call_space, call_model, fetch_dataset, search_spaces, search_models, search_datasets, get_dataset_schema, call_fn, list_bound_fns, save_workflow, ] with self: if get_space() is not None and os.getenv("OAUTH_CLIENT_ID"): gr.LoginButton(visible=False) WorkflowCanvas( value=_load_initial, server_functions=server_functions, ) demo = LockedWorkflow( bind=[generate_prompt, show_nemotron_prompt, generate_z_image], edges=[ ("generate_prompt", "show_nemotron_prompt.prompt"), ("show_nemotron_prompt", "generate_z_image.prompt"), ] ) if __name__ == "__main__": demo.launch()