
    5h;                        % S r SSKJrJr  SSKJr  SSKJrJrJ	r	  SSK
JrJr  SSKrSSKJr  SSKrSSKrSSKJr  \R(                  " \R*                  S	S
9  \R,                  " \5      rSq\\   \S'   \S\4S j5       r\" SSS\S9r\R;                  \S/SS/S/S9   " S S\5      r " S S\5      r " S S\5      r  " S S\5      r!\RE                  SS/S9S  5       r#\RI                  S!\ S"/S#9S$\4S% j5       r%\RE                  S&\!S'/S#9S( 5       r&\RE                  S)S'/S9S* 5       r'\RQ                  \)5      S+ 5       r*g),z
PROFESSIONAL IMAGE GENERATION API (Optimized for SD15 Only)
===========================================================
FastAPI server using Stable Diffusion v1.5 with VRAM-safe memory handling.
    )FastAPIHTTPException)CORSMiddleware)	BaseModelField	validator)ListOptionalN)asynccontextmanager)LocalImageGeneratorz4%(asctime)s - %(name)s - %(levelname)s - %(message)s)levelformat	generatorappc                (  #    [         R                  S5         [        S/S9q[         R                  S5        S7v   [         R                  S5        [        (       ag  [        R                  S5        B[        R                  " 5         [        R                  R                  5         [        R                  R                  5         [         R                  S	5        g! [         a+  n[         R                  SU 35        [        SU 35      eSnAff = f7f)
z|
Lifespan context manager for initializing and cleaning up models.
Only SD15 is preloaded to prevent GPU VRAM accumulation.
u1   🚀 Starting Image Generation API (SD15 only)...sd15)preload_modelsu&   ✅ SD15 model preloaded successfully!u$   ❌ Failed to initialize generator: zModel initialization failed: Nu2   🛑 Shutting down API and releasing GPU memory...u   ✅ Clean shutdown complete!)loggerinfor   r   	ExceptionerrorRuntimeErrorunload_modelgccollecttorchcudaempty_cachesynchronize)r   es     E/var/www/eduai.edurigo.com/generate_image/staging/image_generation.pylifespanr"      s      KKCD@'x@	<=
 

KKDEyv&




 

 
KK./  @;A3?@:1#>??@s(   DC B"D
D$&D

