ollamadiffuser 1.2.2__py3-none-any.whl → 2.0.0__py3-none-any.whl

ollamadiffuser/ui/templates/index.html

@@ -402,7 +402,7 @@
     <div class="container">
         <div class="header">
             <h1>🎨 OllamaDiffuser</h1>
-            <p>Image Generation Model Management Tool</p>
+            <p>Image Generation Model Management Tool &mdash; v2.0.0</p>
         </div>
         
         <div class="status-bar">
@@ -603,6 +603,12 @@
                             {% endif %}
                         </div>
                     </div>
+
+                    <div class="form-group">
+                        <label for="seed">Seed (optional, leave empty for random)</label>
+                        <input type="number" name="seed" id="seed" value="{{ seed or '' }}" min="0" max="4294967295" placeholder="Random">
+                        <small class="help-text">Set a seed for reproducible results</small>
+                    </div>
                     
                     <!-- ControlNet Section -->
                     {% if is_controlnet_model %}
@@ -691,6 +697,46 @@
                         🎨 Generate Image
                     </button>
                 </form>
+
+                <!-- img2img Section -->
+                <details style="margin-top: 20px;">
+                    <summary style="cursor: pointer; font-weight: 600; color: #374151; padding: 10px 0;">🖼️ Image-to-Image (img2img)</summary>
+                    <form method="post" action="/generate/img2img" enctype="multipart/form-data" style="margin-top: 10px;">
+                        <div class="form-group">
+                            <label for="img2img_prompt">Prompt</label>
+                            <textarea name="prompt" id="img2img_prompt" rows="2" placeholder="Describe the desired output..." required></textarea>
+                        </div>
+                        <div class="form-group">
+                            <label for="img2img_image">Input Image</label>
+                            <input type="file" name="image" id="img2img_image" accept="image/*" class="file-input" required>
+                        </div>
+                        <div class="form-row">
+                            <div class="form-group">
+                                <label for="img2img_strength">Strength</label>
+                                <input type="number" name="strength" id="img2img_strength" value="0.75" min="0.0" max="1.0" step="0.05">
+                                <small class="help-text">Higher = more change from original</small>
+                            </div>
+                            <div class="form-group">
+                                <label for="img2img_steps">Steps</label>
+                                <input type="number" name="num_inference_steps" id="img2img_steps" value="{{ model_parameters.get('num_inference_steps', 28) }}" min="1" max="100">
+                            </div>
+                        </div>
+                        <div class="form-row">
+                            <div class="form-group">
+                                <label for="img2img_guidance">Guidance Scale</label>
+                                <input type="number" name="guidance_scale" id="img2img_guidance" value="{{ model_parameters.get('guidance_scale', 3.5) }}" min="0" max="20" step="0.1">
+                            </div>
+                            <div class="form-group">
+                                <label for="img2img_seed">Seed</label>
+                                <input type="number" name="seed" id="img2img_seed" placeholder="Random" min="0" max="4294967295">
+                            </div>
+                        </div>
+                        <input type="hidden" name="negative_prompt" value="low quality, bad anatomy, worst quality, low resolution">
+                        <button type="submit" class="btn" {{ 'disabled' if not model_loaded }}>
+                            🖼️ Generate img2img
+                        </button>
+                    </form>
+                </details>
             </div>
             
             <div class="result-section">
@@ -722,6 +768,21 @@
                 <p style="color: #6b7280; margin-top: 15px; text-align: center;">
                     <strong>Prompt:</strong> {{ prompt }}
                 </p>
+                {% elif input_image_b64 and image_b64 %}
+                <!-- img2img Results -->
+                <div class="controlnet-results">
+                    <div>
+                        <h3 style="color: #374151; margin-bottom: 10px; text-align: center;">🖼️ Input Image</h3>
+                        <img src="data:image/png;base64,{{ input_image_b64 }}" alt="Input Image" class="result-image">
+                    </div>
+                    <div>
+                        <h3 style="color: #374151; margin-bottom: 10px; text-align: center;">🎨 Generated Image</h3>
+                        <img src="data:image/png;base64,{{ image_b64 }}" alt="Generated Image" class="result-image">
+                    </div>
+                </div>
+                <p style="color: #6b7280; margin-top: 15px; text-align: center;">
+                    <strong>Prompt:</strong> {{ prompt }}
+                </p>
                 {% elif image_b64 %}
                 <!-- Regular Generation Result -->
                 <img src="data:image/png;base64,{{ image_b64 }}" alt="Generated Image" class="result-image">

ollamadiffuser/ui/web.py

