File size: 2,305 Bytes
f402022
85b51b6
f2b7ae6
2564711
 
 
e3cbb09
d034098
c4311e1
e3cbb09
 
afdb3f9
2564711
e3cbb09
c4311e1
 
 
2564711
 
c4311e1
e3cbb09
 
2564711
 
 
 
 
ed2ecef
e3cbb09
ed2ecef
e3cbb09
 
 
 
ed2ecef
2564711
 
ed2ecef
2564711
 
 
 
 
 
 
 
 
 
 
 
e3cbb09
2564711
e3cbb09
 
2564711
 
e3cbb09
2564711
 
 
 
 
 
 
 
 
c4311e1
2564711
 
 
 
e3cbb09
59a6b22
2564711
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
import os
import gradio as gr
import requests
from PIL import Image
from io import BytesIO
import base64
import logging

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def generate_image(prompt):
    try:
        RUNPOD_API_KEY = os.getenv("RUNPOD_API_KEY")
        RUNPOD_ENDPOINT = os.getenv("RUNPOD_ENDPOINT")
        
        if not all([RUNPOD_API_KEY, RUNPOD_ENDPOINT]):
            return None, "API credentials not configured"
        
        payload = {
            "input": {
                "prompt": prompt,
                "width": 768,
                "height": 768,
                "num_inference_steps": 20,
                "return_base64": True
            }
        }
        
        headers = {
            "Authorization": f"Bearer {RUNPOD_API_KEY}",
            "Content-Type": "application/json"
        }
        
        logger.info(f"Sending request to RunPod: {payload}")
        response = requests.post(RUNPOD_ENDPOINT, json=payload, headers=headers, timeout=30)
        
        if response.status_code != 200:
            return None, f"API Error {response.status_code}: {response.text}"
        
        data = response.json()
        
        if "image" in data:
            image_data = base64.b64decode(data["image"].split(",")[1])
            return Image.open(BytesIO(image_data)), "Success!"
        elif "output" in data and data["output"]:
            image_url = data["output"][0]
            img_data = requests.get(image_url, timeout=30).content
            return Image.open(BytesIO(img_data)), "Success!"
        else:
            return None, "Unexpected response format"
            
    except Exception as e:
        logger.error(f"Error: {str(e)}")
        return None, f"Error: {str(e)}"

with gr.Blocks() as app:
    gr.Markdown("# 🎨 RunPod Image Generator")
    
    with gr.Row():
        prompt = gr.Textbox(label="Enter prompt", value="A comic book superhero")
        generate_btn = gr.Button("Generate")
    
    output = gr.Image(label="Generated Image")
    status = gr.Textbox(label="Status")
    
    generate_btn.click(
        generate_image,
        inputs=[prompt],
        outputs=[output, status]
    )

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)