DDz(Professional Image Generation API (SD15)u_   🚀 GPU-based Self-Hosted Image Generator optimized for VRAM efficiency (Stable Diffusion 1.5)7.0.0)titledescriptionversionr"   *T)allow_originsallow_credentialsallow_methodsallow_headersc                       \ rS rSr% \" SSSSS9r\\S'   \" SSS	S
9r\\S'   \" SSSSS9r	\
\S'   \" SSSSS9r\
\S'   \" SSSSS9r\
\S'   \" S5      S 5       rSrg)ImageRequestS   .zText prompt   z!a futuristic city skyline at dusk)r%   
min_lengthexamplepromptz(Visual category (e.g., realistic, anime)	realistic)r%   r1   categoryi   zWidth in pixels   i   )r%   gelewidthzHeight in pixelsheight   zNumber of imagesimage_countc                     / SQnUR                  5       U;  a  [        R                  SU S35        gUR                  5       $ )Nr3   3dcartoonisticcomicanime	cinematicfantasy	cyberpunku   ⚠️ Unknown category 'z', defaulting to 'realistic'r3   )lowerr   warning)clsvaluealloweds      r!   validate_categoryImageRequest.validate_categoryZ   s;    l;;='NN6ug=YZ[{{}     N)__name__
__module____qualname____firstlineno__r   r2   str__annotations__r4   r8   intr9   r;   r   rJ   __static_attributes__rM   rL   r!   r-   r-   S   s    1NqrFCr#+U_jkHcks(9cdKE3K);MFCMQ,>1KKKz rL   r-   c                   *    \ rS rSr% \\S'   \\S'   Srg)	ImageDatac   namebase64rM   N)rN   rO   rP   rQ   rR   rS   rU   rM   rL   r!   rW   rW   c   s    
IKrL   rW   c                   4    \ rS rSr% \\S'   \\S'   \\S'   Srg)ImageResponseh   statusmessagedatarM   NrN   rO   rP   rQ   rR   rS   dictrU   rM   rL   r!   r\   r\   h   s    KL
JrL   r\   c                   *    \ rS rSr% \\S'   \\S'   Srg)HealthResponsen   r^   r`   rM   Nra   rM   rL   r!   rd   rd   n   s    K
JrL   rd   /Info)tagsc            	      $   #    SSSSSSSSS	.S
.$ 7f)Nz!Professional Image Generation APIr#   r   operational/generate-image/health/modelsz/docs)generatehealthmodelsdocs)servicer&   modelr^   	endpointsrM   rM   rL   r!   rootru   v   s.      7)	
    rk   
Generation)response_modelrh   requestc           	        #    [         c
  [        SSS9e[        R                  SU R                   SU R
                   SU R                   35        [        R                  SU R                  S S  S	35         [         R                  U R                  U R                  S
U R
                  U R                  U R                  S9nU(       d
  [        SSS9e[        R                  S[        U5       S35        SS[        U5       S3US
S.S.[        R                  R                  5         [        R                  " 5         $ ! [          a<  n[        R#                  S[%        U5       35        [        SS[%        U5       3S9eS nAff = f! [        R                  R                  5         [        R                  " 5         f = f7f)N  z*Generator not initialized. Restart server.status_codedetailu   📥 Generating z image(s) with SD15 | xz   Prompt: d   z...r   )r2   r4   
model_namer8   r9   r;     zNo images generated.u   ✅ Successfully generated z image(s) using SD15successz
Generated z image(s) successfully)images
model_used)r^   r_   r`   u   ❌ Generation error: zError generating image: )r   r   r   r   r;   r8   r9   r2   generate_multiple_imagesr4   lenr   r   r   r   r   r   r   rR   )ry   r   r    s      r!   generate_imager      s~    4`aa
KK"7#6#6"77Mgmm_\]^e^l^l]mno
KK+gnnTc2337833>>%%-->>++ 4 
 C8NOO1#f+>RST#CK=0FG%V<
 	

 


  Y-c!fX676NsSTvh4WXXY 	

 


s7   A4F>7BD= 	4F>=
F7E>>FF 5F;;F>rl   Systemc                     #    [         c
  [        SSS9e S[         R                  5       S.$ ! [         a  n [        S[	        U 5      S9eS n A ff = f7f)Nr{   zGenerator not initializedr|   r   r^   r`   r   )r   r   check_healthr   rR   )r    s    r!   health_checkr      sR     4OPP<#Y-C-C-EFF <CF;;<s$   A, A
AA		AArm   c                  $   #    SS/S// SQS.S.$ 7f)Nr   r   r=   )available_modelspreloaded_modelssupported_categoriesr   rM   rM   rL   r!   list_modelsr      s(      !'!'	%
 rv   c                 (   #    S[        U5      SS.$ 7f)Nr   validation_error)r^   r_   type)rR   )ry   excs     r!   value_error_handlerr      s     #c(<NOOs   )+__doc__fastapir   r   fastapi.middleware.corsr   pydanticr   r   r   typingr	   r
   logging
contextlibr   r   r   local_image_generatorr   basicConfigINFO	getLoggerrN   r   r   rS   r"   r   add_middlewarer-   rW   r\   rd   getru   postr   r   r   exception_handler
ValueErrorr   rM   rL   r!   <module>r      s   + 2 0 0 !  * 	  5   
,,A 
		8	$
 ,0	8'( / 0 0 0@ 
4q	   %%%  9  	 
I Y  F8  
MO ,   P F >
C< D< ($ %. z"P #PrL   