Initial setup for image generation UI v2

Dockerfile

@@ -0,0 +1,26 @@
+FROM python:3.10-slim
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+git git-lfs ffmpeg libsm6 libxext6 cmake rsync libgl1-mesa-glx \
+&& rm -rf /var/lib/apt/lists/* \
+&& git lfs install
+
+# Set working directory
+WORKDIR /app
+
+# Install Python dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy application files
+COPY . .
+
+# Set user
+RUN useradd -m -u 1000 user && chown -R user:user /app
+USER user
+
+# Expose port
+EXPOSE 7860
+
+# Run the app
+CMD ["python", "app.py"]

README.md

@@ -1,10 +1,18 @@
----
-title: Image Generation Ui V2
-emoji: 🚀
-colorFrom: indigo
-colorTo: pink
-sdk: docker
-pinned: false
----
+Image Generation UI v2
+A custom Hugging Face Space for generating images using SSD-1B and Stable Diffusion v1-5 models.
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+ ## Setup
+ - Clone the repository: `git clone https://huggingface.co/spaces/remiai3/image-generation-ui-v2`.
+ - Install dependencies: `pip install -r requirements.txt`.
+ - Set `HF_TOKEN` environment variable with your Hugging Face Access Token.
+ - Run: `python app.py`.
+
+ ## Features
+ - Select between SSD-1B and Stable Diffusion v1-5 models.
+ - Choose image ratios: 1:1, 3:4, 16:9.
+ - Generate up to 4 images (SSD-1B limited to 1).
+ - Custom UI with logo.
+
+ ## Models
+ - [SSD-1B](https://huggingface.co/remiai3/ssd-1b)
+ - [Stable Diffusion v1-5](https://huggingface.co/remiai3/stable-diffusion-v1-5)

app.py

@@ -0,0 +1,114 @@
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+from diffusers import StableDiffusionPipeline, DiffusionPipeline, DPMSolverMultistepScheduler
+import torch
+import os
+from PIL import Image
+import base64
+import time
+from accelerate import Accelerator
+import logging
+
+app = Flask(__name__)
+CORS(app)
+
+     # Configure logging
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+     # Initialize Accelerator for CPU
+accelerator = Accelerator(cpu=True)
+
+     # Model cache
+model_cache = {}
+model_paths = {
+         "ssd-1b": "remiai3/ssd-1b",
+         "sd-v1-5": "remiai3/stable-diffusion-v1-5"
+     }
+
+     # Image ratio to dimensions (optimized for CPU)
+ratio_to_dims = {
+         "1:1": (256, 256),
+         "3:4": (192, 256),
+         "16:9": (256, 144)
+     }
+
+def load_model(model_id):
+         if model_id not in model_cache:
+             logger.info(f"Loading model {model_id}...")
+             try:
+                 pipe = StableDiffusionPipeline.from_pretrained(
+                     model_paths[model_id],
+                     torch_dtype=torch.float32,
+                    use_auth_token=os.getenv("HF_TOKEN"),
+                    use_safetensors=True,
+                     low_cpu_mem_usage=True
+                 )
+                 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+                 pipe = accelerator.prepare(pipe)
+                 pipe.enable_attention_slicing()
+                 pipe.enable_sequential_cpu_offload()
+                 model_cache[model_id] = pipe
+                 logger.info(f"Model {model_id} loaded successfully")
+             except Exception as e:
+                 logger.error(f"Error loading model {model_id}: {str(e)}")
+                 raise
+         return model_cache[model_id]
+
+@app.route('/')
+def index():
+         return app.send_static_file('index.html')
+
+@app.route('/generate', methods=['POST'])
+def generate():
+         try:
+             data = request.json
+             model_id = data.get('model', 'ssd-1b')
+             prompt = data.get('prompt', '')
+             ratio = data.get('ratio', '1:1')
+             num_images = min(int(data.get('num_images', 1)), 4)
+             guidance_scale = float(data.get('guidance_scale', 7.5))
+
+             if not prompt:
+                 return jsonify({"error": "Prompt is required"}), 400
+
+             if model_id == 'ssd-1b' and num_images > 1:
+                 return jsonify({"error": "SSD-1B allows only 1 image per generation"}), 400
+             if model_id == 'ssd-1b' and ratio != '1:1':
+                 return jsonify({"error": "SSD-1B supports only 1:1 ratio"}), 400
+             if model_id == 'sd-v1-5' and len(prompt.split()) > 77:
+                 return jsonify({"error": "Prompt exceeds 77 tokens for Stable Diffusion v1.5"}), 400
+
+             width, height = ratio_to_dims.get(ratio, (256, 256))
+             pipe = load_model(model_id)
+
+             images = []
+             for _ in range(num_images):
+                 image = pipe(
+                     prompt=prompt,
+                     height=height,
+                     width=width,
+                     num_inference_steps=20,
+                     guidance_scale=guidance_scale
+                 ).images[0]
+                 images.append(image)
+
+             output_dir = "outputs"
+             os.makedirs(output_dir, exist_ok=True)
+             image_urls = []
+             for i, img in enumerate(images):
+                 img_path = os.path.join(output_dir, f"generated_{int(time.time())}_{i}.png")
+                 img.save(img_path)
+                 with open(img_path, "rb") as f:
+                     img_data = base64.b64encode(f.read()).decode('utf-8')
+                 image_urls.append(f"data:image/png;base64,{img_data}")
+                 os.remove(img_path)
+
+             return jsonify({"images": image_urls})
+
+         except Exception as e:
+             logger.error(f"Image generation failed: {str(e)}")
+             return jsonify({"error": f"Image generation failed: {str(e)}"}), 500
+
+if __name__ == '__main__':
+         app.run(host='0.0.0.0', port=7860)

index.html

@@ -0,0 +1,148 @@
+<!DOCTYPE html>
+     <html lang="en">
+     <head>
+         <meta charset="UTF-8">
+         <meta name="viewport" content="width=device-width, initial-scale=1.0">
+         <title>Image Generation UI</title>
+         <style>
+             body {
+                 font-family: Arial, sans-serif;
+                 text-align: center;
+                 margin: 0;
+                 padding: 20px;
+                 background-color: #f0f0f0;
+             }
+             .container {
+                 max-width: 800px;
+                 margin: 0 auto;
+             }
+             img.logo {
+                 max-width: 150px;
+                 margin-bottom: 20px;
+             }
+             select, input, button {
+                 margin: 10px;
+                 padding: 8px;
+                 font-size: 16px;
+                 border-radius: 4px;
+                 border: 1px solid #ccc;
+             }
+             button {
+                 background-color: #4CAF50;
+                 color: white;
+                 cursor: pointer;
+             }
+             button:hover {
+                 background-color: #45a049;
+             }
+             #prompt {
+                 width: 300px;
+             }
+             #output {
+                 margin-top: 20px;
+                 display: flex;
+                 flex-wrap: wrap;
+                 justify-content: center;
+                 gap: 10px;
+             }
+             .output-image {
+                 max-width: 200px;
+                 border: 1px solid #ddd;
+                 border-radius: 4px;
+             }
+             .error {
+                 color: red;
+                 margin-top: 10px;
+             }
+             .loading {
+                 display: none;
+                 margin-top: 10px;
+                 font-style: italic;
+             }
+         </style>
+     </head>
+     <body>
+         <div class="container">
+             <img src="/assets/logo.png" alt="Logo" class="logo">
+             <h1>Image Generation UI</h1>
+             <div>
+                 <label for="model-select">Select Model:</label>
+                 <select id="model-select">
+                     <option value="ssd-1b">SSD-1B</option>
+                     <option value="sd-v1-5">Stable Diffusion v1-5</option>
+                 </select>
+             </div>
+             <div>
+                 <label for="ratio-select">Image Ratio:</label>
+                 <select id="ratio-select">
+                     <option value="1:1">1:1</option>
+                     <option value="3:4">3:4</option>
+                     <option value="16:9">16:9</option>
+                 </select>
+             </div>
+             <div>
+                 <label for="num-images">Number of Images (1-4):</label>
+                 <input type="number" id="num-images" min="1" max="4" value="1">
+             </div>
+             <div>
+                 <label for="prompt">Prompt:</label>
+                 <input type="text" id="prompt" placeholder="Enter your prompt">
+             </div>
+             <div>
+                 <label for="guidance-scale">Guidance Scale:</label>
+                 <input type="number" id="guidance-scale" min="1" max="20" step="0.5" value="7.5">
+             </div>
+             <button onclick="generateImages()">Generate Images</button>
+             <div id="loading" class="loading">Generating images, please wait...</div>
+             <div id="error" class="error"></div>
+             <div id="output"></div>
+         </div>
+         <script>
+             async function generateImages() {
+                 const model = document.getElementById('model-select').value;
+                 const ratio = document.getElementById('ratio-select').value;
+                 const numImages = document.getElementById('num-images').value;
+                 const prompt = document.getElementById('prompt').value;
+                 const guidanceScale = document.getElementById('guidance-scale').value;
+                 const outputDiv = document.getElementById('output');
+                 const errorDiv = document.getElementById('error');
+                 const loadingDiv = document.getElementById('loading');
+
+                 outputDiv.innerHTML = '';
+                 errorDiv.innerText = '';
+                 loadingDiv.style.display = 'block';
+
+                 try {
+                     const response = await fetch('/generate', {
+                         method: 'POST',
+                         headers: { 'Content-Type': 'application/json' },
+                         body: JSON.stringify({
+                             model,
+                             prompt,
+                             ratio,
+                             num_images: numImages,
+                             guidance_scale: guidanceScale
+                         })
+                     });
+
+                     loadingDiv.style.display = 'none';
+                     const data = await response.json();
+
+                     if (response.ok) {
+                         data.images.forEach(imgSrc => {
+                             const img = document.createElement('img');
+                             img.src = imgSrc;
+                             img.className = 'output-image';
+                             outputDiv.appendChild(img);
+                         });
+                     } else {
+                         errorDiv.innerText = data.error || 'Failed to generate images';
+                     }
+                 } catch (error) {
+                     loadingDiv.style.display = 'none';
+                     errorDiv.innerText = 'Error: ' + error.message;
+                 }
+             }
+         </script>
+     </body>
+     </html>

requirements.txt

@@ -0,0 +1,10 @@
+flask==2.0.1
+werkzeug==2.0.2
+flask-cors==4.0.1
+torch==2.0.1
+diffusers==0.30.0
+transformers==4.31.0
+pillow==10.0.0
+huggingface_hub==0.16.4
+safetensors==0.3.3
+accelerate==0.21.0