@@ -1,12 +1,15 @@
-from fastapi import FastAPI, Request, Form, File, UploadFile
-from fastapi.responses import HTMLResponse, StreamingResponse
-from fastapi.staticfiles import StaticFiles
-from fastapi.templating import Jinja2Templates
+import asyncio
 import io
 import base64
 import logging
 import json
 from pathlib import Path
+from typing import Optional
+
+from fastapi import FastAPI, Request, Form, File, UploadFile
+from fastapi.responses import HTMLResponse, StreamingResponse
+from fastapi.staticfiles import StaticFiles
+from fastapi.templating import Jinja2Templates
 from PIL import Image
 
 from ..core.models.manager import model_manager
@@ -364,72 +367,65 @@ def create_ui_app() -> FastAPI:
         guidance_scale: float = Form(3.5),
         width: int = Form(1024),
         height: int = Form(1024),
+        seed: Optional[int] = Form(None),
         control_image: UploadFile = File(None),
         controlnet_conditioning_scale: float = Form(1.0),
         control_guidance_start: float = Form(0.0),
-        control_guidance_end: float = Form(1.0)
+        control_guidance_end: float = Form(1.0),
     ):
-        """Generate image (Web UI)"""
+        """Generate image (Web UI) - runs generation in thread pool"""
         error_message = None
         image_b64 = None
         control_image_b64 = None
-        
+
         try:
-            # Check if model is actually loaded in memory
             if not model_manager.is_model_loaded():
                 error_message = "No model loaded. Please load a model first using the model management section above."
-            
+
             if not error_message:
-                # Get inference engine
                 engine = model_manager.loaded_model
-                
                 if engine is None:
                     error_message = "Model engine is not available. Please reload the model."
                 else:
                     # Process control image if provided
                     control_image_pil = None
                     if control_image and control_image.filename:
-                        # Initialize ControlNet preprocessors if needed
                         if not controlnet_preprocessor.is_initialized():
-                            logger.info("Initializing ControlNet preprocessors for image processing...")
                             if not controlnet_preprocessor.initialize():
-                                error_message = "Failed to initialize ControlNet preprocessors. Please check your installation."
-                        
+                                error_message = "Failed to initialize ControlNet preprocessors."
+
                         if not error_message:
-                            # Read uploaded image
                             image_data = await control_image.read()
                             control_image_pil = Image.open(io.BytesIO(image_data)).convert('RGB')
-                            
-                            # Convert control image to base64 for display
                             img_buffer = io.BytesIO()
                             control_image_pil.save(img_buffer, format='PNG')
                             img_buffer.seek(0)
                             control_image_b64 = base64.b64encode(img_buffer.getvalue()).decode()
-                    
-                    # Generate image
-                    image = engine.generate_image(
-                        prompt=prompt,
-                        negative_prompt=negative_prompt,
-                        num_inference_steps=num_inference_steps,
-                        guidance_scale=guidance_scale,
-                        width=width,
-                        height=height,
-                        control_image=control_image_pil,
-                        controlnet_conditioning_scale=controlnet_conditioning_scale,
-                        control_guidance_start=control_guidance_start,
-                        control_guidance_end=control_guidance_end
-                    )
-                    
-                    # Convert to base64
-                    img_buffer = io.BytesIO()
-                    image.save(img_buffer, format='PNG')
-                    img_buffer.seek(0)
-                    image_b64 = base64.b64encode(img_buffer.getvalue()).decode()
-                
+
+                    if not error_message:
+                        image = await asyncio.to_thread(
+                            engine.generate_image,
+                            prompt=prompt,
+                            negative_prompt=negative_prompt,
+                            num_inference_steps=num_inference_steps,
+                            guidance_scale=guidance_scale,
+                            width=width,
+                            height=height,
+                            seed=seed,
+                            control_image=control_image_pil,
+                            controlnet_conditioning_scale=controlnet_conditioning_scale,
+                            control_guidance_start=control_guidance_start,
+                            control_guidance_end=control_guidance_end,
+                        )
+
+                        img_buffer = io.BytesIO()
+                        image.save(img_buffer, format='PNG')
+                        img_buffer.seek(0)
+                        image_b64 = base64.b64encode(img_buffer.getvalue()).decode()
+
         except Exception as e:
             error_message = f"Image generation failed: {str(e)}"
