Abe commited on
Commit
8247a04
·
1 Parent(s): caf3333

initial copy

Browse files
Files changed (11) hide show
  1. .env.example +10 -0
  2. .gitignore +3 -0
  3. Project.md +46 -0
  4. api.py +98 -0
  5. api_example.py +88 -0
  6. app.py +100 -0
  7. config.py +22 -0
  8. inference.py +86 -0
  9. main.py +63 -0
  10. requirements.txt +6 -0
  11. spaces_config.json +23 -0
.env.example ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ # Hugging Face token
2
+ HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxxxxxx
3
+
4
+ # API settings
5
+ API_HOST=0.0.0.0
6
+ API_PORT=8000
7
+
8
+ # Gradio settings
9
+ GRADIO_HOST=0.0.0.0
10
+ GRADIO_PORT=7860
.gitignore ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ .venv
2
+ *.pyc
3
+ __pycache__
Project.md ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Diffusion Models App
2
+
3
+ A Python application that uses Hugging Face inference endpoints for text-to-image and image-to-image generation with a Gradio UI and API endpoints.
4
+
5
+ ## Features
6
+
7
+ - Text-to-image generation
8
+ - Image-to-image transformation with optional prompt
9
+ - Gradio UI for interactive use
10
+ - API endpoints for integration with other applications
11
+ - Configurable models via text input
12
+
13
+ ## Project Structure
14
+
15
+ - `main.py` - Entry point that can run both UI and API
16
+ - `app.py` - Gradio UI implementation
17
+ - `api.py` - FastAPI server for API endpoints
18
+ - `inference.py` - Core functionality for HF inference
19
+ - `config.py` - Configuration and settings
20
+ - `requirements.txt` - Dependencies
21
+
22
+ ## Setup & Usage
23
+
24
+ 1. Clone the repository
25
+ 2. Create a .env file with your Hugging Face token (copy from .env.example)
26
+ 3. Install dependencies: `pip install -r requirements.txt`
27
+ 4. Run the application: `python main.py`
28
+
29
+ ## Running Options
30
+
31
+ - Run both UI and API: `python main.py`
32
+ - Run only the API: `python main.py --mode api`
33
+ - Run only the UI: `python main.py --mode ui`
34
+
35
+ ## API Endpoints
36
+
37
+ - `POST /text-to-image` - Generate an image from text
38
+ - `POST /image-to-image` - Transform an image with optional prompt
39
+
40
+ ## Environment Variables
41
+
42
+ - `HF_TOKEN` - Your Hugging Face API token
43
+ - `API_HOST` - Host for the API server (default: 0.0.0.0)
44
+ - `API_PORT` - Port for the API server (default: 8000)
45
+ - `GRADIO_HOST` - Host for the Gradio UI (default: 0.0.0.0)
46
+ - `GRADIO_PORT` - Port for the Gradio UI (default: 7860)
api.py ADDED
@@ -0,0 +1,98 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, File, UploadFile, Form, HTTPException
2
+ from fastapi.responses import Response
3
+ from PIL import Image
4
+ import io
5
+ import uvicorn
6
+ import config
7
+ from inference import DiffusionInference
8
+
9
+ app = FastAPI(title="Diffusion Models API")
10
+
11
+ # Initialize the inference class
12
+ inference = DiffusionInference()
13
+
14
+ @app.get("/")
15
+ async def root():
16
+ return {"message": "Diffusion Models API is running"}
17
+
18
+ @app.post("/text-to-image")
19
+ async def text_to_image(
20
+ prompt: str = Form(...),
21
+ model: str = Form(None),
22
+ negative_prompt: str = Form(None),
23
+ guidance_scale: float = Form(7.5),
24
+ num_inference_steps: int = Form(50)
25
+ ):
26
+ """
27
+ Generate an image from a text prompt
28
+ """
29
+ try:
30
+ # Use default model if not specified
31
+ if not model:
32
+ model = config.DEFAULT_TEXT2IMG_MODEL
33
+
34
+ # Call the inference module
35
+ image = inference.text_to_image(
36
+ prompt=prompt,
37
+ model_name=model,
38
+ negative_prompt=negative_prompt,
39
+ guidance_scale=guidance_scale,
40
+ num_inference_steps=num_inference_steps
41
+ )
42
+
43
+ # Convert PIL image to bytes
44
+ img_byte_arr = io.BytesIO()
45
+ image.save(img_byte_arr, format='PNG')
46
+ img_byte_arr = img_byte_arr.getvalue()
47
+
48
+ return Response(content=img_byte_arr, media_type="image/png")
49
+ except Exception as e:
50
+ raise HTTPException(status_code=500, detail=str(e))
51
+
52
+ @app.post("/image-to-image")
53
+ async def image_to_image(
54
+ image: UploadFile = File(...),
55
+ prompt: str = Form(None),
56
+ model: str = Form(None),
57
+ negative_prompt: str = Form(None),
58
+ guidance_scale: float = Form(7.5),
59
+ num_inference_steps: int = Form(50)
60
+ ):
61
+ """
62
+ Generate a new image from an input image and optional prompt
63
+ """
64
+ try:
65
+ # Read and convert input image
66
+ contents = await image.read()
67
+ input_image = Image.open(io.BytesIO(contents))
68
+
69
+ # Use default model if not specified
70
+ if not model:
71
+ model = config.DEFAULT_IMG2IMG_MODEL
72
+
73
+ # Call the inference module
74
+ result = inference.image_to_image(
75
+ image=input_image,
76
+ prompt=prompt,
77
+ model_name=model,
78
+ negative_prompt=negative_prompt,
79
+ guidance_scale=guidance_scale,
80
+ num_inference_steps=num_inference_steps
81
+ )
82
+
83
+ # Convert PIL image to bytes
84
+ img_byte_arr = io.BytesIO()
85
+ result.save(img_byte_arr, format='PNG')
86
+ img_byte_arr = img_byte_arr.getvalue()
87
+
88
+ return Response(content=img_byte_arr, media_type="image/png")
89
+ except Exception as e:
90
+ raise HTTPException(status_code=500, detail=str(e))
91
+
92
+ if __name__ == "__main__":
93
+ uvicorn.run(
94
+ "api:app",
95
+ host=config.API_HOST,
96
+ port=config.API_PORT,
97
+ reload=True
98
+ )
api_example.py ADDED
@@ -0,0 +1,88 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import requests
2
+ import io
3
+ from PIL import Image
4
+ import os
5
+ from dotenv import load_dotenv
6
+
7
+ # Load environment variables from .env file
8
+ load_dotenv()
9
+
10
+ # Hugging Face API token (need to set in .env or environment)
11
+ HF_TOKEN = os.getenv("HF_TOKEN")
12
+
13
+ # API base URL
14
+ API_BASE = "http://localhost:8000"
15
+
16
+ def text_to_image(prompt, model=None, negative_prompt=None):
17
+ """
18
+ Generate image from text using the API
19
+ """
20
+ url = f"{API_BASE}/text-to-image"
21
+
22
+ # Prepare form data
23
+ data = {
24
+ "prompt": prompt,
25
+ }
26
+
27
+ if model:
28
+ data["model"] = model
29
+
30
+ if negative_prompt:
31
+ data["negative_prompt"] = negative_prompt
32
+
33
+ # Make API request
34
+ response = requests.post(url, data=data)
35
+
36
+ if response.status_code == 200:
37
+ # Convert response to PIL image
38
+ image = Image.open(io.BytesIO(response.content))
39
+ return image
40
+ else:
41
+ print(f"Error: {response.status_code}")
42
+ print(response.text)
43
+ return None
44
+
45
+ def image_to_image(image_path, prompt=None, model=None):
46
+ """
47
+ Transform image using the API
48
+ """
49
+ url = f"{API_BASE}/image-to-image"
50
+
51
+ # Prepare form data and files
52
+ data = {}
53
+ if prompt:
54
+ data["prompt"] = prompt
55
+
56
+ if model:
57
+ data["model"] = model
58
+
59
+ files = {
60
+ "image": open(image_path, "rb")
61
+ }
62
+
63
+ # Make API request
64
+ response = requests.post(url, data=data, files=files)
65
+
66
+ if response.status_code == 200:
67
+ # Convert response to PIL image
68
+ image = Image.open(io.BytesIO(response.content))
69
+ return image
70
+ else:
71
+ print(f"Error: {response.status_code}")
72
+ print(response.text)
73
+ return None
74
+
75
+ if __name__ == "__main__":
76
+ # Example usage
77
+ print("Text to Image example:")
78
+ image = text_to_image("A beautiful mountain landscape at sunset")
79
+ if image:
80
+ image.save("text2img_output.png")
81
+ print("Image saved as text2img_output.png")
82
+
83
+ print("Image to Image example (requires an input image):")
84
+ # Uncomment and modify path to run:
85
+ # result = image_to_image("input.png", "Turn this into a fantasy scene")
86
+ # if result:
87
+ # result.save("img2img_output.png")
88
+ # print("Image saved as img2img_output.png")
app.py ADDED
@@ -0,0 +1,100 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import config
3
+ from inference import DiffusionInference
4
+ from PIL import Image
5
+ import io
6
+
7
+ # Initialize the inference class
8
+ inference = DiffusionInference()
9
+
10
+ def text_to_image_fn(prompt, model, negative_prompt=None, guidance_scale=7.5, num_inference_steps=50):
11
+ """
12
+ Handle text to image generation request
13
+ """
14
+ try:
15
+ if not model:
16
+ model = config.DEFAULT_TEXT2IMG_MODEL
17
+
18
+ # Call the inference module
19
+ image = inference.text_to_image(
20
+ prompt=prompt,
21
+ model_name=model,
22
+ negative_prompt=negative_prompt,
23
+ guidance_scale=guidance_scale,
24
+ num_inference_steps=num_inference_steps
25
+ )
26
+
27
+ return image, None
28
+ except Exception as e:
29
+ return None, str(e)
30
+
31
+ def image_to_image_fn(image, prompt, model, negative_prompt=None, guidance_scale=7.5, num_inference_steps=50):
32
+ """
33
+ Handle image to image transformation request
34
+ """
35
+ try:
36
+ if not model:
37
+ model = config.DEFAULT_IMG2IMG_MODEL
38
+
39
+ # Call the inference module
40
+ result = inference.image_to_image(
41
+ image=image,
42
+ prompt=prompt,
43
+ model_name=model,
44
+ negative_prompt=negative_prompt,
45
+ guidance_scale=guidance_scale,
46
+ num_inference_steps=num_inference_steps
47
+ )
48
+
49
+ return result, None
50
+ except Exception as e:
51
+ return None, str(e)
52
+
53
+ # Create Gradio UI
54
+ with gr.Blocks(title="Diffusion Models") as app:
55
+ gr.Markdown("# Hugging Face Diffusion Models")
56
+
57
+ with gr.Tab("Text to Image"):
58
+ with gr.Row():
59
+ with gr.Column():
60
+ txt2img_prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
61
+ txt2img_negative = gr.Textbox(label="Negative Prompt (Optional)", placeholder="What to exclude from the image")
62
+ txt2img_model = gr.Textbox(label="Model", placeholder=f"Enter model name (default: {config.DEFAULT_TEXT2IMG_MODEL})")
63
+ txt2img_guidance = gr.Slider(minimum=1.0, maximum=20.0, value=7.5, step=0.5, label="Guidance Scale")
64
+ txt2img_steps = gr.Slider(minimum=10, maximum=100, value=50, step=1, label="Inference Steps")
65
+ txt2img_button = gr.Button("Generate Image")
66
+
67
+ with gr.Column():
68
+ txt2img_output = gr.Image(type="pil", label="Generated Image")
69
+ txt2img_error = gr.Textbox(label="Error", visible=True)
70
+
71
+ txt2img_button.click(
72
+ fn=text_to_image_fn,
73
+ inputs=[txt2img_prompt, txt2img_model, txt2img_negative, txt2img_guidance, txt2img_steps],
74
+ outputs=[txt2img_output, txt2img_error]
75
+ )
76
+
77
+ with gr.Tab("Image to Image"):
78
+ with gr.Row():
79
+ with gr.Column():
80
+ img2img_input = gr.Image(type="pil", label="Input Image")
81
+ img2img_prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")
82
+ img2img_negative = gr.Textbox(label="Negative Prompt (Optional)", placeholder="What to exclude from the image")
83
+ img2img_model = gr.Textbox(label="Model", placeholder=f"Enter model name (default: {config.DEFAULT_IMG2IMG_MODEL})")
84
+ img2img_guidance = gr.Slider(minimum=1.0, maximum=20.0, value=7.5, step=0.5, label="Guidance Scale")
85
+ img2img_steps = gr.Slider(minimum=10, maximum=100, value=50, step=1, label="Inference Steps")
86
+ img2img_button = gr.Button("Transform Image")
87
+
88
+ with gr.Column():
89
+ img2img_output = gr.Image(type="pil", label="Generated Image")
90
+ img2img_error = gr.Textbox(label="Error", visible=True)
91
+
92
+ img2img_button.click(
93
+ fn=image_to_image_fn,
94
+ inputs=[img2img_input, img2img_prompt, img2img_model, img2img_negative, img2img_guidance, img2img_steps],
95
+ outputs=[img2img_output, img2img_error]
96
+ )
97
+
98
+ # Launch the Gradio app
99
+ if __name__ == "__main__":
100
+ app.launch(server_name=config.GRADIO_HOST, server_port=config.GRADIO_PORT)
config.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from dotenv import load_dotenv
3
+
4
+ # Load environment variables from .env file
5
+ load_dotenv()
6
+
7
+ # Hugging Face API token
8
+ HF_TOKEN = os.getenv("HF_TOKEN", "")
9
+
10
+ # Default model for text to image
11
+ DEFAULT_TEXT2IMG_MODEL = "stabilityai/stable-diffusion-2-1"
12
+
13
+ # Default model for image to image
14
+ DEFAULT_IMG2IMG_MODEL = "lllyasviel/sd-controlnet-depth"
15
+
16
+ # API settings
17
+ API_HOST = os.getenv("API_HOST", "0.0.0.0")
18
+ API_PORT = int(os.getenv("API_PORT", "8000"))
19
+
20
+ # Gradio settings
21
+ GRADIO_HOST = os.getenv("GRADIO_HOST", "0.0.0.0")
22
+ GRADIO_PORT = int(os.getenv("GRADIO_PORT", "7860"))
inference.py ADDED
@@ -0,0 +1,86 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from huggingface_hub import InferenceClient
2
+ from PIL import Image
3
+ import io
4
+ import config
5
+
6
+
7
+ class DiffusionInference:
8
+ def __init__(self, api_key=None):
9
+ """
10
+ Initialize the inference client with the Hugging Face API token.
11
+ """
12
+ self.api_key = api_key or config.HF_TOKEN
13
+ self.client = InferenceClient(
14
+ provider="hf-inference",
15
+ api_key=self.api_key,
16
+ )
17
+
18
+ def text_to_image(self, prompt, model_name=None, negative_prompt=None, **kwargs):
19
+ """
20
+ Generate an image from a text prompt.
21
+
22
+ Args:
23
+ prompt (str): The text prompt to guide image generation
24
+ model_name (str, optional): The model to use for inference
25
+ negative_prompt (str, optional): What not to include in the image
26
+ **kwargs: Additional parameters to pass to the model
27
+
28
+ Returns:
29
+ PIL.Image: The generated image
30
+ """
31
+ model = model_name or config.DEFAULT_TEXT2IMG_MODEL
32
+
33
+ # Set up parameters dictionary
34
+ params = {"prompt": prompt}
35
+
36
+ if negative_prompt:
37
+ params["negative_prompt"] = negative_prompt
38
+
39
+ # Add any additional parameters
40
+ params.update(kwargs)
41
+
42
+ try:
43
+ image = self.client.text_to_image(model=model, **params)
44
+ return image
45
+ except Exception as e:
46
+ print(f"Error generating image: {e}")
47
+ raise
48
+
49
+ def image_to_image(self, image, prompt=None, model_name=None, negative_prompt=None, **kwargs):
50
+ """
51
+ Generate a new image from an input image and optional prompt.
52
+
53
+ Args:
54
+ image (PIL.Image or str): Input image or path to image
55
+ prompt (str, optional): Text prompt to guide the transformation
56
+ model_name (str, optional): The model to use for inference
57
+ negative_prompt (str, optional): What not to include in the image
58
+ **kwargs: Additional parameters to pass to the model
59
+
60
+ Returns:
61
+ PIL.Image: The generated image
62
+ """
63
+ model = model_name or config.DEFAULT_IMG2IMG_MODEL
64
+
65
+ # Convert image path to PIL Image if needed
66
+ if isinstance(image, str):
67
+ image = Image.open(image)
68
+
69
+ # Set up parameters dictionary
70
+ params = {"image": image}
71
+
72
+ if prompt:
73
+ params["prompt"] = prompt
74
+
75
+ if negative_prompt:
76
+ params["negative_prompt"] = negative_prompt
77
+
78
+ # Add any additional parameters
79
+ params.update(kwargs)
80
+
81
+ try:
82
+ result = self.client.image_to_image(model=model, **params)
83
+ return result
84
+ except Exception as e:
85
+ print(f"Error transforming image: {e}")
86
+ raise
main.py ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import uvicorn
3
+ import threading
4
+ import os
5
+ import config
6
+ from app import app as gradio_app
7
+ from api import app as api_app
8
+
9
+ def run_api():
10
+ """Run the FastAPI server"""
11
+ uvicorn.run(
12
+ api_app,
13
+ host=config.API_HOST,
14
+ port=config.API_PORT
15
+ )
16
+
17
+ def run_gradio():
18
+ """Run the Gradio interface"""
19
+ gradio_app.launch(
20
+ server_name=config.GRADIO_HOST,
21
+ server_port=config.GRADIO_PORT,
22
+ share=False
23
+ )
24
+
25
+ def main():
26
+ parser = argparse.ArgumentParser(description="Run Diffusion Models App")
27
+ parser.add_argument(
28
+ "--mode",
29
+ type=str,
30
+ default="all",
31
+ choices=["all", "api", "ui"],
32
+ help="Which component to run: 'all' (default), 'api', or 'ui'"
33
+ )
34
+
35
+ args = parser.parse_args()
36
+
37
+ # Check if HF_TOKEN is set
38
+ if not config.HF_TOKEN:
39
+ print("Warning: HF_TOKEN environment variable is not set. Please set it for API access.")
40
+ print("You can create a .env file with HF_TOKEN=your_token or set it in your environment.")
41
+
42
+ if args.mode == "all":
43
+ # Run both API and UI in separate threads
44
+ api_thread = threading.Thread(target=run_api)
45
+ api_thread.daemon = True
46
+ api_thread.start()
47
+
48
+ print(f"API server running at http://{config.API_HOST}:{config.API_PORT}")
49
+ print(f"Starting Gradio UI at http://{config.GRADIO_HOST}:{config.GRADIO_PORT}")
50
+
51
+ # Run Gradio in the main thread
52
+ run_gradio()
53
+
54
+ elif args.mode == "api":
55
+ print(f"Starting API server at http://{config.API_HOST}:{config.API_PORT}")
56
+ run_api()
57
+
58
+ elif args.mode == "ui":
59
+ print(f"Starting Gradio UI at http://{config.GRADIO_HOST}:{config.GRADIO_PORT}")
60
+ run_gradio()
61
+
62
+ if __name__ == "__main__":
63
+ main()
requirements.txt ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ gradio
2
+ huggingface_hub
3
+ Pillow
4
+ fastapi
5
+ uvicorn
6
+ python-dotenv
spaces_config.json ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "sdk": "gradio",
3
+ "sdk_version": "3.50.2",
4
+ "app_file": "app.py",
5
+ "models": [
6
+ {
7
+ "model_name": "stabilityai/stable-diffusion-2-1",
8
+ "model_class": "diffusers"
9
+ },
10
+ {
11
+ "model_name": "lllyasviel/sd-controlnet-depth",
12
+ "model_class": "diffusers"
13
+ }
14
+ ],
15
+ "resources": {
16
+ "accelerator": "gpu",
17
+ "gpu": {
18
+ "count": 1,
19
+ "vendor": "nvidia",
20
+ "memory": "16GB"
21
+ }
22
+ }
23
+ }