app.py

import gradio as gr
import torch
import numpy as np
from PIL import Image, ImageOps
import io
import base64
import cv2
import os
import requests
from fastapi import FastAPI, Request
import uvicorn
from gradio_client import Client, handle_file
from huggingface_hub import snapshot_download

# Configuration
MODEL_REPO = "LiuZichen/MagicQuill-models"
DRAW_N_GUESS_REPO = "LiuZichen/DrawNGuess"  # For prompt guessing (if used)
MAX_CONCURRENT_REQUESTS = 2  # Adjust based on your resources

# Model and Client Initialization
model_dir = "models"
os.makedirs(model_dir, exist_ok=True)
try:
    snapshot_download(repo_id=MODEL_REPO, repo_type="model", local_dir=model_dir)
    client = Client(DRAW_N_GUESS_REPO)  # For prompt guessing
    from MagicQuill import folder_paths  # Import *after* model download
    from MagicQuill.scribble_color_edit import ScribbleColorEditModel
    scribbleColorEditModel = ScribbleColorEditModel() # Initialize the model
except Exception as e:
    print(f"Error initializing models/client: {e}")
    exit(1)

# Helper Functions (as before - read_base64_image, load_and_preprocess_image, etc.)
# ... (Include all the helper functions from previous examples)

# Model Inference Functions
def generate(ckpt_name, total_mask, original_image, add_color_image, add_edge_image, remove_edge_image, positive_prompt, negative_prompt, grow_size, stroke_as_edge, fine_edge, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler):
    try:
      add_color_image, original_image, total_mask, add_edge_mask, remove_edge_mask = prepare_images_and_masks(total_mask, original_image, add_color_image, add_edge_image, remove_edge_image)
      # ... (Rest of the generate function - same logic as before)
    except Exception as e:
      print(f"Error in generate function: {e}")
      return None # Or handle the error as needed

def generate_image_handler(x, ckpt_name, negative_prompt, fine_edge, grow_size, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler):
    try:
        if seed == -1:
            seed = random.randint(0, 2**32 - 1)
        ms_data = x['from_frontend']
        positive_prompt = x['from_backend']['prompt']
        stroke_as_edge = "enable"
        res = generate(ckpt_name, ms_data['total_mask'], ms_data['original_image'], ms_data['add_color_image'], ms_data['add_edge_image'], ms_data['remove_edge_image'], positive_prompt, negative_prompt, grow_size, stroke_as_edge, fine_edge, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler)
        x["from_backend"]["generated_image"] = res
        return x
    except Exception as e:
        print(f"Error in generate_image_handler: {e}")
        return {"error": str(e)}, 500  # Return an error

# Gradio Interface
# ... (Gradio code - similar to before, but make sure the MagicQuill component is correctly integrated)
# ... (Use the generate_image_handler in the btn.click)

demo.queue(max_size=20, status_update_rate=0.1)

# FastAPI App
# ... (FastAPI code - same structure as before, with error handling)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)