Create yml/generate-flux.py

yml/generate-flux.py

@@ -0,0 +1,30 @@
+import os, sys, random, requests
+from gradio_client import Client
+
+prompt = os.environ.get('PROMPT', '')
+width = int(os.environ.get('WIDTH', 1024))
+height = int(os.environ.get('HEIGHT', 1024))
+run_id = os.environ.get('RUN_ID')
+space_url = os.environ.get('SPACE_URL')
+github_run_id = os.environ.get('GITHUB_RUN_ID')
+
+print('1. Connecting to Flux...')
+try:
+    client = Client('black-forest-labs/FLUX.1-dev')
+    seed = random.randint(1, 2147483647)
+    
+    result = client.predict(prompt, seed, True, width, height, 3.5, 28, api_name='/infer')
+    image_path = result[0] if isinstance(result, (tuple, list)) else result
+    
+    print(f'2. Uploading back to Docker Space...')
+    with open(image_path, 'rb') as f:
+        res = requests.post(f'{space_url}/api/webhook/upload', data={'run_id': run_id, 'github_run_id': github_run_id}, files={'file': f})
+    
+    if res.status_code == 200:
+        print('3. SUCCESS!')
+    else:
+        sys.exit(1)
+        
+except Exception as e:
+    print(f'ERROR: {e}')
+    sys.exit(1)