-        
-        # Return result page
+
         context = get_template_context(request)
         context.update({
             "prompt": prompt,
@@ -438,14 +434,80 @@ def create_ui_app() -> FastAPI:
             "guidance_scale": guidance_scale,
             "width": width,
             "height": height,
+            "seed": seed,
             "controlnet_conditioning_scale": controlnet_conditioning_scale,
             "control_guidance_start": control_guidance_start,
             "control_guidance_end": control_guidance_end,
             "image_b64": image_b64,
             "control_image_b64": control_image_b64,
-            "error_message": error_message
+            "error_message": error_message,
         })
-        
+
+        return templates.TemplateResponse("index.html", context)
+
+    @app.post("/generate/img2img")
+    async def generate_img2img_ui(
+        request: Request,
+        prompt: str = Form(...),
+        negative_prompt: str = Form("low quality, bad anatomy, worst quality, low resolution"),
+        num_inference_steps: int = Form(28),
+        guidance_scale: float = Form(3.5),
+        seed: Optional[int] = Form(None),
+        strength: float = Form(0.75),
+        image: UploadFile = File(...),
+    ):
+        """Image-to-image generation (Web UI)"""
+        error_message = None
+        image_b64 = None
+        input_image_b64 = None
+
+        try:
+            if not model_manager.is_model_loaded():
+                error_message = "No model loaded."
+            else:
+                engine = model_manager.loaded_model
+                image_data = await image.read()
+                input_image = Image.open(io.BytesIO(image_data)).convert('RGB')
+
+                # Show the input image
+                buf = io.BytesIO()
+                input_image.save(buf, format='PNG')
+                buf.seek(0)
+                input_image_b64 = base64.b64encode(buf.getvalue()).decode()
+
+                result = await asyncio.to_thread(
+                    engine.generate_image,
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    num_inference_steps=num_inference_steps,
+                    guidance_scale=guidance_scale,
+                    width=input_image.width,
+                    height=input_image.height,
+                    seed=seed,
+                    image=input_image,
+                    strength=strength,
+                )
+
+                buf = io.BytesIO()
+                result.save(buf, format='PNG')
+                buf.seek(0)
+                image_b64 = base64.b64encode(buf.getvalue()).decode()
+
+        except Exception as e:
+            error_message = f"img2img generation failed: {str(e)}"
+
+        context = get_template_context(request)
+        context.update({
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "num_inference_steps": num_inference_steps,
+            "guidance_scale": guidance_scale,
+            "seed": seed,
+            "image_b64": image_b64,
+            "input_image_b64": input_image_b64,
+            "error_message": error_message,
+        })
+
         return templates.TemplateResponse("index.html", context)
     
     @app.post("/preprocess_control_image")
@@ -482,25 +544,24 @@ def create_ui_app() -> FastAPI:
     
     @app.post("/load_model")
     async def load_model_ui(request: Request, model_name: str = Form(...)):
-        """Load model (Web UI)"""
+        """Load model (Web UI) - runs loading in thread pool"""
         success = False
         error_message = None
-        
+
         try:
-            if model_manager.load_model(model_name):
+            if await asyncio.to_thread(model_manager.load_model, model_name):
                 success = True
             else:
                 error_message = f"Failed to load model {model_name}"
         except Exception as e:
             error_message = f"Error loading model: {str(e)}"
-        
-        # Return result page
+
         context = get_template_context(request)
         context.update({
             "success_message": f"Model {model_name} loaded successfully!" if success else None,
-            "error_message": error_message
+            "error_message": error_message,
         })
-        
+
         return templates.TemplateResponse("index.html", context)
     
     @app.post("/unload_model")
@@ -573,12 +634,13 @@ def create_ui_app() -> FastAPI:
         """Pull LoRA from Hugging Face Hub (Web UI)"""
         success = False
         error_message = None
-        
+        final_name = repo_id
+
         try:
             # Use alias if provided, otherwise use repo_id
             lora_alias = alias if alias.strip() else None
             weight_file = weight_name if weight_name.strip() else None
-            
+
             if lora_manager.pull_lora(repo_id, weight_name=weight_file, alias=lora_alias):
                 success = True
                 final_name = lora_alias if lora_alias else repo_id.replace('/', '_')
@@ -586,11 +648,11 @@ def create_ui_app() -> FastAPI:
                 error_message = f"Failed to download LoRA {repo_id}"
         except Exception as e:
             error_message = f"Error downloading LoRA: {str(e)}"
-        
+
         # Return result page
         context = get_template_context(request)
         context.update({
-            "success_message": f"LoRA {final_name if success else repo_id} downloaded successfully!" if success else None,
+            "success_message": f"LoRA {final_name} downloaded successfully!" if success else None,
             "error_message": error_message
         })